pipework 0.282.0 → 0.283.0

package/src/data/db/pool.ts

@@ -166,7 +166,10 @@ export function createManagedConnection(
   observability?: ResolvedObservability,
   applicationName?: string,
 ): ManagedConnection {
-  const baseDebug = observability ? createDebugFn(observability) : undefined
+  const statementWatch = observability
+    ? createStatementWatch({ database: config.name, url: config.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs, onSettled: observability.onStatementSettled })
+    : undefined
+  const baseDebug = observability ? createDebugFn(observability, statementWatch?.noteConnection) : undefined
   const slot: ProbeSlot | undefined = baseDebug ? { current: null } : undefined
   const debug = baseDebug && slot
     ? ((connection, query, parameters) => {
@@ -187,9 +190,6 @@ export function createManagedConnection(
   const observerOnly = observability !== undefined && isObserverOnly(observability)
   if (observerOnly) markObservedConnection(client)
 
-  const statementWatch = observability
-    ? createStatementWatch({ database: config.name, url: config.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs })
-    : undefined
   const db = splitLargeInserts(drizzle(client, statementWatch ? { statementWatch: statementWatch.watch } : {}))
   const canceller = createCanceller(config.url)
 

package/src/data/db/query-watch.ts

@@ -1,7 +1,9 @@
 import postgres from 'postgres'
 import { getLogger } from '../../core/logging/logger.js'
 import { currentQueryLabel } from './query-label.js'
+import { DrizzleQueryError } from '../query/errors.js'
 import type { StatementWatch } from '../query/postgres-js/session.js'
+import type { StatementSettledError, StatementSettledObserver } from '../../core/config/statement-observer.js'
 
 const LABEL_SQL_CHARS = 120
 const SAMPLE_TIMEOUT_MS = 2_000
@@ -23,13 +25,42 @@ export interface StatementWatchOptions {
   readonly slowQueryMs: number
   /** Undefined: no timer is ever created. */
   readonly inFlightQueryMs: number | undefined
+  /** `observability.onStatementSettled`; undefined: nothing extra is tracked or fired. */
+  readonly onSettled?: StatementSettledObserver | undefined
 }
 
 export interface StatementWatchHandle {
   readonly watch: StatementWatch
+  /**
+   * Pass to `createDebugFn` for the same pool: the driver's debug hook is the only place that
+   * learns a statement's connection. Undefined when no `onSettled` hook is configured.
+   */
+  readonly noteConnection: ((connectionId: number, parameters: readonly unknown[]) => void) | undefined
   close(): Promise<void>
 }
 
+/** Rows a settled statement returned, or affected for a write: the driver's command-tag count, else the row count. */
+function rowsOf(result: unknown): number {
+  if (typeof result !== 'object' || result === null) return 0
+  const count = (result as { count?: unknown }).count
+  if (typeof count === 'number') return count
+  return Array.isArray(result) ? result.length : 0
+}
+
+/** Name, message and code only: the driver's error object carries the statement text and bind parameters. */
+function settledError(raw: unknown): StatementSettledError {
+  // Drizzle wraps a driver failure in an error whose message is `Failed query: <sql>\nparams: <values>`: report the cause.
+  let err = raw
+  while (err instanceof DrizzleQueryError && err.cause !== undefined) err = err.cause
+  if (typeof err !== 'object' || err === null) return { name: 'Error', message: String(err) }
+  const { name, message, code } = err as { name?: unknown; message?: unknown; code?: unknown }
+  return {
+    name: typeof name === 'string' ? name : 'Error',
+    message: typeof message === 'string' ? message : '',
+    ...(typeof code === 'string' ? { code } : {}),
+  }
+}
+
 function labelFor(sqlText: string): string {
   return currentQueryLabel() ?? sqlText.replace(/\s+/g, ' ').trim().slice(0, LABEL_SQL_CHARS)
 }
@@ -41,7 +72,10 @@ function labelFor(sqlText: string): string {
  * dedicated connection, never the statement's own and never the app pool.
  */
 export function createStatementWatch(options: StatementWatchOptions): StatementWatchHandle {
-  const { database, url, slowQueryMs, inFlightQueryMs } = options
+  const { database, url, slowQueryMs, inFlightQueryMs, onSettled } = options
+  // Keyed on the parameter array the driver builds for a statement: the same array reaches the
+  // debug hook (with the connection ordinal) and sits on the pending query the session hands us.
+  const connections = onSettled ? new WeakMap<object, number>() : undefined
   let sampler: postgres.Sql | undefined
   let sampling = false
   let closed = false
@@ -72,8 +106,24 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
     return Promise.race([run, timeout]).catch(() => UNKNOWN_BACKEND).finally(() => { clearTimeout(timer) })
   }
 
-  const watch: StatementWatch = async <T>(sqlText: string, run: () => Promise<T>): Promise<T> => {
+  function fireSettled(event: Parameters<StatementSettledObserver>[0]): void {
+    if (!onSettled) return
+    const failed = (err: unknown): void => {
+      getLogger().warn({ err, database }, 'statement:settled-hook-failed')
+    }
+    try {
+      const out = onSettled(event)
+      if (out !== undefined && typeof (out as Promise<void>).then === 'function') void Promise.resolve(out).catch(failed)
+    } catch (err) {
+      failed(err)
+    }
+  }
+
+  const watch: StatementWatch = async <T>(sqlText: string, run: (seen: (pending: object) => void) => Promise<T>): Promise<T> => {
     const startedMs = Date.now()
+    let pendingQuery: object | undefined
+    let outcome: { result: unknown } | { error: unknown } | undefined
+    const seen = (pending: object): void => { pendingQuery = pending }
     let interval: NodeJS.Timeout | undefined
     let settled = false
     if (inFlightQueryMs !== undefined) {
@@ -93,7 +143,12 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
       interval.unref()
     }
     try {
-      return await run()
+      const result = await run(seen)
+      if (onSettled) outcome = { result }
+      return result
+    } catch (err) {
+      if (onSettled) outcome = { error: err }
+      throw err
     } finally {
       settled = true
       if (interval !== undefined) clearInterval(interval)
@@ -101,11 +156,25 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
       if (slowQueryMs > 0 && durationMs > slowQueryMs) {
         getLogger().warn({ label: labelFor(sqlText), durationMs, database }, 'query:slow')
       }
+      if (onSettled && outcome !== undefined) {
+        const parameters = pendingQuery === undefined ? undefined : (pendingQuery as { parameters?: unknown }).parameters
+        const connectionId = typeof parameters === 'object' && parameters !== null ? connections?.get(parameters) ?? null : null
+        fireSettled({
+          label: currentQueryLabel() ?? null,
+          durationMs,
+          rows: 'result' in outcome ? rowsOf(outcome.result) : 0,
+          connectionId,
+          ...('error' in outcome ? { error: settledError(outcome.error) } : {}),
+        })
+      }
     }
   }
 
   return {
     watch,
+    noteConnection: connections
+      ? (connectionId, parameters) => { connections.set(parameters, connectionId) }
+      : undefined,
     async close() {
       closed = true
       const client = sampler

package/src/data/db/raw-client-warning.ts

@@ -23,7 +23,7 @@ export function warnRawClientIsUnsanitized(database: string): void {
 
   getBaseLogger().warn(
     { database },
-    '[pipework] raw .client read on a connection with observability.onStatement configured. ' +
+    '[pipework] raw .client read on a connection with observability.onStatement or onStatementSettled configured. ' +
     'Queries issued through .client bypass pipework\'s driver error boundary, so their errors ' +
     'serialize bind parameters (postgres.js marks query/parameters/args enumerable whenever its ' +
     'debug hook is installed). Use connection.drizzle — or pipe() inside a managed context — to ' +

package/src/data/query/postgres-js/session.ts

@@ -16,10 +16,19 @@ import { type Assume, mapResultRow } from '../utils.js';
 import { propagateObserved, throughErrorBoundary } from './driver-error-boundary.js';
 
 /** Wraps one statement's execution; installed by the pool to time it and report it while it runs. */
-export type StatementWatch = <T>(query: string, run: () => Promise<T>) => Promise<T>;
+export type StatementWatch = <T>(query: string, run: (seen: (pending: object) => void) => Promise<T>) => Promise<T>;
 
-function watched<T>(watch: StatementWatch | undefined, query: string, run: () => Promise<T>): Promise<T> {
-	return watch === undefined ? run() : watch(query, run);
+function watched<T>(watch: StatementWatch | undefined, query: string, run: (seen: (pending: object) => void) => Promise<T>): Promise<T> {
+	return watch === undefined ? run(noteSeen) : watch(query, run);
+}
+
+function noteSeen(): void {}
+
+/** Starts the statement and tells the watch which pending query it is, so it can learn the connection it runs on. */
+function startUnsafe(client: Sql, query: string, params: unknown[], seen: (pending: object) => void) {
+	const pending = client.unsafe(query, params as any[]);
+	seen(pending);
+	return pending;
 }
 
 export class PostgresJsPreparedQuery<T extends PreparedQueryConfig> extends PgPreparedQuery<T> {
@@ -58,8 +67,8 @@ export class PostgresJsPreparedQuery<T extends PreparedQueryConfig> extends PgPr
 			const { fields, queryString: query, client, joinsNotNullableMap, customResultMapper, watch } = this;
 			if (!fields && !customResultMapper) {
 				return tracer.startActiveSpan('drizzle.driver.execute', () => {
-					return throughErrorBoundary(client, () => watched(watch, query, () => this.queryWithCache(query, params, async () => {
-						return await client.unsafe(query, params as any[]);
+					return throughErrorBoundary(client, () => watched(watch, query, (seen) => this.queryWithCache(query, params, async () => {
+						return await startUnsafe(client, query, params, seen);
 					})));
 				});
 			}
@@ -69,8 +78,8 @@ export class PostgresJsPreparedQuery<T extends PreparedQueryConfig> extends PgPr
 					'drizzle.query.text': query,
 					'drizzle.query.params': JSON.stringify(params),
 				});
-				return throughErrorBoundary(client, () => watched(watch, query, () => this.queryWithCache(query, params, async () => {
-					return await client.unsafe(query, params as any[]).values();
+				return throughErrorBoundary(client, () => watched(watch, query, (seen) => this.queryWithCache(query, params, async () => {
+					return await startUnsafe(client, query, params, seen).values();
 				})));
 			});
 
@@ -95,8 +104,8 @@ export class PostgresJsPreparedQuery<T extends PreparedQueryConfig> extends PgPr
 					'drizzle.query.text': this.queryString,
 					'drizzle.query.params': JSON.stringify(params),
 				});
-				return throughErrorBoundary(this.client, () => watched(this.watch, this.queryString, () => this.queryWithCache(this.queryString, params, async () => {
-					return this.client.unsafe(this.queryString, params as any[]);
+				return throughErrorBoundary(this.client, () => watched(this.watch, this.queryString, (seen) => this.queryWithCache(this.queryString, params, async () => {
+					return startUnsafe(this.client, this.queryString, params, seen);
 				})));
 			});
 		});
@@ -165,14 +174,14 @@ export class PostgresJsSession<
 
 	query(query: string, params: unknown[]): Promise<RowList<Row[]>> {
 		this.logger.logQuery(query, params);
-		return throughErrorBoundary(this.client, () => watched(this.options.watch, query, () => this.client.unsafe(query, params as any[]).values()));
+		return throughErrorBoundary(this.client, () => watched(this.options.watch, query, (seen) => startUnsafe(this.client, query, params, seen).values()));
 	}
 
 	queryObjects<T extends Row>(
 		query: string,
 		params: unknown[],
 	): Promise<RowList<T[]>> {
-		return throughErrorBoundary(this.client, () => watched(this.options.watch, query, () => this.client.unsafe(query, params as any[])));
+		return throughErrorBoundary(this.client, () => watched(this.options.watch, query, (seen) => startUnsafe(this.client, query, params, seen) as unknown as Promise<RowList<T[]>>));
 	}
 
 	override transaction<T>(

package/src/index.ts

@@ -21,7 +21,7 @@ export type { TimeoutOptions } from './data/db/timeouts.js'
 // sanctioned door to a raw connection; keeps the WRAPPED_PACKAGES postgres ban airtight.
 export { tap } from './data/db/pool.js'
 export type { ManagedConnection, ConnectOptions } from './data/db/pool.js'
-export type { StatementObserver, StatementEvent } from './core/config/statement-observer.js'
+export type { StatementObserver, StatementEvent, StatementSettledObserver, StatementSettledEvent, StatementSettledError } from './core/config/statement-observer.js'
 
 export type { DefinedTable, ExternalTable, DomainDefinition, DefinitionMethods, FieldDescriptor, FieldModifiers, FieldMeta, DefineOptions,
   FactoryBuilder, FactoryOptions, TraitDefinition, FieldColumn, ProjectedColumns,

package/src/test/plugin.ts

@@ -127,10 +127,10 @@ async function createIsolatedConnection(
   observability: ResolvedObservability,
   withSystem: boolean,
 ): Promise<PostgresJsDatabase> {
-  const debug = createDebugFn(observability)
+  const watch = createStatementWatch({ database: name, url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs, onSettled: observability.onStatementSettled })
+  const debug = createDebugFn(observability, watch.noteConnection)
   const client = postgres(testDb.url, { max: 1, onnotice: () => {}, ...(debug ? { debug } : {}) })
   if (isObserverOnly(observability)) markObservedConnection(client)
-  const watch = createStatementWatch({ database: name, url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs })
   const db = drizzle(client, { statementWatch: watch.watch })
 
   await db.execute(sql.raw('BEGIN'))
@@ -177,12 +177,12 @@ export function createSetupConnection(
   testDb: TestDatabase,
   observability?: ResolvedObservability,
 ): { drizzle: PostgresJsDatabase; close: () => Promise<void> } {
-  const debug = observability ? createDebugFn(observability) : undefined
-  const client = postgres(testDb.url, { max: 1, onnotice: () => {}, ...(debug ? { debug } : {}) })
-  if (observability && isObserverOnly(observability)) markObservedConnection(client)
   const watch = observability
-    ? createStatementWatch({ database: 'setup', url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs })
+    ? createStatementWatch({ database: 'setup', url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs, onSettled: observability.onStatementSettled })
     : undefined
+  const debug = observability ? createDebugFn(observability, watch?.noteConnection) : undefined
+  const client = postgres(testDb.url, { max: 1, onnotice: () => {}, ...(debug ? { debug } : {}) })
+  if (observability && isObserverOnly(observability)) markObservedConnection(client)
   const db = splitLargeInserts(drizzle(client, watch ? { statementWatch: watch.watch } : {}))
   return { drizzle: db, close: async () => { await watch?.close(); await client.end() } }
 }