pipework 0.291.0 → 0.292.0

package/src/data/write/update-each.ts

@@ -0,0 +1,120 @@
+// `pipe.updateEach(table, { keys, data, where, set })`: one UPDATE for many keys, a count per key.
+//
+//   UPDATE "t" SET ... FROM (select * from unnest($1::uuid[], $2::int[], $3::int[]) as "rows"("id", "ordinal", "_i")) "rows"
+//   WHERE <where> RETURNING "rows"."_i"
+//
+// The keys travel as `pipe.rows` columns (one parameter per column), each with its input index `_i`. The
+// rows the update returns are counted per `_i`, which is what tells a key that matched nothing (it never
+// appears in the result) from one that matched once. A key that did not match exactly once rolls the update
+// back (a savepoint inside the caller's transaction) and throws naming the key.
+
+import type { SQL } from '../db/types.js'
+import { pipe } from '../db/db.js'
+import { rows } from '../db/rows.js'
+import type { RowsColumns, RowsData, RowsSource } from '../db/rows.js'
+import type { PgTable } from '../query/pg-core/index.js'
+import { assertNotExternal } from '../domain/external-registry.js'
+import type { DefinedTable } from '../domain/types.js'
+import type { UpdateSetValues } from './values.js'
+import { executeWrite } from './execute.js'
+import { batchLimits } from './batch-insert.js'
+
+const ORDINAL = '_i'
+
+export interface UpdateEachOptions<T extends PgTable, TKeys extends RowsColumns> {
+  /** The key rowset's columns, each with its SQL type: `{ id: 'uuid', ordinal: 'int' }`. Read in `where` and `set` as `k.id`. */
+  readonly keys: TKeys
+  /** One array per column, all the same length; element i of each is input key i. */
+  readonly data: RowsData<TKeys>
+  /** Which columns identify an input key: they must be distinct across inputs, and name the key in an error. Default: all of `keys`. */
+  readonly key?: ReadonlyArray<keyof TKeys & string>
+  /** The match for one key, as for `pipe.update().where()`: `and(eq(t.id, k.id), isNull(t.closedAt))`. */
+  readonly where: (k: RowsSource<'rows', TKeys>) => SQL | undefined
+  /** What to set; a key column is a valid value: `{ closedAt: k.ordinal }`. */
+  readonly set: (k: RowsSource<'rows', TKeys>) => UpdateSetValues<T>
+  /** `'one'` (default): every key must update exactly one row or the update is refused. `'any'`: no refusal, read the counts. */
+  readonly expect?: 'one' | 'any'
+}
+
+function describeKey(columns: readonly string[], data: Record<string, readonly unknown[]>, i: number): string {
+  const part = (c: string): string => {
+    const v = data[c]![i]
+    return `${c}=${v instanceof Date ? v.toISOString() : JSON.stringify(v, (_k, x: unknown) => (typeof x === 'bigint' ? x.toString() : x))}`
+  }
+  return `{ ${columns.map(part).join(', ')} }`
+}
+
+/**
+ * Updates the rows each input key matches in one statement and returns how many rows each key updated, in
+ * input order. By default a key that updated zero rows or more than one makes the call throw, naming the key,
+ * with the whole update undone; so does a key repeated in `data`. Runs in the caller's transaction and is audited
+ * as `pipe.update()` is (one audit row per row changed). An empty `data` returns `[]` without a statement.
+ */
+export async function updateEach<T extends DefinedTable<any> & PgTable, const TKeys extends RowsColumns>( // eslint-disable-line @typescript-eslint/no-explicit-any
+  table: T,
+  options: UpdateEachOptions<T, TKeys>,
+): Promise<number[]> {
+  assertNotExternal(table, 'update')
+  const names = Object.keys(options.keys)
+  if (names.includes(ORDINAL)) throw new Error(`[pipework] pipe.updateEach: "${ORDINAL}" is reserved for the input index; rename that key column.`)
+  const data = options.data as Record<string, readonly unknown[]>
+  const count = names.length === 0 ? 0 : (data[names[0]!]?.length ?? 0)
+  const identity: readonly string[] = options.key ?? names
+  for (const c of identity) {
+    if (!names.includes(c)) throw new Error(`[pipework] pipe.updateEach: key column "${c}" is not one of the key rowset's columns (${names.join(', ')}).`)
+  }
+  const seen = new Map<string, number>()
+  for (let i = 0; i < count; i++) {
+    const id = describeKey(identity, data, i)
+    const first = seen.get(id)
+    if (first !== undefined) {
+      throw new Error(`[pipework] pipe.updateEach: key ${id} appears at inputs ${String(first)} and ${String(i)}. Pass each key once; a repeated key would update one row twice.`)
+    }
+    seen.set(id, i)
+  }
+  if (count === 0) return []
+
+  // Over `batchLimits` the keys go as consecutive slices, in order, in the one transaction below: a refusal
+  // anywhere undoes every slice.
+  const slices: [number, number][] = []
+  let from = 0
+  let bytes = 0
+  for (let i = 0; i < count; i++) {
+    let rowBytes = 0
+    for (const c of names) { const v = data[c]![i]; rowBytes += typeof v === 'string' ? Buffer.byteLength(v) : 8 }
+    if (i > from && (i - from >= batchLimits.rows || bytes + rowBytes > batchLimits.bytes)) { slices.push([from, i]); from = i; bytes = 0 }
+    bytes += rowBytes
+  }
+  slices.push([from, count])
+
+  const dbName = table.$options.database
+  const runSlice = async (a: number, b: number): Promise<number[]> => {
+    const part = Object.fromEntries(names.map(c => [c, data[c]!.slice(a, b)]))
+    const source = rows({ ...options.keys, [ORDINAL]: 'int' }, { ...part, [ORDINAL]: Array.from({ length: b - a }, (_, i) => i) } as never)
+    const k = source as unknown as RowsSource<'rows', TKeys>
+    const where = options.where(k)
+    if (where === undefined) throw new Error('[pipework] pipe.updateEach: where() returned no condition. Match each key, e.g. eq(t.id, k.id): without one every key would update every row.')
+    return await executeWrite({
+      verb: 'update' as const,
+      table,
+      set: options.set(k) as Record<string, unknown>,
+      where,
+      each: { source, ordinal: (source as unknown as Record<string, SQL>)[ORDINAL]!, count: b - a },
+    }) as number[]
+  }
+  const run = async (): Promise<number[]> => {
+    const counts: number[] = []
+    for (const [a, b] of slices) for (const n of await runSlice(a, b)) counts.push(n)
+    if ((options.expect ?? 'one') === 'one') {
+      const bad = counts.findIndex(n => n !== 1)
+      if (bad !== -1) {
+        throw new Error(
+          `[pipework] pipe.updateEach: key ${describeKey(identity, data, bad)} (input ${String(bad)}) matched ${String(counts[bad])} rows of "${table.$name}", expected exactly one. ` +
+          'Nothing was updated. A key matching none is absent or no longer eligible (check the where); a key matching several needs a narrower where. Pass expect: \'any\' to read the counts instead.',
+        )
+      }
+    }
+    return counts
+  }
+  return dbName === undefined ? pipe.transaction(run) : pipe.transaction(dbName, run)
+}