pipework 0.251.0 → 0.252.1

package/src/infra/quota/handle.ts

@@ -0,0 +1,215 @@
+import { sql } from '../../data/db/sql.js'
+import type { DB } from '../../data/db/index.js'
+import { getContext } from '../../request/context/store.js'
+import { getBaseLogger } from '../../core/logging/index.js'
+import { recordFault, faultMessage } from '../alerts/faults.js'
+import type { ManagedConnection } from '../../data/db/pool.js'
+import type { AcceptReadonly } from '../../data/db/readonly.js'
+import { QuotaExhaustedError } from './errors.js'
+import { retryUnsent, type ConnectRetryOptions } from '../fetch/connect-retry.js'
+
+/** Anything with a case-insensitive `get` — a fetch `Headers` satisfies it. */
+export type QuotaHeaders = { get(name: string): string | null }
+export type QuotaHeaderInput = QuotaHeaders | AcceptReadonly<Record<string, string | string[] | undefined>>
+
+export interface QuotaBudget {
+  readonly remaining: number
+  readonly resetAt: Date
+  readonly limit?: number
+}
+
+export interface QuotaConfig<C extends string> {
+  /** Remaining count each class keeps in reserve for others: a call of class c is refused while remaining <= reserve[c]. Classes absent here have reserve 0 (refused only at zero). */
+  readonly reserve?: Partial<Record<C, number>>
+  /** Reads the budget from a response's headers; default reads GitHub-style headers. Return null when the response carries none. */
+  readonly read?: (headers: QuotaHeaders) => QuotaBudget | null
+  /** Database name; default the instance's default database. */
+  readonly database?: string
+  /** A standalone connection (from `tap()`) for a process outside the manifold; the caller owns its lifecycle. With it, no managed context is needed and `database` is not used. */
+  readonly connection?: ManagedConnection
+}
+
+export interface QuotaState {
+  readonly remaining: number
+  readonly limit: number | null
+  readonly resetAt: Date
+}
+
+export interface QuotaCallOptions<R> {
+  /** Where to read the response headers; default `result.headers` when it has a `get`. */
+  readonly headers?: (result: R) => QuotaHeaderInput | null
+  /**
+   * Sends again when the connection could not be made (see `connectRetry`). One budget unit is taken for the call:
+   * a connect fault never reached the upstream, so the tries cost nothing more. `true` uses the default waits.
+   */
+  readonly retryUnsent?: true | ConnectRetryOptions
+}
+
+export interface Quota<C extends string> {
+  /** Takes one unit of the upstream's budget for this class, sends, and records the budget from the response headers. */
+  call<R>(target: { key: string; class: C }, send: () => Promise<R>, opts?: QuotaCallOptions<R>): Promise<R>
+  /** Records the budget from headers of a call sent through the caller's own client. */
+  record(key: string, headers: QuotaHeaderInput): Promise<void>
+  /** The recorded budget; null when there is none or its window has passed. */
+  state(key: string): Promise<QuotaState | null>
+}
+
+function intHeader(headers: QuotaHeaders, name: string): number | null {
+  const raw = headers.get(name)
+  if (raw === null) return null
+  const n = Number(raw.trim())
+  return raw.trim() !== '' && Number.isFinite(n) ? Math.trunc(n) : null
+}
+
+function defaultRead(headers: QuotaHeaders): QuotaBudget | null {
+  const remaining = intHeader(headers, 'x-ratelimit-remaining')
+  const reset = intHeader(headers, 'x-ratelimit-reset')
+  if (remaining === null || reset === null) return null
+  const limit = intHeader(headers, 'x-ratelimit-limit')
+  return { remaining, resetAt: new Date(reset * 1000), ...(limit !== null ? { limit } : {}) }
+}
+
+function toHeaders(input: QuotaHeaderInput): QuotaHeaders {
+  if (typeof (input as QuotaHeaders).get === 'function') return input as QuotaHeaders
+  const map = new Map<string, string>()
+  for (const [name, value] of Object.entries(input as Record<string, string | readonly string[] | undefined>)) {
+    if (value === undefined) continue
+    map.set(name.toLowerCase(), typeof value === 'string' ? value : value.join(', '))
+  }
+  return { get: (name) => map.get(name.toLowerCase()) ?? null }
+}
+
+function headersOf(value: unknown): QuotaHeaderInput | null {
+  if (typeof value !== 'object' || value === null) return null
+  const headers = (value as { headers?: unknown }).headers
+  if (typeof headers !== 'object' || headers === null) return null
+  return headers as QuotaHeaderInput
+}
+
+function resolveDb(database: string | undefined): DB {
+  const ctx = getContext()
+  if (ctx === undefined) {
+    throw new Error(
+      '[pipework] quota needs a managed context to reach the database. ' +
+      'Call it inside a request handler, a job, or flow.run.',
+    )
+  }
+  const name = database ?? ctx.instance.config.defaultDatabaseName()
+  return ctx.instance.pool.getOrCreate(ctx.instance.config.database(name)).drizzle
+}
+
+const iso = (d: Date): string => d.toISOString()
+
+function logRecordFailure(key: string, err: unknown): void {
+  getBaseLogger().error({ err, key }, 'quota could not record the budget from a response')
+  void recordFault('quota-record', faultMessage(err), { key })
+}
+
+export function createQuota<C extends string>(config: QuotaConfig<C> = {}): Quota<C> {
+  if (config.connection !== undefined && config.database !== undefined) {
+    throw new Error('[pipework] quota.create takes connection or database, not both.')
+  }
+  const read = config.read ?? defaultRead
+  const db = (): DB => config.connection !== undefined ? config.connection.drizzle : resolveDb(config.database)
+
+  async function record(key: string, input: QuotaHeaderInput): Promise<void> {
+    const headers = toHeaders(input)
+    const budget = read(headers)
+    const retryAfter = intHeader(headers, 'retry-after')
+    if (budget === null && retryAfter === null) return
+    const conn = db()
+    if (budget !== null) {
+      // Same window: lowest remaining. Different window: a live response (reset in the future) replaces the row; a past reset is ignored.
+      const limit = budget.limit ?? null
+      await conn.execute(sql`
+        INSERT INTO outbound_quota (key, remaining, "limit", reset_at, updated_at)
+        VALUES (${key}, ${budget.remaining}, ${limit}, ${iso(budget.resetAt)}::timestamptz, clock_timestamp())
+        ON CONFLICT (key) DO UPDATE SET
+          remaining = CASE
+            WHEN date_trunc('second', outbound_quota.reset_at) = date_trunc('second', EXCLUDED.reset_at)
+              THEN LEAST(outbound_quota.remaining, EXCLUDED.remaining)
+            WHEN EXCLUDED.reset_at > now() THEN EXCLUDED.remaining
+            ELSE outbound_quota.remaining END,
+          "limit" = CASE
+            WHEN date_trunc('second', outbound_quota.reset_at) = date_trunc('second', EXCLUDED.reset_at)
+              OR EXCLUDED.reset_at > now()
+              THEN COALESCE(EXCLUDED."limit", outbound_quota."limit")
+            ELSE outbound_quota."limit" END,
+          reset_at = CASE
+            WHEN date_trunc('second', outbound_quota.reset_at) = date_trunc('second', EXCLUDED.reset_at)
+              THEN outbound_quota.reset_at
+            WHEN EXCLUDED.reset_at > now() THEN EXCLUDED.reset_at
+            ELSE outbound_quota.reset_at END,
+          updated_at = clock_timestamp()`)
+    }
+    if (retryAfter !== null) {
+      await conn.execute(sql`
+        INSERT INTO outbound_quota (key, remaining, "limit", reset_at, updated_at)
+        VALUES (${key}, 0, NULL, now() + make_interval(secs => ${Math.max(0, retryAfter)}), clock_timestamp())
+        ON CONFLICT (key) DO UPDATE SET
+          remaining = 0,
+          reset_at = GREATEST(outbound_quota.reset_at, now() + make_interval(secs => ${Math.max(0, retryAfter)})),
+          updated_at = clock_timestamp()`)
+    }
+  }
+
+  async function state(key: string): Promise<QuotaState | null> {
+    const rows = await db().execute(sql`
+      SELECT remaining, "limit", reset_at FROM outbound_quota WHERE key = ${key} AND reset_at > now()`)
+    const row = rows[0]
+    if (row === undefined) return null
+    return {
+      remaining: Number(row['remaining']),
+      limit: row['limit'] === null ? null : Number(row['limit']),
+      resetAt: new Date(row['reset_at'] as string | Date),
+    }
+  }
+
+  async function take(key: string, cls: string): Promise<void> {
+    const conn = db()
+    const reserve = config.reserve?.[cls as C] ?? 0
+    const taken = await conn.execute(sql`
+      UPDATE outbound_quota SET remaining = remaining - 1, updated_at = clock_timestamp()
+      WHERE key = ${key} AND reset_at > now() AND remaining > ${reserve}
+      RETURNING remaining`)
+    if (taken.length > 0) return
+    const rows = await conn.execute(sql`
+      UPDATE outbound_quota SET refused_at = clock_timestamp()
+      WHERE key = ${key} AND reset_at > now()
+      RETURNING remaining, reset_at`)
+    const row = rows[0]
+    if (row === undefined) return // unknown budget: send optimistically
+    throw new QuotaExhaustedError({
+      key, class: cls, reserve,
+      remaining: Number(row['remaining']),
+      resetAt: new Date(row['reset_at'] as string | Date),
+    })
+  }
+
+  return {
+    record,
+    state,
+    async call<R>(target: { key: string; class: C }, send: () => Promise<R>, opts: QuotaCallOptions<R> = {}): Promise<R> {
+      await take(target.key, target.class)
+      let result: R
+      try {
+        const o = opts.retryUnsent
+        result = o === undefined ? await send() : await retryUnsent(send, o === true ? {} : o)
+      } catch (err) {
+        const e = err as { response?: unknown } | null
+        const found = headersOf(e?.response) ?? headersOf(err)
+        if (found !== null) await record(target.key, found).catch((recordErr: unknown) => { logRecordFailure(target.key, recordErr) })
+        throw err
+      }
+      const found = opts.headers !== undefined
+        ? opts.headers(result)
+        : (() => {
+            const h = headersOf(result)
+            return h !== null && typeof (h as QuotaHeaders).get === 'function' ? h : null
+          })()
+      // The call was sent: a failure to record must not reject it, or a retry would send it again.
+      if (found !== null) await record(target.key, found).catch((recordErr: unknown) => { logRecordFailure(target.key, recordErr) })
+      return result
+    },
+  }
+}

package/src/infra/quota/index.ts

@@ -0,0 +1,5 @@
+export { createQuota } from './handle.js'
+export type {
+  Quota, QuotaConfig, QuotaHeaders, QuotaHeaderInput, QuotaBudget, QuotaState, QuotaCallOptions,
+} from './handle.js'
+export { QuotaExhaustedError } from './errors.js'

package/src/infra/quota/namespace.ts

@@ -0,0 +1,7 @@
+import { createQuota } from './handle.js'
+
+/** Quota namespace — callers of one rate-limited upstream share its budget by class, read from the upstream's headers and kept in Postgres. */
+export const quota = {
+  /** Creates a quota handle. */
+  create: createQuota,
+} as const

package/src/infra/quota/schema.ts

@@ -0,0 +1,15 @@
+// One row per upstream key: what the upstream last said was left, and when it resets.
+// Not tenant-scoped — one budget per upstream across the whole app.
+
+import { define } from '../../data/domain/index.js'
+import { field } from '../../data/domain/field.js'
+
+export const outboundQuota = define('outbound_quota', {
+  key: field.text().primaryKey(),
+  remaining: field.integer(),
+  limit: field.integer().nullable(),
+  resetAt: field.timestamp(),
+  updatedAt: field.timestamp().default('now'),
+  /** When a call was last refused against this row; it counts while it is not older than `updatedAt`. */
+  refusedAt: field.timestamp().nullable(),
+})

package/src/infra/storage/backend.ts

@@ -0,0 +1,34 @@
+import type { Readable } from 'node:stream'
+
+/** One stored object as a listing reports it. `key` is relative to the backend's prefix. */
+export interface BackendObject {
+  readonly key: string
+  readonly lastModified: Date
+}
+
+export interface BackendSignedUrlOptions {
+  readonly expiresInSeconds: number
+  readonly disposition: 'attachment' | 'inline'
+  readonly filename?: string
+  readonly mediaType: string
+}
+
+/**
+ * Where the bytes of stored files live. `storage.s3(...)` and `storage.local(...)`
+ * make one; keys are the relative `tenants/<tenantId>/...` names the handle builds.
+ */
+export interface StorageBackend {
+  readonly name: 's3' | 'local'
+  /** Streams `body` to `key`, replacing any object there. A stream error leaves nothing at `key`. */
+  put(key: string, body: AsyncIterable<Buffer | Uint8Array>): Promise<void>
+  exists(key: string): Promise<boolean>
+  copy(from: string, to: string): Promise<void>
+  /** Removes an object; absent is not an error. */
+  remove(key: string): Promise<void>
+  /** Streams an object's bytes. */
+  read(key: string): Promise<Readable>
+  /** Every object whose key starts with `prefix`, in key order. */
+  list(prefix: string): AsyncIterable<BackendObject>
+  /** A time-limited URL that serves the object to a client with no credentials; throws where the backend cannot. */
+  signedUrl(key: string, options: BackendSignedUrlOptions): Promise<string>
+}