@o/database 0.16.18 → 0.16.19-canary.1790191946359

package/src/takeout/createServerHelpers.ts

@@ -0,0 +1,21 @@
+import { getDBClient, type GetDBClientOptions } from './getDBClient.js'
+import { createSql, setDefaultPool } from './sql.js'
+
+import type { Pool } from 'pg'
+
+export type ServerHelpers = {
+  sql: ReturnType<typeof createSql>
+  getDBClient: (
+    options?: Omit<GetDBClientOptions, 'pool' | 'connectionString'>
+  ) => ReturnType<typeof getDBClient>
+}
+
+export function createServerHelpers(pool: Pool): ServerHelpers {
+  const sql = createSql(pool)
+  setDefaultPool(pool)
+
+  return {
+    sql,
+    getDBClient: (options = {}) => getDBClient({ pool, ...options }),
+  }
+}

package/src/takeout/getDBClient.ts

@@ -0,0 +1,254 @@
+import pg, { type Pool, type PoolClient } from 'pg'
+
+// some of this file retry logic taken from:
+// https://github.com/brianc/node-postgres/issues/2718#issuecomment-1074019993
+
+export type GetDBClientOptions = {
+  pool?: Pool
+  connectionString?: string
+  retries?: number
+  onRetry?: (error: Error, attempt: number) => void
+}
+
+const cache = new Map<
+  string,
+  {
+    pool: pg.Pool
+    maxConnections: number | null
+    reservedConnections: number | null
+    openedConnections: number | null
+    openedConnectionsLastUpdate: number | null
+  }
+>()
+
+const createPoolKey = (connectionString: string) => connectionString
+
+const getOrCreatePoolCache = (connectionString: string, config: pg.PoolConfig) => {
+  const key = createPoolKey(connectionString)
+
+  if (!cache.has(key)) {
+    cache.set(key, {
+      pool: new pg.Pool(config),
+      maxConnections: null,
+      reservedConnections: null,
+      openedConnections: null,
+      openedConnectionsLastUpdate: null,
+    })
+  }
+
+  return cache.get(key)!
+}
+
+export async function getDBClient(options: GetDBClientOptions = {}): Promise<PoolClient> {
+  const { pool, connectionString, retries = 8 } = options
+
+  if (!pool && !connectionString) {
+    throw new Error('Either pool or connectionString must be provided')
+  }
+
+  let client: PoolClient | null = null
+
+  try {
+    client = await tryToGetNewClientFromPool(pool, connectionString, retries)
+    return client
+  } catch (error) {
+    console.error(`Failed to get DB client:`, error)
+    throw error
+  }
+}
+
+async function tryToGetNewClientFromPool(
+  providedPool: Pool | undefined,
+  connectionString: string | undefined,
+  retries: number
+): Promise<PoolClient> {
+  const { default: retry } = await import('async-retry')
+  const clientFromPool = await retry(
+    async () => {
+      if (providedPool) {
+        const client = await providedPool.connect()
+        return client
+      }
+
+      if (!connectionString) {
+        throw new Error('No connection string provided')
+      }
+
+      const configurations: pg.PoolConfig = {
+        connectionString,
+        connectionTimeoutMillis: 5_000,
+        // idle_session_timeout set to 35s on server, client timeout at 30s
+        // fix via https://github.com/brianc/node-postgres/issues/2718#issuecomment-2094885323
+        idleTimeoutMillis: 30_000,
+        allowExitOnIdle: true,
+      }
+
+      const poolCache = getOrCreatePoolCache(connectionString, configurations)
+
+      const client = await poolCache.pool.connect()
+      return client
+    },
+    {
+      retries,
+      minTimeout: 300,
+      factor: 2,
+      maxTimeout: 8000,
+    }
+  )
+
+  return clientFromPool
+}
+
+export async function queryDb(
+  queryText: string,
+  params?: any[],
+  options: GetDBClientOptions = {}
+): Promise<pg.QueryResult> {
+  let client: PoolClient | null = null
+
+  try {
+    client = await tryToGetNewClientFromPool(
+      options.pool,
+      options.connectionString,
+      options.retries || 8
+    )
+    return await client.query(queryText, params)
+  } catch (error) {
+    console.error(`Database query failed:`, {
+      query: queryText,
+      error: error instanceof Error ? error.message : String(error),
+    })
+    throw error
+  } finally {
+    if (client && options.connectionString) {
+      const tooManyConnections = await checkForTooManyConnections(
+        client,
+        options.connectionString
+      )
+
+      if (tooManyConnections) {
+        const poolCache = cache.get(createPoolKey(options.connectionString))
+        client.release()
+        await poolCache?.pool.end()
+        if (poolCache) {
+          cache.delete(createPoolKey(options.connectionString))
+        }
+      } else {
+        client.release()
+      }
+    } else if (client) {
+      client.release()
+    }
+  }
+}
+
+async function checkForTooManyConnections(
+  client: PoolClient,
+  connectionString: string
+): Promise<boolean> {
+  const poolCache = cache.get(createPoolKey(connectionString))
+  if (!poolCache) return false
+
+  const currentTime = Date.now()
+  const openedConnectionsMaxAge = 10000
+  const maxConnectionsTolerance = 0.9
+
+  if (poolCache.maxConnections === null || poolCache.reservedConnections === null) {
+    const [maxConnections, reservedConnections] = await getConnectionLimits(client)
+    poolCache.maxConnections = maxConnections
+    poolCache.reservedConnections = reservedConnections
+  }
+
+  if (
+    poolCache.openedConnections === null ||
+    poolCache.openedConnectionsLastUpdate === null ||
+    currentTime - poolCache.openedConnectionsLastUpdate > openedConnectionsMaxAge
+  ) {
+    const openedConnections = await getOpenedConnections(client, connectionString)
+    poolCache.openedConnections = openedConnections
+    poolCache.openedConnectionsLastUpdate = currentTime
+  }
+
+  if (
+    poolCache.openedConnections >
+    (poolCache.maxConnections - poolCache.reservedConnections) * maxConnectionsTolerance
+  ) {
+    console.warn(
+      `Too many connections detected: ${poolCache.openedConnections}/${poolCache.maxConnections - poolCache.reservedConnections}`
+    )
+    return true
+  }
+
+  return false
+}
+
+async function getConnectionLimits(client: PoolClient): Promise<[number, number]> {
+  const maxConnectionsResult = await client.query('SHOW max_connections')
+  const reservedConnectionResult = await client.query(
+    'SHOW superuser_reserved_connections'
+  )
+
+  return [
+    Number.parseInt(maxConnectionsResult.rows[0].max_connections, 10),
+    Number.parseInt(reservedConnectionResult.rows[0].superuser_reserved_connections, 10),
+  ]
+}
+
+async function getOpenedConnections(
+  client: PoolClient,
+  connectionString: string
+): Promise<number> {
+  // For Aurora/RDS, we need to get the database name from connection string
+  const dbName = new URL(connectionString).pathname.slice(1)
+  const openConnectionsResult = await client.query(
+    'SELECT numbackends as opened_connections FROM pg_stat_database WHERE datname = $1',
+    [dbName]
+  )
+  const result = Number.parseInt(
+    openConnectionsResult.rows[0]?.opened_connections || 0,
+    10
+  )
+  return result
+}
+
+export async function getNewClient(options: GetDBClientOptions = {}): Promise<pg.Client> {
+  const { connectionString } = options
+
+  if (!connectionString) {
+    throw new Error('connectionString is required for getNewClient')
+  }
+
+  try {
+    const client = await tryToGetNewClient(connectionString)
+    return client
+  } catch (error) {
+    console.error(`Failed to get new client:`, error)
+    throw error
+  }
+}
+
+async function tryToGetNewClient(connectionString: string): Promise<pg.Client> {
+  const configurations: pg.PoolConfig = {
+    connectionString,
+    connectionTimeoutMillis: 5_000,
+    idleTimeoutMillis: 30_000,
+    allowExitOnIdle: true,
+  }
+
+  const { default: retry } = await import('async-retry')
+  const client = await retry(
+    async () => {
+      const newClient = new pg.Client(configurations)
+      await newClient.connect()
+      return newClient
+    },
+    {
+      retries: 10,
+      minTimeout: 100,
+      factor: 2,
+      maxTimeout: 5000,
+    }
+  )
+
+  return client
+}

package/src/takeout/index.ts

@@ -0,0 +1,38 @@
+export {
+  createDatabase,
+  type DatabaseConfig,
+  type MigrateOptions,
+} from './createDatabase.js'
+export { createPool, type CreatePoolOptions } from './createPool.js'
+export {
+  createDrizzle,
+  type CreateDrizzleOptions,
+  type DrizzleHelpers,
+} from './createDrizzle.js'
+export { createSql, type SqlQuery } from './sql.js'
+export { createServerHelpers, type ServerHelpers } from './createServerHelpers.js'
+export { getDBClient, queryDb, type GetDBClientOptions } from './getDBClient.js'
+export { processInChunks, updateInChunks } from './chunkedQuery.js'
+export { createDb } from './createDb.js'
+export { migrate, loadMigrationsFromDir, type Migration } from './migrate.js'
+export { waitForDatabase } from './waitForDatabase.js'
+export {
+  buildInitSql,
+  defineSeed,
+  defineTrigger,
+  derivePublicTables,
+  ensureZeroDatabases,
+  ensureZeroPublication,
+  ensureZeroSeeds,
+  ensureZeroTriggers,
+  isPrivateTable,
+  loadSeeds,
+  loadTriggers,
+  privateTable,
+  type InitSqlOptions,
+  type SeedDefinition,
+  type TriggerDefinition,
+  type ZeroDatabaseOptions,
+  type ZeroPublicationOptions,
+} from './zero.js'
+export { applySeed, type SeedClient, type SeedData, type SeedRows } from './seed.js'

package/src/takeout/migrate.ts

@@ -0,0 +1,272 @@
+/**
+ * Custom migration script - why?
+ * Well migration scripts aren't so complex, and we haven't found a simple enough library
+ * to deserve a new dependency. We like this setup because we can customize it easily
+ * and it's pretty easy to understand.
+ *
+ * This migration system handles both custom TypeScript migrations and
+ * SQL migrations generated by Drizzle.
+ */
+
+import { readdirSync, readFileSync } from 'node:fs'
+import { basename, resolve } from 'node:path'
+
+import { getDBClient } from './getDBClient.js'
+import { ensureZeroDatabases } from './zero.js'
+
+import type { PoolClient } from 'pg'
+
+export type Migration = {
+  name: string
+  up?: (client: PoolClient) => Promise<void>
+  /** per-migration timeout in ms, overrides defaultTimeout */
+  timeout?: number
+}
+
+/** default timeout for migrations (5 seconds) */
+const DEFAULT_MIGRATION_TIMEOUT = 5_000
+
+export type MigrateOptions = {
+  connectionString: string
+  migrations: Record<string, () => Promise<unknown>>
+  createDatabases?: string[]
+  onMigrationComplete?: () => Promise<void>
+  gitSha?: string
+  cvrDb?: string
+  changeDb?: string
+  /** call process.exit(0) after migrations complete (default: false) */
+  forceExit?: boolean
+  /** default timeout for all migrations in ms (default: 5000) */
+  defaultTimeout?: number
+}
+
+async function withTimeout<T>(promise: Promise<T>, ms: number, name: string): Promise<T> {
+  let timeoutId: ReturnType<typeof setTimeout>
+  const timeoutPromise = new Promise<never>((_, reject) => {
+    timeoutId = setTimeout(() => {
+      reject(new Error(`migration "${name}" timed out after ${ms}ms`))
+    }, ms)
+  })
+  try {
+    return await Promise.race([promise, timeoutPromise])
+  } finally {
+    clearTimeout(timeoutId!)
+  }
+}
+
+export async function migrate(options: MigrateOptions) {
+  const {
+    connectionString,
+    migrations,
+    createDatabases = [],
+    onMigrationComplete,
+    forceExit = false,
+    gitSha,
+    cvrDb,
+    changeDb,
+    defaultTimeout = DEFAULT_MIGRATION_TIMEOUT,
+  } = options
+
+  // create zero + additional databases before running migrations
+  await ensureZeroDatabases({ connectionString, cvrDb, changeDb })
+  for (const dbUrl of createDatabases) {
+    await ensureZeroDatabases({ connectionString, cvrDb: dbUrl })
+  }
+
+  const client = await getDBClient({ connectionString })
+
+  try {
+    await client.query('BEGIN')
+    // force public schema - if a schema matching the db user exists (e.g. startchat),
+    // postgres resolves unqualified names there first, causing tables to land in the
+    // wrong schema. setting search_path ensures all CREATE TABLE etc go to public.
+    await client.query(`SET search_path = public`)
+    await client.query(`
+      CREATE TABLE IF NOT EXISTS public.migrations (
+        id SERIAL PRIMARY KEY,
+        name VARCHAR(255) NOT NULL,
+        run_on TIMESTAMP NOT NULL DEFAULT NOW()
+      )
+    `)
+
+    const appliedMigrations = await client.query('SELECT name FROM public.migrations')
+    const appliedMigrationNames = new Set(appliedMigrations.rows.map((row) => row.name))
+
+    // get TypeScript migrations
+    const tsMigrationsSorted = Object.entries(migrations)
+      .sort(([a], [b]) => a.localeCompare(b))
+      .map(([file, run]) => ({
+        name: basename(file).replace(/\.(ts|sql)$/, ''),
+        run,
+      }))
+      .filter(({ name }) => /^[\d]+/.test(name))
+
+    // process TypeScript migrations
+    const tsMigrations: Migration[] = await Promise.all(
+      tsMigrationsSorted.map(async ({ name, run }) => {
+        if (appliedMigrationNames.has(name)) {
+          return null
+        }
+        try {
+          const migration = (await run()) as object
+          return { ...migration, name }
+        } catch (error) {
+          console.error(`Failed to load TypeScript migration ${name}:`, error)
+          throw error
+        }
+      })
+    ).then((migrations) => migrations.filter(Boolean) as Migration[])
+
+    // combine TypeScript and SQL migrations, sorted by name to ensure correct order
+    const pendingMigrations = [...tsMigrations].sort((a, b) =>
+      a.name.localeCompare(b.name)
+    )
+
+    if (!pendingMigrations.length) {
+      await client.query('COMMIT')
+    } else {
+      for (const migration of pendingMigrations) {
+        if (migration.up) {
+          // use savepoints so "already exists" errors don't abort the transaction.
+          // this handles cases where tables/columns exist but migration tracking
+          // was lost (e.g. db restore, pglite data loss)
+          const sp = `sp_${migration.name.replace(/\W/g, '_')}`
+          await client.query(`SAVEPOINT ${sp}`)
+          try {
+            const timeout = migration.timeout ?? defaultTimeout
+            await withTimeout(migration.up(client), timeout, migration.name)
+            await client.query(`RELEASE SAVEPOINT ${sp}`)
+          } catch (e: any) {
+            // ROLLBACK TO SAVEPOINT can fail in its own right (e.g. when running
+            // against pglite via orez's wire-protocol proxy, a parallel client
+            // sharing the same backend session can corrupt the transaction
+            // state, leaving us with no active transaction). if we don't catch
+            // the rollback failure separately, the original migration error `e`
+            // is replaced by the rollback's error, hiding what actually went
+            // wrong. log the rollback failure and re-throw the original.
+            try {
+              await client.query(`ROLLBACK TO SAVEPOINT ${sp}`)
+            } catch (rollbackErr) {
+              console.error(
+                `[migrate] ${migration.name}: ROLLBACK TO SAVEPOINT failed:`,
+                rollbackErr
+              )
+              console.error(
+                `[migrate] ${migration.name}: original migration error was:`,
+                e
+              )
+              throw e
+            }
+            // only ignore idempotency errors - NOT all 42xxx errors!
+            // 42xxx is the entire "Syntax Error or Access Rule Violation" class
+            // which includes undefined_table (42P01), syntax_error (42601), etc.
+            // we only want to skip migrations where objects already/don't exist:
+            const idempotentErrorCodes = [
+              '42710', // duplicate_object
+              '42P07', // duplicate_table
+              '42701', // duplicate_column
+              '42P16', // invalid_table_definition (e.g. constraint already exists)
+              '42704', // undefined_object (e.g. DROP CONSTRAINT on non-existent constraint)
+            ]
+            if (idempotentErrorCodes.includes(e?.code)) {
+              console.info(
+                `[migrate] ${migration.name}: ${e.message} (${e.code}), recording as applied`
+              )
+            } else {
+              // any other error should fail the migration
+              throw e
+            }
+          }
+        }
+
+        await client.query('INSERT INTO public.migrations (name) VALUES ($1)', [
+          migration.name,
+        ])
+      }
+
+      await client.query('COMMIT')
+    }
+  } catch (e) {
+    console.error(`Migration failed, rolling back:`, e)
+    await client.query('ROLLBACK')
+    try {
+      client.release(false)
+    } catch (releaseErr) {
+      console.error(`Error releasing connection after rollback:`, releaseErr)
+    }
+    throw e
+  }
+
+  if (onMigrationComplete) {
+    await onMigrationComplete()
+  }
+
+  try {
+    // don't destroy the connection forcefully - let it return to pool gracefully
+    client.release(false)
+  } catch (err) {
+    console.error(`Error releasing connection gracefully, trying to destroy:`, err)
+    try {
+      client.release(true)
+    } catch (destroyErr) {
+      console.error(`Error destroying connection:`, destroyErr)
+    }
+  }
+
+  if (forceExit) {
+    // flush stdio before exiting — process.exit() doesn't wait for streams to drain
+    await new Promise<void>((resolve) => {
+      process.stderr.write('', () => {
+        process.stdout.write('', () => resolve())
+      })
+    })
+    process.exit(0)
+  }
+}
+
+/** load migration files from a directory (for non-vite contexts like bun scripts) */
+export function loadMigrationsFromDir(
+  dir: string
+): Record<string, () => Promise<unknown>> {
+  const files = readdirSync(dir)
+    .filter((f) => /^\d+/.test(f) && (f.endsWith('.ts') || f.endsWith('.sql')))
+    .sort()
+
+  // deduplicate: prefer .ts over .sql for same base name
+  const seen = new Set<string>()
+  const deduped: string[] = []
+  for (const f of files) {
+    if (f.endsWith('.ts')) {
+      seen.add(f.replace(/\.ts$/, ''))
+      deduped.push(f)
+    }
+  }
+  for (const f of files) {
+    if (f.endsWith('.sql') && !seen.has(f.replace(/\.sql$/, ''))) {
+      deduped.push(f)
+    }
+  }
+  deduped.sort()
+
+  return Object.fromEntries(
+    deduped.map((f) => {
+      if (f.endsWith('.sql')) {
+        return [
+          `./${f}`,
+          async () => {
+            const sql = readFileSync(resolve(dir, f), 'utf-8')
+            return {
+              up: async (client: PoolClient) => {
+                for (const stmt of sql.split('--> statement-breakpoint')) {
+                  const trimmed = stmt.trim()
+                  if (trimmed) await client.query(trimmed)
+                }
+              },
+            }
+          },
+        ]
+      }
+      return [`./${f}`, () => import(/* @vite-ignore */ resolve(dir, f))]
+    })
+  )
+}

package/src/takeout/scripts/build-migrations.ts

@@ -0,0 +1,66 @@
+import { execSync } from 'node:child_process'
+import path from 'node:path'
+
+export type BuildMigrationsOptions = {
+  migrationsDir: string
+  outFile?: string
+  target?: string
+  aliases?: Record<string, string>
+}
+
+export async function buildMigrations(options: BuildMigrationsOptions) {
+  const {
+    migrationsDir,
+    outFile = 'migrate-dist.js',
+    target = 'node22',
+    aliases = {},
+  } = options
+
+  const { build } = await import('vite')
+
+  const migrateFile = path.join(migrationsDir, '..', 'migrate.ts')
+
+  const result = await build({
+    configFile: false,
+    resolve: {
+      alias: aliases,
+    },
+    define: {
+      'process.env.GIT_SHA': JSON.stringify(
+        execSync('git rev-parse HEAD').toString().trim()
+      ),
+    },
+    build: {
+      outDir: path.dirname(migrateFile),
+      target,
+      minify: false,
+      emptyOutDir: false,
+      copyPublicDir: false,
+      lib: {
+        name: 'migrate',
+        formats: ['es'],
+        entry: migrateFile,
+      },
+      rollupOptions: {
+        external: (id) => {
+          // externalize all node modules and node: imports
+          if (id.startsWith('node:') || id === 'pg') return true
+          // externalize all absolute paths (like aliases)
+          if (id.startsWith('/')) return false
+          // externalize node_modules
+          if (!id.startsWith('.') && !id.startsWith('/')) return true
+          return false
+        },
+        output: {
+          format: 'es',
+          inlineDynamicImports: true,
+          exports: 'named',
+          entryFileNames: outFile,
+        },
+      },
+    },
+  })
+
+  console.info(`✓ Built migration bundle: ${outFile}`)
+  return result
+}