pipework 0.282.0 → 0.284.0

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

@@ -8,7 +8,7 @@ import { runInContext } from '../request/context/store.js'
 import type { Flow } from '../request/context/types.js'
 import { randomUUID } from 'node:crypto'
 import { createDebugFn, isObserverOnly } from '../data/db/instrumentation.js'
-import { createStatementWatch } from '../data/db/query-watch.js'
+import { createStatementWatch, largeQueryBudgets } from '../data/db/query-watch.js'
 import { markObservedConnection } from '../data/query/postgres-js/driver-error-boundary.js'
 import type { ResolvedObservability } from '../core/config/load.js'
 import { collapseTransactionsToSavepoints } from './savepoints.js'
@@ -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, budgets: largeQueryBudgets(observability) })
+  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, budgets: largeQueryBudgets(observability) })
     : 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() } }
 }