WRNexusJS/packages/db/src/performance.ts

import type { Db, ExecResult, Row } from "./driver.ts";
import type { Model } from "./schema.ts";

export interface QueryRecord {
  sql: string;
  paramsCount: number;
  durationMs: number;
  rowCount?: number;
  operation: "all" | "one" | "exec";
  duplicateCount: number;
}

export interface QueryIssue {
  code: string;
  severity: "error" | "warning" | "info";
  message: string;
  sql: string;
}

export interface QueryPolicy {
  slowQueryMs?: number;
  timeoutMs?: number;
  maxRows?: number;
  duplicateWarningCount?: number;
  warnSelectStar?: boolean;
  warnUnboundedSelect?: boolean;
  onQuery?: (record: QueryRecord) => void | Promise<void>;
  onIssue?: (issue: QueryIssue) => void | Promise<void>;
}

const recentQueries: QueryRecord[] = [];
const recentQueryIssues: QueryIssue[] = [];
const TELEMETRY_LIMIT = 200;

export function getDbPerformanceSnapshot(): { queries: QueryRecord[]; issues: QueryIssue[] } {
  return { queries: structuredClone(recentQueries), issues: structuredClone(recentQueryIssues) };
}

export function resetDbPerformanceSnapshot(): void {
  recentQueries.length = 0;
  recentQueryIssues.length = 0;
}

function remember<T>(items: T[], value: T): void {
  items.push(structuredClone(value));
  if (items.length > TELEMETRY_LIMIT) items.splice(0, items.length - TELEMETRY_LIMIT);
}

function normalizedSql(sql: string): string {
  return sql
    .replace(/--.*$/gm, " ")
    .replace(/\/\*[\s\S]*?\*\//g, " ")
    .replace(/\s+/g, " ")
    .trim();
}

function operationName(sql: string): string {
  return normalizedSql(sql).split(" ")[0]?.toUpperCase() ?? "UNKNOWN";
}

async function withTimeout<T>(promise: Promise<T>, timeoutMs: number, sql: string): Promise<T> {
  let timer: ReturnType<typeof setTimeout> | undefined;
  try {
    return await Promise.race([
      promise,
      new Promise<T>((_resolve, reject) => {
        timer = setTimeout(
          () =>
            reject(
              new Error(`WRN-DB-TIMEOUT: query exceeded ${timeoutMs} ms: ${sql.slice(0, 120)}`),
            ),
          timeoutMs,
        );
      }),
    ]);
  } finally {
    if (timer) clearTimeout(timer);
  }
}

export function instrumentDb(db: Db, policy: QueryPolicy = {}): Db {
  const counts = new Map<string, number>();
  const slowQueryMs = policy.slowQueryMs ?? 100;
  const timeoutMs = policy.timeoutMs ?? 30_000;
  const maxRows = policy.maxRows ?? 10_000;
  const duplicateWarningCount = policy.duplicateWarningCount ?? 5;

  const inspect = async (sql: string): Promise<void> => {
    const normalized = normalizedSql(sql);
    if (policy.warnSelectStar !== false && /^SELECT\s+\*/i.test(normalized)) {
      const issue: QueryIssue = {
        code: "WRN-DB-SELECT-STAR",
        severity: "warning",
        message: "Avoid SELECT * in production queries.",
        sql,
      };
      remember(recentQueryIssues, issue);
      await policy.onIssue?.(issue);
    }
    if (
      policy.warnUnboundedSelect !== false &&
      /^SELECT\b/i.test(normalized) &&
      !/\bLIMIT\b/i.test(normalized) &&
      !/\bCOUNT\s*\(/i.test(normalized)
    ) {
      const issue: QueryIssue = {
        code: "WRN-DB-UNBOUNDED-SELECT",
        severity: "warning",
        message: "SELECT query has no LIMIT. Prefer cursor pagination for large datasets.",
        sql,
      };
      remember(recentQueryIssues, issue);
      await policy.onIssue?.(issue);
    }
  };

  const record = async (
    sql: string,
    params: unknown[],
    operation: QueryRecord["operation"],
    durationMs: number,
    rowCount?: number,
  ): Promise<void> => {
    const key = `${operation}:${normalizedSql(sql)}:${JSON.stringify(params)}`;
    const duplicateCount = (counts.get(key) ?? 0) + 1;
    counts.set(key, duplicateCount);
    const queryRecord: QueryRecord = {
      sql,
      paramsCount: params.length,
      durationMs,
      ...(rowCount === undefined ? {} : { rowCount }),
      operation,
      duplicateCount,
    };
    remember(recentQueries, queryRecord);
    await policy.onQuery?.(queryRecord);
    if (durationMs >= slowQueryMs) {
      const issue: QueryIssue = {
        code: "WRN-DB-SLOW-QUERY",
        severity: durationMs >= slowQueryMs * 5 ? "error" : "warning",
        message: `Query took ${durationMs.toFixed(2)} ms.`,
        sql,
      };
      remember(recentQueryIssues, issue);
      await policy.onIssue?.(issue);
    }
    if (duplicateCount === duplicateWarningCount) {
      const issue: QueryIssue = {
        code: "WRN-DB-DUPLICATE-QUERY",
        severity: "warning",
        message: `The same query ran ${duplicateCount} times in one request scope (possible N+1).`,
        sql,
      };
      remember(recentQueryIssues, issue);
      await policy.onIssue?.(issue);
    }
  };

  const wrapped: Db = {
    driver: db.driver,
    async all<T = Row>(sql: string, params: unknown[] = [], model?: Model<T>): Promise<T[]> {
      await inspect(sql);
      const started = performance.now();
      const rows = await withTimeout(db.all(sql, params, model), timeoutMs, sql);
      const duration = performance.now() - started;
      if (rows.length > maxRows) {
        await policy.onIssue?.({
          code: "WRN-DB-MAX-ROWS",
          severity: "error",
          message: `Query returned ${rows.length} rows; maximum is ${maxRows}.`,
          sql,
        });
        throw new Error(`WRN-DB-MAX-ROWS: query returned ${rows.length} rows.`);
      }
      await record(sql, params, "all", duration, rows.length);
      return rows;
    },
    async one<T = Row>(sql: string, params: unknown[] = [], model?: Model<T>): Promise<T | null> {
      await inspect(sql);
      const started = performance.now();
      const row = await withTimeout(db.one(sql, params, model), timeoutMs, sql);
      await record(sql, params, "one", performance.now() - started, row ? 1 : 0);
      return row;
    },
    async exec(sql: string, params: unknown[] = []): Promise<ExecResult> {
      const started = performance.now();
      const result = await withTimeout(db.exec(sql, params), timeoutMs, sql);
      await record(sql, params, "exec", performance.now() - started, result.changes);
      return result;
    },
    tx<T>(fn: (transaction: Db) => Promise<T>): Promise<T> {
      return db.tx((transaction) => fn(instrumentDb(transaction, policy)));
    },
    createTable(model) {
      return db.createTable(model);
    },
    close() {
      return db.close();
    },
  };
  return wrapped;
}

export function queryOperation(sql: string): string {
  return operationName(sql);
}