perf(zqlite): fetch-template cache, direct SQL formatting, single-shot iterator logging

packages/zqlite/src/db.test.ts

@@ -125,6 +125,60 @@ test('slow queries are logged', () => {
   ]);
 });
 
+test('an iterator logs its slow-query epilogue exactly once', () => {
+  const sink = new TestLogSink();
+  const lc = new LogContext('debug', undefined, sink);
+
+  // threshold is 0 so every query is logged
+  const db = new Database(lc, ':memory:', undefined, 0);
+  db.exec('CREATE TABLE foo (id INTEGER PRIMARY KEY)');
+  db.exec('INSERT INTO foo (id) VALUES (1), (2), (3)');
+  const stmt = db.prepare('SELECT * FROM foo');
+
+  function iterateLogsFor(use: (it: IterableIterator<unknown>) => void) {
+    sink.messages.length = 0;
+    use(stmt.iterate());
+    return sink.messages
+      .filter(([, context]) => context?.method === 'iterate')
+      .map(([, context]) => context?.type);
+  }
+
+  // Exhausted and then closed by the caller, which is what `TableSource`
+  // does: the iterator sees `done` and the `finally` closes it anyway.
+  expect(
+    iterateLogsFor(it => {
+      for (const _ of it) {
+        // drain
+      }
+      it.return?.();
+    }),
+  ).toEqual(['total', 'sqlite']);
+
+  // Closed early, before the native iterator reported `done`.
+  expect(
+    iterateLogsFor(it => {
+      it.next();
+      it.return?.();
+    }),
+  ).toEqual(['total', 'sqlite']);
+
+  // Aborted, then closed.
+  expect(
+    iterateLogsFor(it => {
+      it.next();
+      try {
+        it.throw?.(new Error('boom'));
+      } catch {
+        // the native iterator rethrows
+      }
+      it.return?.();
+    }),
+  ).toEqual(['total', 'sqlite']);
+
+  // The statement is still usable, so every iterator was closed.
+  expect(stmt.all()).toHaveLength(3);
+});
+
 test('sql errors are annotated with sql', () => {
   const sink = new TestLogSink();
   const lc = new LogContext('debug', undefined, sink);

packages/zqlite/src/db.ts

@@ -134,12 +134,7 @@ export class Database implements Disposable {
       }
       throw e;
     } finally {
-      logIfSlow(
-        this.#lc.withContext('method', method),
-        performance.now() - start,
-        {method},
-        this.#threshold,
-      );
+      logIfSlow(this.#lc, performance.now() - start, this.#threshold, method);
     }
   }
 
@@ -207,10 +202,11 @@ export class Statement {
     const start = performance.now();
     const ret = this.#stmt.run(...params);
     logIfSlow(
-      this.#lc.withContext('method', 'run'),
+      this.#lc,
       performance.now() - start,
-      {...this.#attrs, method: 'run'},
       this.#threshold,
+      'run',
+      this.#attrs,
     );
     return ret;
   }
@@ -219,10 +215,11 @@ export class Statement {
     const start = performance.now();
     const ret = this.#stmt.get(...params);
     logIfSlow(
-      this.#lc.withContext('method', 'get'),
+      this.#lc,
       performance.now() - start,
-      {...this.#attrs, method: 'get'},
       this.#threshold,
+      'get',
+      this.#attrs,
     );
     return ret as T;
   }
@@ -231,17 +228,18 @@ export class Statement {
     const start = performance.now();
     const ret = this.#stmt.all(...params);
     logIfSlow(
-      this.#lc.withContext('method', 'all'),
+      this.#lc,
       performance.now() - start,
-      {...this.#attrs, method: 'all'},
       this.#threshold,
+      'all',
+      this.#attrs,
     );
     return ret as T[];
   }
 
   iterate<T>(...params: unknown[]): IterableIterator<T> {
     return new LoggingIterableIterator(
-      this.#lc.withContext('method', 'iterate'),
+      this.#lc,
       this.#attrs,
       this.#stmt.iterate(...params),
       this.#threshold,
@@ -256,6 +254,7 @@ class LoggingIterableIterator<T> implements IterableIterator<T> {
   readonly #attrs: Attributes;
   #start: number;
   #sqliteRowTimeSum: number;
+  #logged = false;
 
   constructor(
     lc: LogContext,
@@ -282,18 +281,32 @@ class LoggingIterableIterator<T> implements IterableIterator<T> {
     return ret;
   }
 
+  /**
+   * Callers that exhaust an iterator still close it afterwards -- `#fetch`
+   * closes in a `finally` -- so an ordinary completion reaches here through
+   * `next()` and then again through `return()`. Only the first one describes
+   * the query, so the epilogue runs once per iterator.
+   */
   #log() {
+    if (this.#logged) {
+      return;
+    }
+    this.#logged = true;
     logIfSlow(
-      this.#lc.withContext('type', 'total'),
+      this.#lc,
       performance.now() - this.#start,
-      {...this.#attrs, type: 'total', method: 'iterate'},
       this.#threshold,
+      'iterate',
+      this.#attrs,
+      'total',
     );
     logIfSlow(
-      this.#lc.withContext('type', 'sqlite'),
+      this.#lc,
       this.#sqliteRowTimeSum,
-      {...this.#attrs, type: 'sqlite', method: 'iterate'},
       this.#threshold,
+      'iterate',
+      this.#attrs,
+      'sqlite',
     );
   }
 
@@ -314,19 +327,29 @@ class LoggingIterableIterator<T> implements IterableIterator<T> {
   }
 }
 
+/**
+ * Building the LogContext and the attribute object is only worth it when the
+ * query is actually slow, which is the rare case: this runs on every statement
+ * execution, and twice more for every iterator that finishes.
+ */
 function logIfSlow(
   lc: LogContext,
   elapsed: number,
-  attrs: Attributes,
   threshold: number,
+  method: string,
+  attrs?: Attributes,
+  type?: string,
 ): void {
-  if (elapsed >= threshold) {
-    for (const [key, value] of Object.entries(attrs)) {
-      lc = lc.withContext(key, value);
-    }
-    lc.warn?.('Slow SQLite query', elapsed);
-    manualSpan(tracer, 'db.slow-query', elapsed, attrs);
+  if (elapsed < threshold) {
+    return;
+  }
+  const allAttrs: Attributes =
+    type === undefined ? {...attrs, method} : {...attrs, method, type};
+  for (const [key, value] of Object.entries(allAttrs)) {
+    lc = lc.withContext(key, value);
   }
+  lc.warn?.('Slow SQLite query', elapsed);
+  manualSpan(tracer, 'db.slow-query', elapsed, allAttrs);
 }
 
 /**

packages/zqlite/src/internal/sql.test.ts

@@ -1,5 +1,7 @@
+import {escapeSQLiteIdentifier} from '@databases/escape-identifier';
+import type {FormatConfig} from '@databases/sql';
 import {expect, test} from 'vitest';
-import {compile, sql} from './sql.ts';
+import {compile, format, sql} from './sql.ts';
 
 test('can do empty slots', () => {
   const str = compile(sql`INSERT INTO foo (id, name) VALUES (?, ?)`);
@@ -17,3 +19,39 @@ test('escapes identifiers as advertised', () => {
   const str = compile(sql`SELECT * FROM ${sql.ident('foo"bar')}`);
   expect(str).toMatchInlineSnapshot(`"SELECT * FROM "foo""bar""`);
 });
+
+// `format` walks the SQL items itself instead of using `@databases/sql`'s
+// generic `formatStandard`, so it has to keep producing exactly what
+// `formatStandard` produces for every shape the query builder emits.
+test('matches the generic @databases/sql formatter', () => {
+  const generic: FormatConfig = {
+    escapeIdentifier: str => escapeSQLiteIdentifier(str),
+    formatValue: value => ({placeholder: '?', value}),
+  };
+
+  for (const query of [
+    sql`SELECT 1`,
+    sql`SELECT ${sql.join(
+      ['a', 'b'].map(c => sql.ident(c)),
+      sql`,`,
+    )} FROM ${sql.ident('foo')} WHERE ${sql.ident('a')} = ${'x'} AND ${sql.ident(
+      'b',
+    )} IS ${null} ORDER BY ${sql.ident('a')} ${sql.__dangerous__rawValue(
+      'desc',
+    )}`,
+    sql`SELECT * FROM ${sql.ident('sch"ema', 'ta ble')} WHERE ${sql.ident(
+      'n',
+    )} IN (${sql.join(
+      [1, 2, 3].map(v => sql`${v}`),
+      sql`,`,
+    )})`,
+    sql`
+      SELECT ${sql.ident('a')}
+        FROM ${sql.ident('foo')}
+        WHERE ${sql.ident('a')} = ${1}
+        ORDER BY ${sql.ident('a')}`,
+    sql`   `,
+  ]) {
+    expect(format(query)).toEqual(query.format(generic));
+  }
+});

packages/zqlite/src/internal/sql.ts

@@ -1,18 +1,85 @@
 import {escapeSQLiteIdentifier} from '@databases/escape-identifier';
-import type {FormatConfig, SQLQuery} from '@databases/sql';
-import sql from '@databases/sql';
+import type {SQLItem, SQLQuery} from '@databases/sql';
+import sql, {SQLItemType} from '@databases/sql';
 
-const sqliteFormat: FormatConfig = {
-  escapeIdentifier: str => escapeSQLiteIdentifier(str),
-  formatValue: value => ({placeholder: '?', value}),
-};
+/**
+ * Identifiers are table and column names, so the working set is the size of
+ * the schema. The bound is only there so that a client-influenced query shape
+ * cannot grow the map without limit; clearing is cheaper than tracking
+ * recency for something that refills in a handful of fetches.
+ */
+const MAX_CACHED_IDENTIFIERS = 10_000;
+
+const escapedIdentifiers = new Map<string, string>();
+
+function escapeIdentifier(name: string): string {
+  let escaped = escapedIdentifiers.get(name);
+  if (escaped === undefined) {
+    escaped = escapeSQLiteIdentifier(name);
+    if (escapedIdentifiers.size >= MAX_CACHED_IDENTIFIERS) {
+      escapedIdentifiers.clear();
+    }
+    escapedIdentifiers.set(name, escaped);
+  }
+  return escaped;
+}
+
+/**
+ * SQLite-specific replacement for `@databases/sql`'s generic `formatStandard`.
+ *
+ * It produces byte-identical output, but skips the generic formatter's
+ * dedent pass (split/filter/regex/`Math.min` over every line) for the
+ * single-line SQL that the query builder emits, and reuses escaped
+ * identifiers instead of re-escaping the same column names on every fetch.
+ * `TableSource.#fetch` formats a query for every fetch, so this runs on the
+ * hot read path.
+ */
+function formatSQLite(items: readonly SQLItem[]): {
+  text: string;
+  values: unknown[];
+} {
+  let text = '';
+  const values: unknown[] = [];
+  for (const item of items) {
+    switch (item.type) {
+      case SQLItemType.RAW:
+        text += item.text;
+        break;
+      case SQLItemType.VALUE:
+        text += '?';
+        values.push(item.value);
+        break;
+      case SQLItemType.IDENTIFIER:
+        text +=
+          item.names.length === 1
+            ? escapeIdentifier(item.names[0])
+            : item.names.map(name => escapeIdentifier(name)).join('.');
+        break;
+    }
+  }
+  // Multi-line templates are dedented by the common indent, as
+  // `formatStandard` does. Single-line text is unaffected by that pass, since
+  // the trailing `trim()` removes the leading whitespace either way.
+  if (text.includes('\n') && text.trim()) {
+    const lines = text.split('\n');
+    const min = Math.min(
+      ...lines
+        .filter(line => line.trim() !== '')
+        .map(line => line.length - line.trimStart().length),
+    );
+    if (min) {
+      text = lines.map(line => line.substring(min)).join('\n');
+    }
+  }
+  return {text: text.trim(), values};
+}
 
 export function compile(sql: SQLQuery): string {
-  return sql.format(sqliteFormat).text;
+  return sql.format(formatSQLite).text;
 }
 
 export function format(sql: SQLQuery) {
-  return sql.format(sqliteFormat);
+  return sql.format(formatSQLite);
 }
 
 export {sql};

packages/zqlite/src/internal/statement-cache.ts

@@ -107,7 +107,15 @@ export class StatementCache {
    * @returns
    */
   get(sql: string): CachedStatement {
-    sql = normalizeWhitespace(sql);
+    return this.getNormalized(normalizeWhitespace(sql));
+  }
+
+  /**
+   * {@linkcode get} for callers that already hold normalized sql, so that a
+   * caller reusing one canonical string across many lookups does not re-run
+   * the normalizing regex every time.
+   */
+  getNormalized(sql: string): CachedStatement {
     const statements = this.#cache.get(sql);
     if (statements && statements.length > 0) {
       const statement = statements.pop()!;
@@ -175,6 +183,6 @@ export class StatementCache {
   }
 }
 
-function normalizeWhitespace(sql: string) {
+export function normalizeWhitespace(sql: string) {
   return sql.replaceAll(/\s+/g, ' ');
 }