Adding a statement or node to the beginning of a query

[charlie-maloney]

I am trying to implement row level security in my project to better isolate my tenant's data. I am looking for a way to add the following statement to the beginning of query:

SET app.current_tenant = ${sql.raw("${tenantId}")}`

The entire query might be:

SET app.current_tenant = "tenant-1"; SELECT * from "tenants" WHERE id = "tenant-1"; 

I can do this manually for each query, but I want to find a way to get a connection from the pool that already has this prepended to the query based on the request context. I'm not sure of the best way to do this in Kysely. Maybe I need to create a custom plugin?

I was thinking of something like below, but after looking at the docs for extending Kysely, I wasn't sure how to actually add a node to the beginning of a query.

export class ConnectionService {
  constructor(@InjectDb() private db: Database) {}

  getConnection() {
    // returns an instance of kysely with the SET app.current_tenant pre-pended to every query
    return this.db.withPlugin(new TenantContextPlugin());
  }
}

Any help would be much appreciated. Thanks for this awesome lib!

[igalklebanov]

Hey 👋

Have you checked https://github.com/ben-pr-p/kysely-access-control for RLS?

Seeing that you're trying to send 2 queries batched, and Kysely's core dialects only ever use the first query's results, I'm afraid you might have to implement your own driver/connection and combine it with existing core dialect components.

Also, IIRC there's currently a restriction in place denying plugins from changing the root node's kind. This takes away the ability to turn the query into a raw node that represents set ...; {originalQuery}.

[charlie-maloney]

Hi @igalklebanov --

Thanks for your response. I did check out that library, and it does look interesting. While it would be helpful for implementing more robustness / security to my queries, it doesn't actually use Postgres RLS (as far as I understood).

I will look into the driver/connection logic. Glad I asked though--I wanted to make sure I wasn't missing something simple here.

[charlie-maloney]

Following up here @igalklebanov. If you are interested, here is what I came up with. Let me know if you have any suggestions / comments.

I end up making two queries (one to set the app.current_tenant and one for the actual query) but I think this is unavoidable.

type Callback = (db: Database) => any;

@Injectable()
export class ConnectionService {
  constructor(@InjectDb() private db: Database) {}

  getConnection() {
    const connection = this.db.connection();

    return async <T extends Callback>(
      callback: Callback,
    ): Promise<ReturnType<T>> => {
      const setTenant = sql`SET app.current_tenant = ${sql.raw(ctx.getTenant())}`;

      return await connection.execute(async (db) => {
        await db.executeQuery(setTenant.compile(db));

        return callback(db);
      });
    };
  }
}

// Use in application
class SomeService {
  constructor(private connectionService: ConnectionService) {}

  async someMethod() {
    const withConnection = this.connectionService.getConnection();

    const queryFn = async (db: Database) => {
      return db.selectFrom('users').selectAll().executeTakeFirst();
    };

    const result = await withConnection<typeof queryFn>(queryFn);

    console.log('result', result);
  }
}