pipework 0.283.0 → 0.285.0

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

@@ -1,11 +1,14 @@
 import postgres from 'postgres'
+import { createHash } from 'node:crypto'
 import { getLogger } from '../../core/logging/logger.js'
 import { currentQueryLabel } from './query-label.js'
 import { DrizzleQueryError } from '../query/errors.js'
+import { parameterBytes, setOperationArms } from './query-size.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 STATEMENT_ID_HEX_CHARS = 16
 const SAMPLE_TIMEOUT_MS = 2_000
 const PG_TRUNCATED_QUERY_LENGTH = 1_000
 
@@ -18,6 +21,24 @@ interface BackendSample {
 
 const UNKNOWN_BACKEND: BackendSample = { pid: null, state: null, waitEventType: null, waitEvent: null }
 
+/** `observability.largeQuery*` budgets; 0 disables a budget and nothing is measured for it. */
+export interface LargeQueryBudgets {
+  readonly parameterBytes: number
+  readonly sqlChars: number
+  readonly arms: number
+  readonly rows: number
+}
+
+/** The budgets of a resolved `observability` section. */
+export function largeQueryBudgets(o: {
+  readonly largeQueryParamBytes: number
+  readonly largeQuerySqlChars: number
+  readonly largeQueryArms: number
+  readonly largeQueryRows: number
+}): LargeQueryBudgets {
+  return { parameterBytes: o.largeQueryParamBytes, sqlChars: o.largeQuerySqlChars, arms: o.largeQueryArms, rows: o.largeQueryRows }
+}
+
 /** What a pool's watch needs: where it logs, how slow is slow, and how to reach the server for a wait-event sample. */
 export interface StatementWatchOptions {
   readonly database: string
@@ -27,6 +48,8 @@ export interface StatementWatchOptions {
   readonly inFlightQueryMs: number | undefined
   /** `observability.onStatementSettled`; undefined: nothing extra is tracked or fired. */
   readonly onSettled?: StatementSettledObserver | undefined
+  /** Undefined: no `query:large` check. */
+  readonly budgets?: LargeQueryBudgets | undefined
 }
 
 export interface StatementWatchHandle {
@@ -61,8 +84,22 @@ function settledError(raw: unknown): StatementSettledError {
   }
 }
 
-function labelFor(sqlText: string): string {
-  return currentQueryLabel() ?? sqlText.replace(/\s+/g, ' ').trim().slice(0, LABEL_SQL_CHARS)
+function collapse(sqlText: string): string {
+  return sqlText.replace(/\s+/g, ' ').trim()
+}
+
+/** The first 120 characters of the whitespace-collapsed statement text: the `sql` of the log lines and `statementPrefix`. */
+function sqlPrefix(sqlText: string): string {
+  return collapse(sqlText).slice(0, LABEL_SQL_CHARS)
+}
+
+/**
+ * `v1:` + the first 16 hex characters (64 bits) of the SHA-256 of the whitespace-collapsed
+ * parameterized text. 64 bits keeps a collision among an app's few thousand statement shapes
+ * negligible (birthday bound near 4 billion shapes) while the id stays short in a log line.
+ */
+function statementIdOf(sqlText: string): string {
+  return `v1:${createHash('sha256').update(collapse(sqlText)).digest('hex').slice(0, STATEMENT_ID_HEX_CHARS)}`
 }
 
 /**
@@ -72,7 +109,9 @@ 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, onSettled } = options
+  const { database, url, slowQueryMs, inFlightQueryMs, onSettled, budgets } = options
+  const watchesSize = budgets !== undefined && (budgets.parameterBytes > 0 || budgets.sqlChars > 0 || budgets.arms > 0 || budgets.rows > 0)
+  const needsOutcome = onSettled !== undefined || (budgets?.rows ?? 0) > 0
   // 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
@@ -119,6 +158,25 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
     }
   }
 
+  /** One `query:large` per statement over any enabled budget, carrying every enabled measurement. Never throws. */
+  function warnIfLarge(sqlText: string, parameters: unknown, outcome: { result: unknown } | { error: unknown } | undefined, b: LargeQueryBudgets): void {
+    try {
+      const measured: Record<string, number> = {}
+      const over: string[] = []
+      const check = (name: string, budget: number, measure: () => number): void => {
+        if (budget <= 0) return
+        const value = measure()
+        measured[name] = value
+        if (value > budget) over.push(name)
+      }
+      check('parameterBytes', b.parameterBytes, () => parameterBytes(parameters))
+      check('sqlChars', b.sqlChars, () => sqlText.length)
+      check('arms', b.arms, () => setOperationArms(sqlText))
+      check('rows', b.rows, () => (outcome !== undefined && 'result' in outcome ? rowsOf(outcome.result) : 0))
+      if (over.length > 0) getLogger().warn({ label: currentQueryLabel() ?? null, sql: sqlPrefix(sqlText), database, over, ...measured }, 'query:large')
+    } catch { /* a measurement must never affect the statement */ }
+  }
+
   const watch: StatementWatch = async <T>(sqlText: string, run: (seen: (pending: object) => void) => Promise<T>): Promise<T> => {
     const startedMs = Date.now()
     let pendingQuery: object | undefined
@@ -128,13 +186,14 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
     let settled = false
     if (inFlightQueryMs !== undefined) {
       const log = getLogger()
-      const label = labelFor(sqlText)
+      const label = currentQueryLabel() ?? null
+      const sql = sqlPrefix(sqlText)
       interval = setInterval(() => {
         const elapsedS = Math.round((Date.now() - startedMs) / 100) / 10
         void sampleBackend(sqlText, startedMs).then(backend => {
           if (settled) return
           log.warn({
-            label, elapsedS, database,
+            label, sql, elapsedS, database,
             pid: backend.pid, state: backend.state,
             waitEventType: backend.waitEventType, waitEvent: backend.waitEvent,
           }, 'query:in-flight')
@@ -144,23 +203,26 @@ export function createStatementWatch(options: StatementWatchOptions): StatementW
     }
     try {
       const result = await run(seen)
-      if (onSettled) outcome = { result }
+      if (needsOutcome) outcome = { result }
       return result
     } catch (err) {
-      if (onSettled) outcome = { error: err }
+      if (needsOutcome) outcome = { error: err }
       throw err
     } finally {
       settled = true
       if (interval !== undefined) clearInterval(interval)
       const durationMs = Date.now() - startedMs
       if (slowQueryMs > 0 && durationMs > slowQueryMs) {
-        getLogger().warn({ label: labelFor(sqlText), durationMs, database }, 'query:slow')
+        getLogger().warn({ label: currentQueryLabel() ?? null, sql: sqlPrefix(sqlText), durationMs, database }, 'query:slow')
       }
+      const parameters = pendingQuery === undefined ? undefined : (pendingQuery as { parameters?: unknown }).parameters
+      if (watchesSize) warnIfLarge(sqlText, parameters, outcome, budgets)
       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,
+          statementId: statementIdOf(sqlText),
+          statementPrefix: sqlPrefix(sqlText),
           durationMs,
           rows: 'result' in outcome ? rowsOf(outcome.result) : 0,
           connectionId,

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,7 +127,7 @@ async function createIsolatedConnection(
   observability: ResolvedObservability,
   withSystem: boolean,
 ): Promise<PostgresJsDatabase> {
-  const watch = createStatementWatch({ database: name, url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs, onSettled: observability.onStatementSettled })
+  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)
@@ -178,7 +178,7 @@ export function createSetupConnection(
   observability?: ResolvedObservability,
 ): { drizzle: PostgresJsDatabase; close: () => Promise<void> } {
   const watch = observability
-    ? createStatementWatch({ database: 'setup', url: testDb.url, slowQueryMs: observability.slowQueryMs, inFlightQueryMs: observability.inFlightQueryMs, onSettled: observability.onStatementSettled })
+    ? 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 } : {}) })