supabase local db pooler with drizzle

it was a bit tedious to setup so I'm posting what worked for me in case it helps anyone:

config.toml

[db.pooler]
enabled = true
# Port to use for the local connection pooler.
port = 54329
# Specifies when a server connection can be reused by other clients.
# Configure one of the supported pooler modes: `transaction`, `session`.
pool_mode = "transaction"
# How many server connections to allow per user/database pair.
default_pool_size = 20
# Maximum number of client connections allowed.
max_client_conn = 100

these connection strings worked for me:

ADMIN_DATABASE_URL=postgresql://postgres.pooler-dev:[email protected]:54329/postgres
DATABASE_URL=postgresql://rls_client.pooler-dev:[email protected]:54329/postgres

and to create a client with drizzle:

import { type DrizzleConfig, sql } from 'drizzle-orm';
import { drizzle as drizzlePostgres } from 'drizzle-orm/node-postgres';
import * as schema from 'kotoba-core/schema';
import { serverSupabaseSession } from '#root/server/serverSupabaseClient';
import { decode } from './jwt';
import { QueryLogger } from 'kotoba-core/database';
import { fileLogger } from 'kotoba-core/server';
import pg from 'pg';

const adminPool = new pg.Pool({
  connectionString: process.env.ADMIN_DATABASE_URL,
  max: 5,
});

const rlsPool = new pg.Pool({
  connectionString: process.env.DATABASE_URL,
  max: 5,
});

// https://github.com/orgs/supabase/discussions/23224
// Should be secure because we use the access token that is signed, and not the data read directly from the storage
export function createDrizzleSupabaseClient(session: Awaited<ReturnType<typeof serverSupabaseSession>>) {
  const config = {
    casing: 'snake_case',
    schema: schema,
    logger: process.env.NODE_ENV === 'production' ? true : new QueryLogger(fileLogger({})),
  } satisfies DrizzleConfig<typeof schema>;

  const admin = drizzlePostgres({
    client: adminPool,
    ...config,
  });

  const client = drizzlePostgres({
    client: rlsPool,
    ...config,
  });

  if (!session?.access_token) {
    return {
      admin,
      rls: (async (transaction, ...rest) => {
        return await client.transaction(async (tx) => {
          try {
            return await transaction(tx);
          } catch (e) {
            console.error('Transaction failed: ', e);
          } finally {

          }
        }, ...rest);
      }) as typeof client.transaction,
    };
  }

  const token = decode(session.access_token);

  return {
    admin,
    rls: (async (transaction, ...rest) => {
      return await client.transaction(async (tx) => {
        // Supabase exposes auth.uid() and auth.jwt()
        // https://supabase.com/docs/guides/database/postgres/row-level-security#helper-functions
        try {
          await tx.execute(
              sql`
---------------------------- CONFIG START 
select set_config('request.jwt.claims', '${sql.raw(JSON.stringify(token, null, 2))}', TRUE);
select set_config('request.jwt.claim.sub', '${sql.raw(token.sub ?? '')}', TRUE);
set local role ${sql.raw(token.role ?? 'anon')};
---------------------------- CONFIG END
          `,
          );
          return await transaction(tx);
        } catch (e) {
          console.error('Transaction failed: ', e);
        } finally {
          await tx.execute(
              sql`
---------------------------- RESET START
select set_config('request.jwt.claims', NULL, TRUE);
select set_config('request.jwt.claim.sub', NULL, TRUE);
reset role;
---------------------------- RESET END
          `,
          );
        }
      }, ...rest);
    }) as typeof client.transaction,
  };
}

marpe/supabase_local_drizzle_pool.md