Multi-tenant apps with Postgres RLS

schema.ts

import {
  varchar,
  pgPolicy,
  pgRole,
  pgSchema,
  type AnyPgColumn,
  type PgRole,
} from 'drizzle-orm/pg-core'
import { createId } from '@/lib/id'
import { sql, type SQL } from 'drizzle-orm'

// optional, define schema in "ui" namespace
const schema = pgSchema('ui')

// roles and postgres users are the same thing
const userRole = pgRole('ui_user').existing()

const crudPolicy = ({ role, read, modify }: { role: PgRole; read: SQL; modify: SQL }) => {
  return [
    pgPolicy(`crud-${role.name}-policy-select`, {
      for: 'select',
      to: role,
      using: read,
    }),
    pgPolicy(`crud-${role.name}-policy-insert`, {
      for: 'insert',
      to: role,
      withCheck: modify,
    }),
    pgPolicy(`crud-${role.name}-policy-update`, {
      for: 'update',
      to: role,
      using: modify,
      withCheck: modify,
    }),
    pgPolicy(`crud-${role.name}-policy-delete`, {
      for: 'delete',
      to: role,
      using: modify,
    }),
  ]
}

const onlyOwnOrg = (orgIdColumn: AnyPgColumn) => {
  return sql`${orgIdColumn} IS NOT NULL AND ${orgIdColumn} = current_setting('auth.org_id')`
}

const orgIdColumn = varchar('org_id')
  .notNull()
  .default(sql`current_setting('auth.org_id')`)

export const casesTable = schema.table(
  'cases',
  {
    id: varchar('id')
      .primaryKey()
      .$defaultFn(() => createId('case')),
    orgId: orgIdColumn,
    
    // ...
  },
  (table) => [
    crudPolicy({
      role: userRole,
      read: onlyOwnOrg(table.orgId),
      modify: onlyOwnOrg(table.orgId),
    }),
  ],
)

//
// APLICATION LAYER
//

import { auth } from '@clerk/nextjs/server'

const db = drizzle(/* ... */)

type Transaction = Parameters<Parameters<(typeof db)['transaction']>[0]>[0]

export async function withAuth<T>(callback: (db: Transaction) => Promise<T>): Promise<T> {
  const { userId, orgId } = await auth()

  if (!userId) {
    throw new Error('User not authenticated')
  }

  if (!orgId) {
    throw new Error('Organization not found')
  }

  return await db.transaction(async (tx) => {
    await tx.execute(
      sql.raw(`
        SET LOCAL auth.user_id = '${userId}';
        SET LOCAL auth.org_id = '${orgId}';
      `),
    )
    return await callback(tx)
  })
}