javascriptnode.jspostgresqlknex.js

Knexjs pooling - set session variables for a connection


I'm working on a nodejs server using Knex.js to connect to a postgres DB. On some of the tables I've set some RLS policies based on the current session values (app.user_id and app.tenant_id) to isolate data correctly.

I would like to be able to pass the userId and tenantId from my model/service methods to a db util method that grabs a pool connection, sets the session variables and returns the connection to be used (all in Knex ideally).

Currently I have the following:

const buildSetSessionQuery = ({userId, tenantId}) => {
    let queryString = "";
    if(userId) {
        queryString += format(`SET SESSION app.user_id=%L;`, userId);
    }

    if(tenantId) {
        queryString += format(`SET SESSION app.tenant_id=%L;`, tenantId);
    }

    return queryString;
};

exports.getUserConnection = ({userId, tenantId}={}) => {
    const conn = knex({
        client: 'pg',
        connection: getConfig(), // Assume this returns the correct connection config
        pool: {
            min: 0,
            max: 40,
            afterCreate: function (conn, done) {
                const setSessionQuery = buildSetSessionQuery({userId, tenantId});
                conn.query(setSessionQuery);
                done (null, conn);
            }
        }
    });

    return conn;
};

But this code creates a whole Knex pool for each query that's executed (as opposed to grabbing an available pool connection from a single pool), which quickly exhausts the max_connections limit on my pg server.

Is there a way to make it so getUserConnection() returns a knex pool connection with the session variables set, so that I can use it easily like so:

exports.getData = async ({tenantId, userId}) => {
    const conn = db.getUserConnection({ tenantId, userId });

    const rows = await conn
        .select('*')
        .from('data')
        .orderBy([{
            column: 'timestamp',
            order: 'desc',
            nulls: 'last'
        }]);

    return rows;
};

Solution

  • You can get around this by creating a method that initiates a database transaction and sets the session details.

    The method below takes a call back function that expects the transaction connection mentioned above to be passed in.

    This does seem like a little bit of overhead to create a transaction for each tenant query, but it seems to work.

    /**
    * Establishes a database connection with session settings for the specified user and tenant.
            *
            * @param {Object} params - The parameters object.
            * @param {string} params.userId - The ID of the user.
            * @param {string} params.tenantId - The ID of the tenant.
            * @returns {Object} - The Knex instance representing the database connection.
            */
            exports.queryWithRLS = async ({userId, tenantId}={}, query) => {
                const conn = await userConn.transaction();
                try {
                    await conn.raw(buildSetSessionQuery({userId, tenantId}))
                    const data = await query(conn);
                    await conn.commit();
                    return data;
                } catch (error) {
                    await conn.rollback();
                    throw error;
                }
            };
             
             
             //Service File
            const rows = await db.queryWithRLS({userId: 1, tenantId: 1,}, async (conn) => {
                        return await conn.select('*').from(`data`)
                        .where({'tenant_id': 1})
                    });
                    return res.status(200).json({rows})
            }