@o/database 0.16.18 → 0.16.19-canary.1790191946359

package/src/takeout/better-auth.ts

@@ -0,0 +1,110 @@
+import { drizzleAdapter, type DrizzleAdapterConfig } from 'better-auth/adapters/drizzle'
+
+import {
+  createSQLiteDrizzle,
+  type CreateSQLiteDatabaseOptions,
+  type SQLiteSchema,
+  type SQLiteTransactionExecutor,
+  type SQLiteTransactionProvider,
+} from './sqlite.js'
+
+import type { AnyRelations, EmptyRelations } from 'drizzle-orm'
+
+export type CreateBetterAuthSQLiteAdapterOptions<
+  TSchema extends SQLiteSchema,
+  TRelations extends AnyRelations = EmptyRelations,
+> = Pick<
+  CreateSQLiteDatabaseOptions<TSchema, TRelations>,
+  'relations' | 'schema' | 'transactionProvider'
+> & {
+  /**
+   * runs the adapter's read-only operations (`findOne`, `findMany`, `count`).
+   * the SQLite host decides what a read session means; on a host that admits
+   * one writer at a time and readers together, a session lookup no longer
+   * waits for the writer turn. defaults to `transactionProvider`.
+   */
+  readTransactionProvider?: SQLiteTransactionProvider
+} & Omit<DrizzleAdapterConfig, 'provider' | 'schema' | 'transaction'>
+
+// the adapter operations that only SELECT. every other operation, and the
+// explicit `transaction`, keeps the write provider.
+const READ_OPERATIONS: ReadonlySet<string> = new Set(['findOne', 'findMany', 'count'])
+
+/**
+ * gives Better Auth its official Drizzle adapter inside the transaction
+ * callback owned by the SQLite host. the Drizzle adapter's transaction mode
+ * stays disabled because it must never emit its own BEGIN or COMMIT.
+ */
+export function createBetterAuthSQLiteAdapter<
+  TSchema extends SQLiteSchema,
+  TRelations extends AnyRelations = EmptyRelations,
+>(
+  options: CreateBetterAuthSQLiteAdapterOptions<TSchema, TRelations>
+): ReturnType<typeof drizzleAdapter> {
+  const {
+    relations,
+    schema,
+    transactionProvider,
+    readTransactionProvider = transactionProvider,
+    ...drizzleOptions
+  } = options
+  const adapterConfig = {
+    ...drizzleOptions,
+    provider: 'sqlite',
+    schema,
+    transaction: false,
+  } satisfies DrizzleAdapterConfig
+  const createAdapter = (
+    executor: SQLiteTransactionExecutor,
+    authOptions: Parameters<ReturnType<typeof drizzleAdapter>>[0]
+  ) =>
+    drizzleAdapter(
+      createSQLiteDrizzle(executor, { relations, schema }),
+      adapterConfig
+    )(authOptions)
+
+  return (authOptions) => {
+    const unavailableExecutor: SQLiteTransactionExecutor = {
+      exec() {
+        throw new Error('Better Auth SQLite operations require a transaction callback')
+      },
+      execMany() {
+        throw new Error('Better Auth SQLite operations require a transaction callback')
+      },
+      query() {
+        throw new Error('Better Auth SQLite operations require a transaction callback')
+      },
+      queryAst() {
+        throw new Error('Better Auth SQLite operations require a transaction callback')
+      },
+    }
+    const target = createAdapter(unavailableExecutor, authOptions)
+    const transaction: typeof target.transaction = (work) =>
+      transactionProvider(async (executor) => work(createAdapter(executor, authOptions)))
+
+    return new Proxy(target, {
+      get(_target, property, receiver) {
+        if (property === 'transaction') return transaction
+
+        const value = Reflect.get(target, property, receiver)
+        if (typeof value !== 'function') return value
+
+        const provider =
+          typeof property === 'string' && READ_OPERATIONS.has(property)
+            ? readTransactionProvider
+            : transactionProvider
+        return (...args: unknown[]) =>
+          provider(async (executor) => {
+            const adapter = createAdapter(executor, authOptions)
+            const operation = Reflect.get(adapter, property)
+            if (typeof operation !== 'function') {
+              throw new TypeError(
+                `Better Auth adapter operation ${String(property)} is unavailable`
+              )
+            }
+            return Reflect.apply(operation, adapter, args)
+          })
+      },
+    })
+  }
+}

package/src/takeout/build.ts

@@ -0,0 +1,2 @@
+// build utilities that depend on vite
+export { buildMigrations } from './scripts/build-migrations.js'

package/src/takeout/chunkedQuery.ts

@@ -0,0 +1,91 @@
+import type { PoolClient, QueryResultRow } from 'pg'
+
+interface ChunkedQueryOptions {
+  chunkSize?: number
+  onProgress?: (processed: number, total: number) => void
+}
+
+/**
+ * Process database records in chunks to avoid memory issues with large datasets
+ */
+export async function processInChunks<T extends QueryResultRow = QueryResultRow>(
+  client: PoolClient,
+  query: string,
+  processor: (rows: T[]) => Promise<void>,
+  options: ChunkedQueryOptions = {}
+): Promise<void> {
+  const { chunkSize = 1000, onProgress } = options
+  let offset = 0
+  let hasMore = true
+  let totalProcessed = 0
+
+  // first get total count for progress reporting
+  const countQuery = query
+    .replace(/SELECT .+ FROM/, 'SELECT COUNT(*) FROM')
+    .replace(/ORDER BY .+/, '')
+  const countResult = await client.query(countQuery)
+  const totalCount = Number.parseInt(countResult.rows[0].count, 10)
+
+  while (hasMore) {
+    const paginatedQuery = `${query} LIMIT ${chunkSize} OFFSET ${offset}`
+    const result = await client.query<T>(paginatedQuery)
+
+    if (result.rows.length === 0) {
+      hasMore = false
+      break
+    }
+
+    await processor(result.rows)
+
+    totalProcessed += result.rows.length
+    offset += chunkSize
+
+    if (onProgress) {
+      onProgress(totalProcessed, totalCount)
+    }
+
+    // check if we've processed all records
+    if (result.rows.length < chunkSize) {
+      hasMore = false
+    }
+  }
+}
+
+/**
+ * Update records in chunks with a transformer function
+ */
+export async function updateInChunks<T extends QueryResultRow & { id: string }>(
+  client: PoolClient,
+  tableName: string,
+  selectQuery: string,
+  transformer: (row: T) => Promise<Partial<T> | null>,
+  options: ChunkedQueryOptions = {}
+): Promise<number> {
+  let totalUpdated = 0
+
+  await processInChunks<T>(
+    client,
+    selectQuery,
+    async (rows) => {
+      for (const row of rows) {
+        const updates = await transformer(row)
+
+        if (updates && Object.keys(updates).length > 0) {
+          // build update query dynamically
+          const setClause = Object.keys(updates)
+            .map((key, index) => `${key} = $${index + 2}`)
+            .join(', ')
+
+          const values = [row.id, ...Object.values(updates)]
+
+          await client.query(`UPDATE ${tableName} SET ${setClause} WHERE id = $1`, values)
+
+          totalUpdated++
+        }
+      }
+    },
+    options
+  )
+
+  return totalUpdated
+}

package/src/takeout/cli.ts

@@ -0,0 +1,126 @@
+#!/usr/bin/env node
+import { join } from 'node:path'
+
+import { defineCommand, runMain } from 'citty'
+
+import { buildMigrations as buildMigrationsRun } from './scripts/build-migrations.js'
+import { syncDrizzleMigrations } from './scripts/drizzle-migrations-sync.js'
+import { addMigration } from './scripts/migration-add.js'
+import { runPgDump } from './scripts/pg_dump.js'
+import { runPsql } from './scripts/psql.js'
+
+const syncDrizzle = defineCommand({
+  meta: {
+    name: 'sync-drizzle',
+    description: 'Sync Drizzle SQL migrations to TypeScript wrappers',
+  },
+  args: {
+    dir: {
+      type: 'string',
+      description: 'Migrations directory',
+      required: false,
+      default: './src/database/migrations',
+    },
+  },
+  async run({ args }) {
+    const migrationsDir = join(process.cwd(), args.dir)
+    console.info(`Syncing migrations in ${migrationsDir}`)
+    await syncDrizzleMigrations({ migrationsDir })
+  },
+})
+
+const migrationAdd = defineCommand({
+  meta: {
+    name: 'migrate:add',
+    description: 'Create a new custom TypeScript migration',
+  },
+  args: {
+    name: {
+      type: 'positional',
+      description: 'Migration name',
+      required: false,
+    },
+    dir: {
+      type: 'string',
+      description: 'Migrations directory',
+      required: false,
+      default: './src/database/migrations',
+    },
+  },
+  async run({ args }) {
+    const migrationsDir = join(process.cwd(), args.dir)
+    addMigration({ migrationsDir, name: args.name })
+  },
+})
+
+const psql = defineCommand({
+  meta: {
+    name: 'psql',
+    description: 'Connect to PostgreSQL database with psql',
+  },
+  args: {
+    connectionString: {
+      type: 'string',
+      description: 'PostgreSQL connection string',
+      required: false,
+    },
+    query: {
+      type: 'string',
+      description: 'Query to execute',
+      required: false,
+    },
+  },
+  async run({ args }) {
+    const connectionString = args.connectionString || process.env.ZERO_UPSTREAM_DB
+    if (!connectionString) {
+      console.error(
+        'No connection string provided. Set ZERO_UPSTREAM_DB or pass --connectionString'
+      )
+      process.exit(1)
+    }
+    const exitCode = runPsql({ connectionString, query: args.query })
+    process.exit(exitCode || 0)
+  },
+})
+
+const pgDump = defineCommand({
+  meta: {
+    name: 'pg_dump',
+    description: 'Dump PostgreSQL database using pg_dump',
+  },
+  args: {
+    connectionString: {
+      type: 'string',
+      description: 'PostgreSQL connection string',
+      required: false,
+    },
+  },
+  async run({ args }) {
+    const connectionString = args.connectionString || process.env.ZERO_UPSTREAM_DB
+    if (!connectionString) {
+      console.error(
+        'No connection string provided. Set ZERO_UPSTREAM_DB or pass --connectionString'
+      )
+      process.exit(1)
+    }
+    const cliArgs = process.argv.slice(3) // get args after command name
+    const exitCode = runPgDump({ connectionString, args: cliArgs })
+    process.exit(exitCode || 0)
+  },
+})
+
+const main = defineCommand({
+  meta: {
+    name: 'database',
+    description: 'Database utilities and migration tools',
+    version: '0.0.1',
+  },
+  subCommands: {
+    'sync-drizzle': syncDrizzle,
+    'migrate:add': migrationAdd,
+    psql,
+    pg_dump: pgDump,
+  },
+})
+
+runMain(main)

package/src/takeout/createDatabase.ts

@@ -0,0 +1,164 @@
+import { getTableName } from 'drizzle-orm'
+import { PgTable } from 'drizzle-orm/pg-core'
+
+import { createDrizzle } from './createDrizzle.js'
+import { createPool, type CreatePoolOptions } from './createPool.js'
+import { createServerHelpers } from './createServerHelpers.js'
+import { getDBClient } from './getDBClient.js'
+import { migrate } from './migrate.js'
+import { waitForDatabase } from './waitForDatabase.js'
+import {
+  ensureZeroDatabases,
+  ensureZeroPublication,
+  ensureZeroSeeds,
+  ensureZeroTriggers,
+  type SeedDefinition,
+  type TriggerDefinition,
+} from './zero.js'
+
+import type { AnyRelations, EmptyRelations } from 'drizzle-orm'
+import type { Pool } from 'pg'
+
+export type DatabaseConfig<
+  TSchema extends Record<string, unknown>,
+  TRelations extends AnyRelations | Record<string, unknown> = EmptyRelations,
+> = {
+  /**
+   * The PG connection string. Required unless a pre-built `pool` is provided
+   * (it is still used by `migrate()` / `waitForDatabase()`, which need a string).
+   */
+  connectionString?: string
+  schema: TSchema
+  relations?: TRelations
+  /**
+   * Optional pre-built Pool-like client (e.g. a DoBackend-backed Pool for the
+   * Cloudflare Durable Object runtime, where there is no external Postgres).
+   * When provided, `createPool` is skipped and this pool is used directly.
+   * `connectionString` is then only needed for the migration helpers.
+   */
+  pool?: Pool
+  /** Options forwarded to the default `createPool` (ignored when `pool` is set). */
+  poolOptions?: Omit<CreatePoolOptions, 'connectionString'>
+  zero?: {
+    publicationName: string
+    publicTables: string[]
+    triggers?: TriggerDefinition[]
+    seeds?: SeedDefinition[]
+  }
+}
+
+export type MigrateOptions = {
+  migrations: Record<string, () => Promise<unknown>>
+  cvrDb?: string
+  changeDb?: string
+  createDatabases?: string[]
+  gitSha?: string
+  defaultTimeout?: number
+  onMigrationComplete?: () => Promise<void>
+}
+
+function stripQueryParams(connStr: string | undefined): string | undefined {
+  if (!connStr) return connStr
+  return connStr.split('?')[0]
+}
+
+export function createDatabase<
+  TSchema extends Record<string, unknown>,
+  TRelations extends AnyRelations = EmptyRelations,
+>(config: DatabaseConfig<TSchema, TRelations>) {
+  const { connectionString, schema, relations, zero } = config
+
+  // use a caller-provided pool (e.g. the Cloudflare DO-backed pool) when given;
+  // otherwise build the default real-pg pool from the connection string.
+  const pool =
+    config.pool ??
+    createPool({
+      connectionString: connectionString ?? '',
+      ...config.poolOptions,
+    })
+
+  const drizzle = createDrizzle({
+    pool,
+    schema,
+    relations: relations as TRelations & AnyRelations,
+  })
+
+  const { sql, getDBClient: getBoundDBClient } = createServerHelpers(pool)
+
+  const allTableNames = Object.values(schema)
+    .filter((v): v is PgTable => v instanceof PgTable)
+    .map((t) => getTableName(t))
+
+  const publicTableSet = new Set(zero?.publicTables || [])
+  const privateTableNames = allTableNames.filter((n) => !publicTableSet.has(n))
+
+  async function close() {
+    try {
+      await pool.end()
+    } catch (e) {
+      console.error('[database] error closing:', e)
+    }
+  }
+
+  async function runMigrate(options: MigrateOptions) {
+    if (!connectionString) {
+      throw new Error(
+        'createDatabase: migrate() requires a `connectionString` (it is not derivable from a pre-built `pool`)'
+      )
+    }
+    console.info('[database] waiting for database to be ready...')
+    await waitForDatabase(connectionString)
+
+    console.info('[database] running migrations...')
+    await migrate({
+      connectionString,
+      migrations: options.migrations,
+      cvrDb: stripQueryParams(options.cvrDb),
+      changeDb: stripQueryParams(options.changeDb),
+      createDatabases: options.createDatabases,
+      gitSha: options.gitSha,
+      defaultTimeout: options.defaultTimeout,
+      onMigrationComplete: async () => {
+        if (zero) {
+          await ensureZeroPublication({
+            connectionString,
+            publicationName: zero.publicationName,
+            privateTableNames,
+          })
+
+          if (zero.triggers?.length || zero.seeds?.length) {
+            const client = await getDBClient({ connectionString })
+            try {
+              if (zero.triggers?.length) {
+                await ensureZeroTriggers(client, zero.triggers)
+              }
+              if (zero.seeds?.length) {
+                await ensureZeroSeeds(client, zero.seeds)
+              }
+            } finally {
+              client.release()
+            }
+          }
+        }
+
+        if (options.onMigrationComplete) {
+          await options.onMigrationComplete()
+        }
+      },
+    })
+    console.info('[database] migrations complete')
+    await close()
+    process.exit(0)
+  }
+
+  return {
+    pool,
+    drizzle,
+    sql,
+    getDBClient: getBoundDBClient,
+    close,
+    migrate: runMigrate,
+    privateTableNames,
+    config,
+  }
+}

package/src/takeout/createDb.ts

@@ -0,0 +1,15 @@
+import { type NodePgDatabase, drizzle } from 'drizzle-orm/node-postgres'
+
+import { createPool } from './createPool.js'
+
+export const createDb = <TSchema extends Record<string, unknown>>(
+  connectionString: string,
+  schema: TSchema
+): NodePgDatabase<TSchema> => {
+  const pool = createPool({ connectionString })
+  return drizzle({
+    client: pool,
+    schema: schema as Record<string, unknown>,
+    logger: false,
+  } as any) as unknown as NodePgDatabase<TSchema>
+}

package/src/takeout/createDrizzle.ts

@@ -0,0 +1,127 @@
+import { and, eq, gt, gte, lt, lte, ne } from 'drizzle-orm'
+import { type NodePgDatabase, drizzle } from 'drizzle-orm/node-postgres'
+import { PgTimestampString } from 'drizzle-orm/pg-core'
+
+import type { AnyRelations, EmptyRelations, InferInsertModel, SQL } from 'drizzle-orm'
+import type { PgTable } from 'drizzle-orm/pg-core'
+import type { Pool } from 'pg'
+
+let patched = false
+function patchTimestampForZero() {
+  if (patched) return
+  patched = true
+  const orig = PgTimestampString.prototype.mapToDriverValue
+  PgTimestampString.prototype.mapToDriverValue = function (value: any) {
+    if (typeof value === 'number') return new Date(value).toISOString()
+    return orig.call(this, value)
+  }
+}
+
+const operators = { eq, ne, lt, lte, gt, gte } as const
+type Op = keyof typeof operators
+type WhereOp = Partial<Record<Op, any>>
+type WhereValue = string | number | boolean | null | WhereOp
+type WithSQL<T> = { [K in keyof T]: T[K] | SQL }
+
+type TableName<TSchema> = {
+  [K in keyof TSchema]: TSchema[K] extends PgTable ? K : never
+}[keyof TSchema]
+
+type WhereClause<TSchema, T extends TableName<TSchema>> = Partial<
+  Record<
+    keyof InferInsertModel<TSchema[T] extends PgTable ? TSchema[T] : never> & string,
+    WhereValue
+  >
+>
+
+type InsertMap<TSchema> = {
+  [K in TableName<TSchema>]: Partial<
+    WithSQL<InferInsertModel<TSchema[K] extends PgTable ? TSchema[K] : never>>
+  >
+}
+
+function resolveWhere(table: any, where: Record<string, WhereValue>) {
+  const conditions = Object.entries(where).map(([key, value]) => {
+    if (value !== null && typeof value === 'object') {
+      const entries = Object.entries(value) as [Op, any][]
+      if (entries.length === 1) {
+        return operators[entries[0]![0]](table[key], entries[0]![1])
+      }
+      return and(...entries.map(([op, v]) => operators[op](table[key], v)))!
+    }
+    return eq(table[key], value)
+  })
+  return conditions.length === 1 ? conditions[0]! : and(...conditions)!
+}
+
+export type CreateDrizzleOptions<
+  TSchema extends Record<string, unknown>,
+  TRelations extends AnyRelations = EmptyRelations,
+> = {
+  pool: Pool
+  schema: TSchema
+  relations?: TRelations
+}
+
+type ExtractTable<TSchema, T extends TableName<TSchema>> = TSchema[T] extends PgTable
+  ? TSchema[T]
+  : never
+
+export type DrizzleHelpers<
+  TSchema extends Record<string, unknown> = Record<string, unknown>,
+  TDb extends NodePgDatabase<any, any> = NodePgDatabase<TSchema>,
+> = {
+  insertRow<T extends TableName<TSchema>>(
+    tableName: T,
+    values: InsertMap<TSchema>[T]
+  ): ReturnType<ReturnType<TDb['insert']>['values']>
+  updateRow<T extends TableName<TSchema>>(
+    tableName: T,
+    where: WhereClause<TSchema, T>,
+    set: InsertMap<TSchema>[T]
+  ): ReturnType<ReturnType<ReturnType<TDb['update']>['set']>['where']>
+  deleteRow<T extends TableName<TSchema>>(
+    tableName: T,
+    where: WhereClause<TSchema, T>
+  ): ReturnType<ReturnType<TDb['delete']>['where']>
+}
+
+export function createDrizzle<
+  TSchema extends Record<string, unknown>,
+  TRelations extends AnyRelations = EmptyRelations,
+>(
+  options: CreateDrizzleOptions<TSchema, TRelations>
+): NodePgDatabase<TSchema, TRelations> &
+  DrizzleHelpers<TSchema, NodePgDatabase<TSchema, TRelations>> {
+  patchTimestampForZero()
+
+  // drizzle v1 overloads don't support generic schema+relations together at the call site
+  const instance = drizzle({
+    client: options.pool,
+    schema: options.schema as Record<string, unknown>,
+    relations: options.relations,
+    logger: process.env.DEBUG_SQL === '1',
+  } as any) as unknown as NodePgDatabase<TSchema, TRelations>
+
+  const helpers: DrizzleHelpers<TSchema, NodePgDatabase<TSchema, TRelations>> = {
+    insertRow(tableName, values) {
+      const table = options.schema[tableName] as any
+      return instance.insert(table).values(values as Record<string, unknown>)
+    },
+    updateRow(tableName, where, set) {
+      const table = options.schema[tableName] as any
+      return instance
+        .update(table)
+        .set(set as Record<string, unknown>)
+        .where(resolveWhere(table, where as Record<string, WhereValue>))
+    },
+    deleteRow(tableName, where) {
+      const table = options.schema[tableName] as any
+      return instance
+        .delete(table)
+        .where(resolveWhere(table, where as Record<string, WhereValue>))
+    },
+  }
+
+  return Object.assign(instance, helpers)
+}

package/src/takeout/createPool.ts

@@ -0,0 +1,58 @@
+import pg, { type Pool } from 'pg'
+
+export type CreatePoolOptions = {
+  connectionString: string
+  max?: number
+  idleTimeoutMillis?: number
+  connectionTimeoutMillis?: number
+  keepAlive?: boolean
+  keepAliveInitialDelayMillis?: number
+}
+
+export function createPool(options: CreatePoolOptions): Pool {
+  const {
+    connectionString,
+    max = 20,
+    idleTimeoutMillis = 30_000,
+    connectionTimeoutMillis = 5_000,
+    keepAlive = true,
+    keepAliveInitialDelayMillis = 10_000,
+  } = options
+
+  const ssl = connectionString.includes('sslmode=require')
+    ? { rejectUnauthorized: false }
+    : undefined
+
+  const pool = new pg.Pool({
+    connectionString,
+    max,
+    idleTimeoutMillis,
+    connectionTimeoutMillis,
+    keepAlive,
+    keepAliveInitialDelayMillis,
+    ssl,
+  })
+
+  pool.on('error', (error) => {
+    console.error(`[database] pool error`, error.message, error.stack)
+  })
+
+  pool.on('connect', (client) => {
+    client.on('error', (error) => {
+      console.error(`[database] client error`, error.message, error.stack)
+    })
+
+    // pg-protocol's parser can throw synchronously inside the stream 'data' handler
+    // (e.g. RangeError when reading corrupted/stale buffer data from a zombie connection).
+    // wrapping the stream with its own error handler catches them at the source.
+    const stream = (client as any).connection?.stream
+    if (stream && !stream.__pgErrorPatched) {
+      stream.__pgErrorPatched = true
+      stream.on('error', (error: Error) => {
+        console.error(`[database] stream error`, error.message, error.stack)
+      })
+    }
+  })
+
+  return pool
+}