@take-out/postgres 0.0.28

package/package.json

@@ -0,0 +1,67 @@
+{
+  "name": "@take-out/postgres",
+  "version": "0.0.28",
+  "sideEffects": false,
+  "source": "src/index.ts",
+  "main": "dist/cjs",
+  "module": "dist/esm",
+  "type": "module",
+  "types": "./src/index.ts",
+  "bin": {
+    "postgres": "./cli.cjs"
+  },
+  "files": [
+    "src",
+    "types",
+    "dist",
+    "cli.cjs"
+  ],
+  "publishConfig": {
+    "access": "public"
+  },
+  "scripts": {
+    "build": "tamagui-build",
+    "watch": "tamagui-build --watch",
+    "lint": "biome check src",
+    "lint:fix": "biome check --write src",
+    "clean": "tamagui-build clean",
+    "clean:build": "tamagui-build clean:build"
+  },
+  "exports": {
+    "./package.json": "./package.json",
+    ".": {
+      "types": "./src/index.ts",
+      "import": "./dist/esm/index.mjs",
+      "require": "./dist/cjs/index.cjs",
+      "default": "./dist/esm/index.mjs"
+    },
+    "./migrate": {
+      "types": "./src/migrate.ts",
+      "import": "./dist/esm/migrate.mjs",
+      "require": "./dist/cjs/migrate.cjs"
+    },
+    "./build": {
+      "types": "./src/build.ts",
+      "import": "./dist/esm/build.mjs",
+      "require": "./dist/cjs/build.cjs"
+    }
+  },
+  "dependencies": {
+    "@take-out/helpers": "0.0.28",
+    "async-retry": "^1.3.3",
+    "citty": "^0.1.6",
+    "vite": "^6.0.11"
+  },
+  "peerDependencies": {
+    "drizzle-orm": "*",
+    "pg": "^8.16.3"
+  },
+  "devDependencies": {
+    "@biomejs/biome": "^2.1.2",
+    "@tamagui/build": "*",
+    "@types/async-retry": "^1.4.8",
+    "@types/node": "24.0.3",
+    "drizzle-kit": "^0.30.5",
+    "drizzle-orm": "^0.40.0"
+  }
+}

package/src/build.ts

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

package/src/cli.ts

@@ -0,0 +1,153 @@
+#!/usr/bin/env node
+import { defineCommand, runMain } from 'citty'
+import { join } from 'node:path'
+import { buildMigrations as buildMigrationsRun } from './scripts/build-migrations'
+import { syncDrizzleMigrations } from './scripts/drizzle-migrations-sync'
+import { addMigration } from './scripts/migration-add'
+import { runPgDump } from './scripts/pg_dump'
+import { runPsql } from './scripts/psql'
+
+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 buildMigrations = defineCommand({
+  meta: {
+    name: 'build-migrations',
+    description: 'Build migration bundle for deployment',
+  },
+  args: {
+    dir: {
+      type: 'string',
+      description: 'Migrations directory',
+      required: false,
+      default: './src/database/migrations',
+    },
+    out: {
+      type: 'string',
+      description: 'Output file name',
+      required: false,
+      default: 'migrate-dist.js',
+    },
+  },
+  async run({ args }) {
+    const migrationsDir = join(process.cwd(), args.dir)
+    await buildMigrationsRun({
+      migrationsDir,
+      outFile: args.out,
+    })
+  },
+})
+
+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: 'postgres',
+    description: 'PostgreSQL database utilities and migration tools',
+    version: '0.0.1',
+  },
+  subCommands: {
+    'sync-drizzle': syncDrizzle,
+    'build-migrations': buildMigrations,
+    'migrate:add': migrationAdd,
+    psql,
+    pg_dump: pgDump,
+  },
+})
+
+runMain(main)

package/src/createServerHelpers.ts

@@ -0,0 +1,20 @@
+import type { Pool } from 'pg'
+import { createSql, setDefaultPool } from './sql'
+import { getDBClient, type GetDBClientOptions } from './helpers/getDBClient'
+
+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/helpers/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/helpers/getDBClient.ts

@@ -0,0 +1,264 @@
+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) {
+        console.info(`Connecting to provided pool...`)
+        const client = await providedPool.connect()
+        console.info(`Connected to pool`)
+        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)
+
+      console.info(`Connecting to pool ${connectionString}...`)
+      const client = await poolCache.pool.connect()
+      console.info(`Connected to pool`)
+      return client
+    },
+    {
+      retries,
+      minTimeout: 300,
+      factor: 2,
+      maxTimeout: 8000,
+    }
+  )
+
+  return clientFromPool
+}
+
+export async function queryDb(
+  queryText: string,
+  params?: any[],
+  options: GetDBClientOptions = {}
+): Promise<pg.QueryResult<any>> {
+  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)
+    console.info(
+      `Max connections: ${maxConnections}, Reserved connections: ${reservedConnections}`
+    )
+    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]> {
+  console.info(`Getting connection limits...`)
+  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)
+  console.info(`Getting opened connections...`)
+  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
+  )
+  console.info(`Opened connections: ${result}`)
+  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/index.ts

@@ -0,0 +1,36 @@
+import type { NodePgDatabase } from 'drizzle-orm/node-postgres'
+import { drizzle } from 'drizzle-orm/node-postgres'
+import type { Pool } from 'pg'
+import pg from 'pg'
+
+export const createPool = (connectionString: string): Pool => {
+  return new pg.Pool({
+    connectionString,
+  })
+}
+
+export const createDb = <TSchema extends Record<string, unknown>>(
+  connectionString: string,
+  schema: TSchema
+): NodePgDatabase<TSchema> => {
+  const pool = createPool(connectionString)
+  return drizzle(pool, {
+    schema,
+    logger: false,
+  }) as NodePgDatabase<TSchema>
+}
+
+export type { NodePgDatabase }
+
+// re-export for convenience
+export { createSql, type SqlQuery } from './sql'
+
+export { createServerHelpers, type ServerHelpers } from './createServerHelpers'
+
+export {
+  getDBClient,
+  queryDb,
+  type GetDBClientOptions,
+} from './helpers/getDBClient'
+
+export { processInChunks, updateInChunks } from './helpers/chunkedQuery'