Cache statements in @powersync/web

Author: simolus3

packages/web/src/db/adapters/options.ts

@@ -88,6 +88,12 @@ export interface WebSpecificOpenOptions {
    * A worker will be initialized if none is provided
    */
   workerPort?: MessagePort | undefined;
+
+  /**
+   * If set to a value greater than zero, the worker will cache prepared statements to avoid preparing them every time
+   * a query runs.
+   */
+  preparedStatementsCache?: number;
 }
 
 export interface ResolvedWebSQLOpenOptions extends SQLOpenOptions, WebSpecificOpenOptions {}

packages/web/src/db/adapters/wa-sqlite/RawSqliteConnection.ts

@@ -3,6 +3,7 @@ import { Factory as WaSqliteFactory, SQLITE_ROW } from '@journeyapps/wa-sqlite';
 import { loadModuleAndVfs, WASQLiteVFS } from './vfs.js';
 import { TemporaryStorageOption } from '../options.js';
 import { RawQueryResult, SqliteValue } from '@powersync/common';
+import { PreparedStatementCache } from './StatementCache.js';
 
 export interface RawWebResult extends Required<RawQueryResult> {
   autocommit: boolean;
@@ -23,6 +24,10 @@ export interface RawWaSqliteDatabaseOptions {
   encryptionKey: string | undefined;
   temporaryStorage: TemporaryStorageOption;
   cacheSizeKb: number;
+  /**
+   * The amount of prepared statements to cache, or 0 to disable caching.
+   */
+  preparedStatementsCache: number;
 }
 
 /**
@@ -33,12 +38,18 @@ export interface RawWaSqliteDatabaseOptions {
  */
 export class RawSqliteConnection {
   private _sqliteAPI: SQLiteAPI | null = null;
+  private sqlite3_stmt_isexplain!: (stmt: number) => 0 | 1 | 2;
+
   /**
    * The `sqlite3*` connection pointer.
    */
   private db: number = 0;
+  private statementCache: PreparedStatementCache | null;
 
-  constructor(readonly options: RawWaSqliteDatabaseOptions) {}
+  constructor(readonly options: RawWaSqliteDatabaseOptions) {
+    this.statementCache =
+      options.preparedStatementsCache > 0 ? new PreparedStatementCache(options.preparedStatementsCache) : null;
+  }
 
   get isOpen(): boolean {
     return this.db != 0;
@@ -62,6 +73,7 @@ export class RawSqliteConnection {
 
   private async openSQLiteAPI(): Promise<SQLiteAPI> {
     const { module, vfs } = await loadModuleAndVfs(this.options);
+    this.sqlite3_stmt_isexplain = module.cwrap('sqlite3_stmt_isexplain', 'int', ['int']);
     const sqlite3 = WaSqliteFactory(module);
     sqlite3.vfs_register(vfs, true);
     /**
@@ -141,7 +153,7 @@ export class RawSqliteConnection {
   async executeRaw(sql: string, bindings?: any[]): Promise<RawResultSet[]> {
     const results = [];
     const api = this.requireSqlite();
-    for await (const stmt of api.statements(this.db, sql)) {
+    for await (const stmt of this.cachedStatements(api, sql)) {
       let columns;
 
       const rs = await this.stepThroughStatement(api, stmt, bindings ?? [], columns);
@@ -192,8 +204,56 @@ export class RawSqliteConnection {
 
   async close() {
     if (this.isOpen) {
-      await this.requireSqlite().close(this.db);
+      const api = this.requireSqlite();
+
+      if (this.statementCache) {
+        for (const stmt of this.statementCache.drain()) {
+          await api.finalize(stmt);
+        }
+      }
+
+      await api.close(this.db);
       this.db = 0;
     }
   }
+
+  async *cachedStatements(api: SQLiteAPI, sql: string): AsyncIterable<number> {
+    {
+      const existing = this.statementCache?.lookup(sql);
+      if (existing != null) {
+        yield existing;
+        return;
+      }
+    }
+
+    const inner = api.statements(this.db, sql, { unscoped: true });
+    const preparedStatements: number[] = [];
+
+    try {
+      for await (const stmt of inner) {
+        preparedStatements.push(stmt);
+        yield stmt;
+      }
+    } finally {
+      // We can only cache statements if the sql text corresponds to a single statement, otherwise it's not clear what
+      // portion of the original sql text to use as a key.
+      if (preparedStatements.length == 1 && this.statementCache) {
+        const stmt = preparedStatements[0];
+        // Don't cache EXPLAIN statements, their result becomes invalid after schema changes.
+        if (this.sqlite3_stmt_isexplain(stmt) == 0) {
+          const evicted = this.statementCache.addStatement(sql, stmt);
+          if (evicted != null) {
+            await api.finalize(evicted);
+          }
+
+          return;
+        }
+      }
+
+      // We're not caching statements, so finalize them.
+      for (const stmt of preparedStatements) {
+        await api.finalize(stmt);
+      }
+    }
+  }
 }

packages/web/src/db/adapters/wa-sqlite/StatementCache.ts

@@ -0,0 +1,49 @@
+export class PreparedStatementCache {
+  readonly #size: number;
+  // Note that Map preserves insertion order, which allows using it as an LRU
+  // cache (with the first element being the first element to evict).
+  readonly #statements = new Map<string, number>();
+
+  constructor(size: number) {
+    this.#size = size;
+  }
+
+  /**
+   * Attempts to look up the cached sql statement, if it's currently cached.
+   */
+  lookup(sql: string): number | null {
+    const foundStatement = this.#statements.get(sql);
+    if (foundStatement != null) {
+      // Insert again to mark it as used.
+      this.#statements.set(sql, foundStatement);
+      return foundStatement;
+    }
+
+    return null;
+  }
+
+  /**
+   * Adds a new statement into the cache.
+   *
+   * If that exceeds the target size of the statement cache, returns an old statement to evict.
+   * The caller is responsible for freeing that statement.
+   */
+  addStatement(sql: string, statement: number): number | null {
+    this.#statements.set(sql, statement);
+
+    if (this.#statements.size > this.#size) {
+      for (const [k, v] of this.#statements.entries()) {
+        this.#statements.delete(k);
+        return v;
+      }
+    }
+
+    return null;
+  }
+
+  drain(): number[] {
+    const values = [...this.#statements.values()];
+    this.#statements.clear();
+    return values;
+  }
+}

packages/web/src/db/adapters/wa-sqlite/WASQLiteOpenFactory.ts

@@ -60,8 +60,16 @@ export class WASQLiteOpenFactory implements SQLOpenFactory {
   }
 
   async openConnection(): Promise<PoolConnection> {
-    const { enableMultiTabs, useWebWorker, vfs, dbFilename, encryptionKey, temporaryStorage, cacheSizeKb } =
-      this.options;
+    const {
+      enableMultiTabs,
+      useWebWorker,
+      vfs,
+      dbFilename,
+      encryptionKey,
+      temporaryStorage,
+      cacheSizeKb,
+      preparedStatementsCache
+    } = this.options;
 
     if (!enableMultiTabs) {
       this.logger.log({ level: LogLevels.warn, message: 'Multiple tabs are not enabled in this browser' });
@@ -78,7 +86,9 @@ export class WASQLiteOpenFactory implements SQLOpenFactory {
         vfs,
         encryptionKey,
         temporaryStorage,
-        cacheSizeKb
+        cacheSizeKb,
+        // TODO: Enable prepared statement cache by default?
+        preparedStatementsCache: preparedStatementsCache ?? 0
       };
     }
 

packages/web/tests/open.test.ts

@@ -72,7 +72,8 @@ describe('Open Methods', { sequential: true }, () => {
       cacheSizeKb: 0,
       encryptionKey: undefined,
       filename: '',
-      readonly: false
+      readonly: false,
+      preparedStatementsCache: 0
     };
     const connection = await server.openConnectionLocally(logger, options);
     const client = new DatabaseClient(

packages/web/tests/src/db/StatementCache.test.ts

@@ -0,0 +1,42 @@
+import { expect, test } from 'vitest';
+import { PreparedStatementCache } from '../../../src/db/adapters/wa-sqlite/StatementCache.js';
+import { generateTestDb } from '../../utils/testDb.js';
+import { Schema } from '@powersync/common';
+
+test('statement cache', () => {
+  const cache = new PreparedStatementCache(10);
+  for (let i = 0; i < 10; i++) {
+    cache.addStatement(`SELECT ${i}`, i);
+  }
+
+  for (let i = 0; i < 10; i++) {
+    expect(cache.lookup(`SELECT ${i}`)).toStrictEqual(i);
+  }
+
+  expect(cache.addStatement('SELECT 10', 10)).toStrictEqual(0);
+  expect(cache.lookup('SELECT 0')).toBeNull();
+  expect(cache.lookup('SELECT 1')).not.toBeNull();
+});
+
+test('does not cache explain statements', async () => {
+  const db = generateTestDb({
+    schema: new Schema({}),
+    database: {
+      dbFilename: `${crypto.randomUUID()}.db`,
+      preparedStatementsCache: 16
+    }
+  });
+
+  await db.execute('create table test(id integer primary key, description text)');
+  await db.execute('create index i1 on test(description)');
+  const { detail: firstPlan } = await db.get<{ detail: string }>(
+    'explain query plan select * from test where description = ?'
+  );
+  expect(firstPlan).toContain('USING COVERING INDEX i1');
+
+  await db.execute('drop index i1');
+  const { detail: secondPlan } = await db.get<{ detail: string }>(
+    'explain query plan select * from test where description = ?'
+  );
+  expect(secondPlan).not.toContain('USING COVERING INDEX i1');
+});