@dotdo/postgres 0.1.0

package/src/rls/generator.ts

@@ -0,0 +1,573 @@
+/**
+ * RLS Policy SQL Generator
+ *
+ * SECURITY CRITICAL: This module generates SQL statements for Row-Level Security policies.
+ * All SQL generation must prevent SQL injection attacks through proper identifier quoting
+ * and validation. Never interpolate untrusted input directly into SQL strings.
+ *
+ * Key security considerations:
+ * - All identifiers (policy names, table names, schema names) must be properly quoted
+ * - Double quotes within identifiers must be escaped by doubling them
+ * - Table names must be validated before use to prevent SQL injection
+ * - Policy expressions (USING, WITH CHECK) are passed through as-is; validation is the
+ *   responsibility of the caller or the validatePolicy function in validator.ts
+ *
+ * @module rls/generator
+ */
+
+import type { RLSPolicy, PGliteLike, ApplyPolicyOptions, EnableRLSOptions, DisableRLSOptions } from './policy'
+import { needsQuoting, validateTableName } from './validator'
+import { InvalidIdentifierError, NotFoundError } from '@dotdo/postgres-shared/errors'
+
+/**
+ * Escape double quotes within an identifier for PostgreSQL.
+ *
+ * PostgreSQL requires double quotes within quoted identifiers to be escaped
+ * by doubling them: " becomes "".
+ *
+ * SECURITY: This is essential for preventing SQL injection through identifiers
+ * that contain quote characters.
+ *
+ * @param identifier - The identifier to escape
+ * @returns The identifier with internal double quotes escaped
+ *
+ * @example
+ * escapeIdentifier('my"table') // returns 'my""table'
+ */
+function escapeIdentifier(identifier: string): string {
+  return identifier.replace(/"/g, '""')
+}
+
+/**
+ * Quote an identifier if necessary for safe SQL generation.
+ *
+ * PostgreSQL identifiers (table names, column names, policy names) must be quoted
+ * if they contain special characters, start with a number, or match reserved keywords.
+ *
+ * SECURITY: Always use this function when embedding identifiers in generated SQL.
+ * Never concatenate raw identifiers directly into SQL strings.
+ *
+ * @param identifier - The identifier to quote
+ * @returns The identifier, quoted with escaped internal quotes if necessary
+ *
+ * @example
+ * quoteIdentifier('users')         // returns 'users' (no quoting needed)
+ * quoteIdentifier('my-table')      // returns '"my-table"' (hyphen requires quoting)
+ * quoteIdentifier('my"table')      // returns '"my""table"' (quote escaped)
+ */
+function quoteIdentifier(identifier: string): string {
+  if (needsQuoting(identifier)) {
+    return `"${escapeIdentifier(identifier)}"`
+  }
+  return identifier
+}
+
+/**
+ * Quote a table name, handling schema.table format.
+ *
+ * Supports both simple table names and schema-qualified names (schema.table).
+ * Each component is individually quoted if necessary.
+ *
+ * SECURITY: Always use this function when embedding table names in generated SQL.
+ * The function properly handles:
+ * - Simple table names: 'users' -> 'users' or '"my-table"' -> '"my-table"'
+ * - Schema-qualified names: 'public.users' -> 'public.users'
+ * - Mixed quoting: 'my-schema.my-table' -> '"my-schema"."my-table"'
+ *
+ * @param table - The table name, optionally schema-qualified
+ * @returns The properly quoted table reference
+ *
+ * @example
+ * quoteTableName('users')              // returns 'users'
+ * quoteTableName('public.users')       // returns 'public.users'
+ * quoteTableName('my-schema.my-table') // returns '"my-schema"."my-table"'
+ */
+function quoteTableName(table: string): string {
+  if (table.includes('.')) {
+    const parts = table.split('.')
+    const schema = parts[0] ?? ''
+    const tableName = parts[1] ?? ''
+    return `${quoteIdentifier(schema)}.${quoteIdentifier(tableName)}`
+  }
+  return quoteIdentifier(table)
+}
+
+/**
+ * Extract the table name from a potentially schema-qualified name.
+ *
+ * Used when querying pg_policies which stores only the table name without schema.
+ *
+ * @param table - The table name, optionally schema-qualified
+ * @returns The table name without schema prefix
+ *
+ * @example
+ * extractTableName('public.users')  // returns 'users'
+ * extractTableName('users')         // returns 'users'
+ */
+function extractTableName(table: string): string {
+  if (table.includes('.')) {
+    return table.split('.')[1] ?? table
+  }
+  return table
+}
+
+/**
+ * Escape a string value for use in SQL single-quoted literals.
+ *
+ * PostgreSQL requires single quotes within string literals to be escaped
+ * by doubling them: ' becomes ''.
+ *
+ * SECURITY: Use this when embedding string values in SQL queries.
+ * Prefer parameterized queries when available.
+ *
+ * @param value - The string value to escape
+ * @returns The string with internal single quotes escaped
+ */
+function escapeStringLiteral(value: string): string {
+  return value.replace(/'/g, "''")
+}
+
+/**
+ * Generate CREATE POLICY SQL statement.
+ *
+ * Creates a PostgreSQL RLS policy with the specified configuration.
+ *
+ * SECURITY CRITICAL: This function generates SQL that controls data access.
+ * - Policy names and table names are properly quoted to prevent injection
+ * - USING and WITH CHECK expressions are passed through as-is; they must be
+ *   validated separately using validatePolicy() before calling this function
+ * - Role names are NOT quoted (PostgreSQL syntax); callers must ensure roles
+ *   are validated before passing to this function
+ *
+ * WARNING: Missing USING or WITH CHECK clauses have security implications:
+ * - SELECT without USING: PostgreSQL defaults to true (allows all rows)
+ * - INSERT without WITH CHECK: PostgreSQL defaults to true (allows all inserts)
+ * - UPDATE without USING: Cannot select rows to update
+ * - UPDATE without WITH CHECK: PostgreSQL defaults to USING clause value
+ *
+ * @param policy - The RLS policy definition
+ * @returns The complete CREATE POLICY SQL statement
+ *
+ * @example
+ * const sql = generatePolicySQL({
+ *   name: 'users_read_own',
+ *   table: 'users',
+ *   operation: 'SELECT',
+ *   using: 'user_id = current_user_id()',
+ *   roles: ['authenticated'],
+ *   type: 'permissive'
+ * })
+ * // Returns: CREATE POLICY users_read_own ON users AS PERMISSIVE FOR SELECT TO authenticated USING (user_id = current_user_id())
+ */
+export function generatePolicySQL(policy: RLSPolicy): string {
+  const parts: string[] = []
+
+  // CREATE POLICY name ON table
+  // SECURITY: Both policy.name and policy.table are quoted to prevent injection
+  parts.push(`CREATE POLICY ${quoteIdentifier(policy.name)} ON ${quoteTableName(policy.table)}`)
+
+  // AS PERMISSIVE | RESTRICTIVE
+  // Note: RESTRICTIVE policies are AND-ed together; PERMISSIVE are OR-ed
+  parts.push(`AS ${policy.type.toUpperCase()}`)
+
+  // FOR operation
+  parts.push(`FOR ${policy.operation}`)
+
+  // TO roles (default to 'public' if empty)
+  // WARNING: Role names are not quoted here per PostgreSQL syntax.
+  // Callers must validate role names before passing to this function.
+  const roles = policy.roles.length > 0 ? policy.roles.join(', ') : 'public'
+  parts.push(`TO ${roles}`)
+
+  // USING clause (not for INSERT)
+  // SECURITY: The USING expression controls which existing rows are visible.
+  // WARNING: If policy.using is empty/undefined, PostgreSQL defaults to true (all rows visible)
+  if (policy.using && policy.operation !== 'INSERT') {
+    parts.push(`USING (${policy.using})`)
+  }
+
+  // WITH CHECK clause (only for INSERT, UPDATE, ALL - not SELECT or DELETE)
+  // SECURITY: The WITH CHECK expression controls what new/modified rows are allowed.
+  // WARNING: If policy.check is empty/undefined, PostgreSQL defaults to the USING expression
+  const allowsWithCheck = policy.operation === 'INSERT' || policy.operation === 'UPDATE' || policy.operation === 'ALL'
+  if (policy.check && allowsWithCheck) {
+    parts.push(`WITH CHECK (${policy.check})`)
+  }
+
+  return parts.join(' ')
+}
+
+/**
+ * Generate DROP POLICY SQL statement.
+ *
+ * Creates SQL to drop an existing RLS policy from a table.
+ *
+ * SECURITY: Policy name and table name are properly quoted to prevent injection.
+ *
+ * WARNING: Dropping a policy removes access controls. Ensure this is intentional
+ * and that replacement policies are in place if needed.
+ *
+ * @param policyName - The name of the policy to drop
+ * @param table - The table the policy is attached to
+ * @param ifExists - If true, adds IF EXISTS to prevent errors if policy doesn't exist
+ * @returns The complete DROP POLICY SQL statement
+ *
+ * @example
+ * generateDropPolicySQL('users_read_own', 'users', true)
+ * // Returns: DROP POLICY IF EXISTS users_read_own ON users
+ */
+export function generateDropPolicySQL(policyName: string, table: string, ifExists: boolean = false): string {
+  const ifExistsClause = ifExists ? 'IF EXISTS ' : ''
+  return `DROP POLICY ${ifExistsClause}${quoteIdentifier(policyName)} ON ${quoteTableName(table)}`
+}
+
+/**
+ * Generate ALTER TABLE ... ENABLE ROW LEVEL SECURITY SQL.
+ *
+ * Creates SQL to enable RLS on a table. Once enabled, all table access
+ * (except by table owners and superusers) is governed by RLS policies.
+ *
+ * SECURITY CRITICAL:
+ * - ENABLE ROW LEVEL SECURITY: Enables RLS but allows table owner to bypass
+ * - FORCE ROW LEVEL SECURITY: Enforces RLS even for table owner
+ *
+ * WARNING: Enabling RLS without any policies will block all access to the table
+ * (except for superusers and table owners unless FORCE is used).
+ *
+ * @param table - The table to enable RLS on
+ * @param options - Optional settings
+ * @param options.force - If true, uses FORCE ROW LEVEL SECURITY (applies to owner too)
+ * @returns The complete ALTER TABLE SQL statement
+ *
+ * @example
+ * generateEnableRLSSQL('users')               // Returns: ALTER TABLE users ENABLE ROW LEVEL SECURITY
+ * generateEnableRLSSQL('users', { force: true }) // Returns: ALTER TABLE users FORCE ROW LEVEL SECURITY
+ */
+export function generateEnableRLSSQL(table: string, options?: EnableRLSOptions): string {
+  if (options?.force) {
+    return `ALTER TABLE ${quoteTableName(table)} FORCE ROW LEVEL SECURITY`
+  }
+  return `ALTER TABLE ${quoteTableName(table)} ENABLE ROW LEVEL SECURITY`
+}
+
+/**
+ * Generate ALTER TABLE ... DISABLE ROW LEVEL SECURITY SQL.
+ *
+ * Creates SQL to disable RLS on a table.
+ *
+ * SECURITY WARNING: Disabling RLS removes all row-level access controls.
+ * All users with table-level permissions will have unrestricted access to all rows.
+ * Only use this when:
+ * - Migrating to a different access control mechanism
+ * - Temporarily debugging access issues (re-enable immediately after)
+ * - The table genuinely requires no row-level restrictions
+ *
+ * @param table - The table to disable RLS on
+ * @param options - Optional settings
+ * @param options.removeForce - If true, uses NO FORCE (removes owner bypass restriction)
+ * @returns The complete ALTER TABLE SQL statement
+ *
+ * @example
+ * generateDisableRLSSQL('users')                     // Returns: ALTER TABLE users DISABLE ROW LEVEL SECURITY
+ * generateDisableRLSSQL('users', { removeForce: true }) // Returns: ALTER TABLE users NO FORCE ROW LEVEL SECURITY
+ */
+export function generateDisableRLSSQL(table: string, options?: DisableRLSOptions): string {
+  if (options?.removeForce) {
+    return `ALTER TABLE ${quoteTableName(table)} NO FORCE ROW LEVEL SECURITY`
+  }
+  return `ALTER TABLE ${quoteTableName(table)} DISABLE ROW LEVEL SECURITY`
+}
+
+/**
+ * Enable Row-Level Security on a table.
+ *
+ * Validates the table name and executes the ENABLE ROW LEVEL SECURITY command.
+ *
+ * SECURITY: Table name is validated to prevent SQL injection before generating SQL.
+ *
+ * @param pglite - The PGlite instance to execute against
+ * @param table - The table to enable RLS on (must pass validation)
+ * @param options - Optional settings (see generateEnableRLSSQL)
+ * @throws {Error} If table name is invalid (contains SQL injection patterns)
+ *
+ * @example
+ * await enableRLS(pglite, 'users')
+ * await enableRLS(pglite, 'users', { force: true })
+ */
+export async function enableRLS(
+  pglite: PGliteLike,
+  table: string,
+  options?: EnableRLSOptions
+): Promise<void> {
+  // SECURITY: Validate table name to prevent SQL injection
+  if (!validateTableName(table)) {
+    throw new InvalidIdentifierError(
+      table,
+      'table',
+      'Must be a valid SQL identifier (letters, numbers, underscores; must start with letter or underscore)'
+    )
+  }
+
+  const sql = generateEnableRLSSQL(table, options)
+  await pglite.exec(sql)
+}
+
+/**
+ * Disable Row-Level Security on a table.
+ *
+ * Validates the table name and executes the DISABLE ROW LEVEL SECURITY command.
+ *
+ * SECURITY WARNING: This removes all row-level access controls from the table.
+ * See generateDisableRLSSQL for security considerations.
+ *
+ * @param pglite - The PGlite instance to execute against
+ * @param table - The table to disable RLS on (must pass validation)
+ * @param options - Optional settings (see generateDisableRLSSQL)
+ * @throws {Error} If table name is invalid (contains SQL injection patterns)
+ *
+ * @example
+ * await disableRLS(pglite, 'users')
+ */
+export async function disableRLS(
+  pglite: PGliteLike,
+  table: string,
+  options?: DisableRLSOptions
+): Promise<void> {
+  // SECURITY: Validate table name to prevent SQL injection
+  if (!validateTableName(table)) {
+    throw new InvalidIdentifierError(
+      table,
+      'table',
+      'Must be a valid SQL identifier (letters, numbers, underscores; must start with letter or underscore)'
+    )
+  }
+
+  const sql = generateDisableRLSSQL(table, options)
+  await pglite.exec(sql)
+}
+
+/**
+ * Apply an RLS policy to a database.
+ *
+ * Creates a new RLS policy, optionally replacing an existing one with the same name.
+ *
+ * SECURITY CRITICAL: This function directly affects data access controls.
+ * - Always validate the policy using validatePolicy() before calling this function
+ * - The replace option drops the existing policy first (brief window of no protection)
+ *
+ * @param pglite - The PGlite instance to execute against
+ * @param policy - The RLS policy definition to apply
+ * @param options - Optional settings
+ * @param options.replace - If true, drops existing policy with same name first
+ * @param options.ifNotExists - If true, adds IF NOT EXISTS (no error if policy exists)
+ *
+ * @example
+ * const policy = createPolicy({
+ *   name: 'users_read_own',
+ *   table: 'users',
+ *   operation: 'SELECT',
+ *   using: 'user_id = current_user_id()'
+ * })
+ * await applyPolicy(pglite, policy)
+ */
+export async function applyPolicy(
+  pglite: PGliteLike,
+  policy: RLSPolicy,
+  options?: ApplyPolicyOptions
+): Promise<void> {
+  // Drop existing policy if replace is true
+  // WARNING: Brief window between DROP and CREATE where no policy exists
+  if (options?.replace) {
+    const dropSQL = generateDropPolicySQL(policy.name, policy.table, true)
+    await pglite.exec(dropSQL)
+  }
+
+  // Generate CREATE POLICY SQL
+  let sql = generatePolicySQL(policy)
+
+  // Add IF NOT EXISTS if specified
+  if (options?.ifNotExists) {
+    sql = sql.replace('CREATE POLICY', 'CREATE POLICY IF NOT EXISTS')
+  }
+
+  await pglite.exec(sql)
+}
+
+/**
+ * Drop a policy by name.
+ *
+ * Removes an RLS policy from a table.
+ *
+ * SECURITY WARNING: Dropping a policy removes access controls. Ensure:
+ * - This is intentional
+ * - Replacement policies are in place if needed
+ * - The table owner bypass is acceptable, or FORCE is still enabled
+ *
+ * @param pglite - The PGlite instance to execute against
+ * @param policyName - The name of the policy to drop
+ * @param table - The table the policy is attached to
+ * @param options - Optional settings
+ * @param options.ifExists - If true, no error if policy doesn't exist
+ *
+ * @example
+ * await dropPolicy(pglite, 'users_read_own', 'users', { ifExists: true })
+ */
+export async function dropPolicy(
+  pglite: PGliteLike,
+  policyName: string,
+  table: string,
+  options?: { ifExists?: boolean }
+): Promise<void> {
+  const sql = generateDropPolicySQL(policyName, table, options?.ifExists)
+  await pglite.exec(sql)
+}
+
+/**
+ * Drop all policies on a table.
+ *
+ * Removes all RLS policies from a table by querying pg_policies and dropping each one.
+ *
+ * SECURITY WARNING: This removes ALL access controls from the table. Use with extreme caution.
+ * Consider disabling RLS entirely if you want to remove all restrictions, or be prepared
+ * to immediately apply new policies.
+ *
+ * @param pglite - The PGlite instance to execute against
+ * @param table - The table to remove all policies from
+ *
+ * @example
+ * // WARNING: This removes all security policies!
+ * await dropAllPolicies(pglite, 'users')
+ */
+export async function dropAllPolicies(pglite: PGliteLike, table: string): Promise<void> {
+  // First get all policies
+  const policies = await listPoliciesInternal(pglite, table)
+
+  // Drop each policy
+  // Note: Policies are dropped sequentially to maintain predictable state
+  for (const policy of policies) {
+    await dropPolicy(pglite, policy.name, table)
+  }
+}
+
+/**
+ * List all policies on a table (internal helper).
+ *
+ * Queries pg_policies system view to retrieve policy information.
+ *
+ * SECURITY: Table name is escaped to prevent SQL injection.
+ *
+ * @param pglite - The PGlite instance to query
+ * @param table - The table to list policies for
+ * @returns Array of policy metadata objects
+ */
+async function listPoliciesInternal(
+  pglite: PGliteLike,
+  table: string
+): Promise<{ name: string; operation: string; using?: string; check?: string }[]> {
+  // SECURITY: Extract and escape the table name for safe use in the query
+  const tableName = escapeStringLiteral(extractTableName(table))
+
+  const result = await pglite.query<{
+    policyname: string
+    cmd: string
+    qual: string | null
+    with_check: string | null
+  }>(`
+    SELECT policyname, cmd, qual, with_check
+    FROM pg_policies
+    WHERE tablename = '${tableName}'
+  `)
+
+  return result.rows.map((row) => {
+    const policy: { name: string; operation: string; using?: string; check?: string } = {
+      name: row.policyname,
+      operation: row.cmd,
+    }
+    if (row.qual !== null) {
+      policy.using = row.qual
+    }
+    if (row.with_check !== null) {
+      policy.check = row.with_check
+    }
+    return policy
+  })
+}
+
+/**
+ * List all policies on a table.
+ *
+ * Queries the pg_policies system view to retrieve complete policy definitions.
+ *
+ * SECURITY: Table name is escaped to prevent SQL injection.
+ *
+ * @param pglite - The PGlite instance to query
+ * @param table - The table to list policies for
+ * @returns Array of complete RLSPolicy objects
+ *
+ * @example
+ * const policies = await listPolicies(pglite, 'users')
+ * console.log(policies) // [{ name: 'users_read_own', table: 'users', operation: 'SELECT', ... }]
+ */
+export async function listPolicies(pglite: PGliteLike, table: string): Promise<RLSPolicy[]> {
+  // SECURITY: Extract and escape the table name for safe use in the query
+  const tableName = escapeStringLiteral(extractTableName(table))
+
+  const result = await pglite.query<{
+    policyname: string
+    cmd: string
+    qual: string | null
+    with_check: string | null
+    roles: string[] | null
+    permissive: string
+  }>(`
+    SELECT policyname, cmd, qual, with_check, roles, permissive
+    FROM pg_policies
+    WHERE tablename = '${tableName}'
+  `)
+
+  return result.rows.map((row) => ({
+    name: row.policyname,
+    table,
+    operation: row.cmd as RLSPolicy['operation'],
+    using: row.qual ?? undefined,
+    check: row.with_check ?? undefined,
+    roles: row.roles ?? ['public'],
+    type: row.permissive === 'PERMISSIVE' ? 'permissive' : 'restrictive',
+  }))
+}
+
+/**
+ * Check if RLS is enabled on a table.
+ *
+ * Queries the pg_class system catalog to check if row-level security is enabled.
+ *
+ * SECURITY: Table name is escaped to prevent SQL injection.
+ *
+ * @param pglite - The PGlite instance to query
+ * @param table - The table to check
+ * @returns True if RLS is enabled, false otherwise
+ * @throws {Error} If the table does not exist
+ *
+ * @example
+ * const isEnabled = await checkRLSEnabled(pglite, 'users')
+ * if (!isEnabled) {
+ *   await enableRLS(pglite, 'users')
+ * }
+ */
+export async function checkRLSEnabled(pglite: PGliteLike, table: string): Promise<boolean> {
+  // SECURITY: Extract and escape the table name for safe use in the query
+  const tableName = escapeStringLiteral(extractTableName(table))
+
+  const result = await pglite.query<{ rowsecurity: boolean }>(`
+    SELECT relrowsecurity as rowsecurity
+    FROM pg_class
+    WHERE relname = '${tableName}'
+  `)
+
+  const row = result.rows[0]
+  if (!row) {
+    throw new NotFoundError('Table', extractTableName(table))
+  }
+
+  return row.rowsecurity
+}

package/src/rls/index.ts

@@ -0,0 +1,128 @@
+/**
+ * RLS (Row-Level Security) Abstraction Layer
+ *
+ * This module provides a type-safe DSL for creating and managing PostgreSQL
+ * Row-Level Security policies.
+ */
+
+// Re-export types from policy.ts
+export type {
+  PolicyOperation,
+  PolicyType,
+  TenantColumnType,
+  RLSPolicy,
+  CreatePolicyOptions,
+  EnableRLSOptions,
+  DisableRLSOptions,
+  ApplyPolicyOptions,
+  TenantPolicyOptions,
+  SetTenantContextOptions,
+  PolicyValidationResult,
+  ValidationOptions,
+  DropPolicyOptions,
+  PGliteLike,
+} from './policy'
+
+// Re-export functions from policy.ts
+export { createPolicy, createTenantPolicy } from './policy'
+
+// Re-export functions from validator.ts
+export { validatePolicy, validatePolicySyntax, validateTenantId } from './validator'
+
+// Re-export functions from generator.ts
+export {
+  generatePolicySQL,
+  enableRLS,
+  disableRLS,
+  applyPolicy,
+  dropPolicy,
+  dropAllPolicies,
+  listPolicies,
+  checkRLSEnabled,
+} from './generator'
+
+// Re-export auth integration types and classes
+export type {
+  RLSUser,
+  RLSContext,
+  SetupRLSFromAuthOptions,
+  AdminBypassOptions,
+  AdminOperationResult,
+} from './auth-integration'
+
+// Re-export auth integration functions
+export {
+  RLSContextManager,
+  createAdminBypassCheck,
+  createTenantPolicyWithBypass,
+  createOwnerPolicyWithBypass,
+  hasRole,
+  hasAnyRole,
+  hasAllRoles,
+} from './auth-integration'
+
+// Import types and functions for internal use
+import type { PGliteLike, TenantPolicyOptions, SetTenantContextOptions } from './policy'
+import { createTenantPolicy } from './policy'
+import { validateTenantId } from './validator'
+import { enableRLS, applyPolicy } from './generator'
+import { InvalidTenantIdError } from '@dotdo/postgres-shared/errors'
+
+/**
+ * Apply tenant policies to one or more tables
+ */
+export async function applyTenantPolicies(
+  pglite: PGliteLike,
+  tables: string | string[],
+  tenantColumn: string,
+  options?: TenantPolicyOptions
+): Promise<void> {
+  const tableList = Array.isArray(tables) ? tables : [tables]
+
+  for (const table of tableList) {
+    // Enable RLS on the table
+    await enableRLS(pglite, table)
+
+    // Create and apply tenant policy
+    const policy = createTenantPolicy(table, tenantColumn, options)
+
+    if (Array.isArray(policy)) {
+      for (const p of policy) {
+        await applyPolicy(pglite, p)
+      }
+    } else {
+      await applyPolicy(pglite, policy)
+    }
+  }
+}
+
+/**
+ * Set the tenant context in the current session
+ */
+export async function setTenantContext(
+  pglite: PGliteLike,
+  tenantId: string,
+  options?: SetTenantContextOptions
+): Promise<void> {
+  // Validate tenant ID to prevent injection
+  if (!validateTenantId(tenantId)) {
+    throw new InvalidTenantIdError(tenantId, 'Tenant ID must be alphanumeric with hyphens or underscores')
+  }
+
+  const sessionVariable = options?.sessionVariable ?? 'app.tenant_id'
+  const localKeyword = options?.local ? 'LOCAL ' : ''
+
+  await pglite.exec(`SET ${localKeyword}${sessionVariable} = '${tenantId}'`)
+}
+
+/**
+ * Clear the tenant context from the current session
+ */
+export async function clearTenantContext(
+  pglite: PGliteLike,
+  options?: SetTenantContextOptions
+): Promise<void> {
+  const sessionVariable = options?.sessionVariable ?? 'app.tenant_id'
+
+  await pglite.exec(`RESET ${sessionVariable}`)
+}