Simplify `SQLSchemaMigrations` by making it synchronous

[mw10013]

Is your feature request related to a problem? Please describe.

The implementation of SQLSchemaMigrations is asynchronous making it's use in constructors not as straight-forward and lacks immediate state accuracy for hasMigrationsToRun.

Describe the solution you'd like

A synchronous implementation of SQLSchemaMigrations.

Additional context

I have a proof-of-concept implementation along with tests. If you are open to it, I would be happy to submit a PR and flesh it out under your guidance.

export interface SQLSchemaMigration {
  idMonotonicInc: number;
  description: string;
  sql?: string;
}

export interface SQLSchemaMigrationsConfig {
  doStorage: DurableObjectStorage;
  migrations: SQLSchemaMigration[];
  kvKey?: string;
}

export class SQLSchemaMigrations {
  #_config: SQLSchemaMigrationsConfig;
  #_migrations: SQLSchemaMigration[];
  #_kvKey: string;
  #_lastMigrationId: number;

  constructor(config: SQLSchemaMigrationsConfig) {
    this.#_config = config;
    this.#_kvKey = config.kvKey ?? "__sql_migrations_lastID";

    const migrations = [...config.migrations];
    migrations.sort((a, b) => a.idMonotonicInc - b.idMonotonicInc);
    const idSeen = new Set<number>();
    migrations.forEach((m) => {
      if (m.idMonotonicInc < 0) {
        throw new Error(
          `Migration ID cannot be negative: ${String(m.idMonotonicInc)}`,
        );
      }
      if (idSeen.has(m.idMonotonicInc)) {
        throw new Error(
          `Duplicate migration ID detected: ${String(m.idMonotonicInc)}`,
        );
      }
      idSeen.add(m.idMonotonicInc);
    });

    this.#_migrations = migrations;

    this.#_lastMigrationId =
      this.#_config.doStorage.kv.get<number>(this.#_kvKey) ?? -1;
  }

  /**
   * Checks if there are any migrations that have not been run yet.
   * @returns `true` if there are pending migrations, `false` otherwise.
   */
  hasMigrationsToRun(): boolean {
    if (!this.#_migrations.length) {
      return false;
    }
    return (
      this.#_lastMigrationId <
      this.#_migrations[this.#_migrations.length - 1].idMonotonicInc
    );
  }

  runAll(sqlGen?: (idMonotonicInc: number) => string): {
    rowsRead: number;
    rowsWritten: number;
  } {
    const result = {
      rowsRead: 0,
      rowsWritten: 0,
    };

    if (!this.hasMigrationsToRun()) {
      return result;
    }

    // Find the index of the last successfully run migration (by exact ID match) to locate the start of pending ones
    const lastRunIndex = this.#_migrations.findIndex(
      (m) => m.idMonotonicInc === this.#_lastMigrationId,
    );
    // If no exact match (e.g., none run yet), run all; otherwise, slice from the next migration
    const migrationsToRun =
      lastRunIndex === -1
        ? this.#_migrations
        : this.#_migrations.slice(lastRunIndex + 1);

    if (migrationsToRun.length > 0) {
      this.#_config.doStorage.transactionSync(() => {
        migrationsToRun.forEach((m) => {
          const query = m.sql ?? sqlGen?.(m.idMonotonicInc);
          if (!query) {
            throw new Error(
              `migration with neither 'sql' nor 'sqlGen' provided: ${String(m.idMonotonicInc)}`,
            );
          }
          const cursor = this.#_config.doStorage.sql.exec(query);
          // Consume the cursor for accurate rowsRead/rowsWritten tracking.
          cursor.toArray();

          result.rowsRead += cursor.rowsRead;
          result.rowsWritten += cursor.rowsWritten;
        });

        this.#_config.doStorage.kv.put(
          this.#_kvKey,
          migrationsToRun[migrationsToRun.length - 1].idMonotonicInc,
        );
      });
      // Only update the instance property after the transaction succeeds.
      // If updated inside the transaction and it fails, the instance property
      // would be ahead of storage, causing hasMigrationsToRun() to incorrectly
      // return false on subsequent calls.
      this.#_lastMigrationId =
        migrationsToRun[migrationsToRun.length - 1].idMonotonicInc;
    }

    return result;
  }
}

import type { SQLSchemaMigration } from "../../workers/sql-schema-migrations";
import type { Env } from "./test-worker";
import { env, listDurableObjectIds, runInDurableObject } from "cloudflare:test";
import { describe, expect, it } from "vitest";
import { SQLSchemaMigrations } from "../../workers/sql-schema-migrations";
import { SQLMigrationsDO } from "./test-worker";

function makeM(state: DurableObjectState, migrations: SQLSchemaMigration[]) {
  return new SQLSchemaMigrations({
    doStorage: state.storage,
    migrations: migrations,
  });
}

declare module "cloudflare:test" {
  // eslint-disable-next-line @typescript-eslint/no-empty-object-type
  interface ProvidedEnv extends Env {}
}

describe("happy paths", () => {
  it("empty initial storage", async () => {
    // Check sending request directly to instance
    const id = env.SQL_MIGRATIONS_DO.idFromName("emptyDO");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        const m = makeM(state, [
          {
            idMonotonicInc: 1,
            description: "test default tables",
            sql: `SELECT * FROM sqlite_master;`,
          },
        ]);
        m.runAll();

        expect(state.storage.sql.databaseSize).toEqual(8192);
        return Promise.resolve();
      },
    );

    // Check direct access to instance fields and storage
    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        expect(await state.storage.get<number>("__sql_migrations_lastID")).toBe(
          1,
        );
      },
    );

    // Check IDs can be listed
    const ids = await listDurableObjectIds(env.SQL_MIGRATIONS_DO);
    expect(ids.length).toBe(1);
    // eslint-disable-next-line @typescript-eslint/no-unsafe-call, @typescript-eslint/no-unsafe-member-access
    expect(ids[0].equals(id)).toBe(true);
  });

  it("multiple DDL", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("emptyDO");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        expect(instance).toBeInstanceOf(SQLMigrationsDO);

        const res = makeM(state, [
          {
            idMonotonicInc: 1,
            description: "tbl1",
            sql: `CREATE TABLE users(name TEXT PRIMARY KEY, age INTEGER);`,
          },
          {
            idMonotonicInc: 2,
            description: "tbl2",
            sql: `CREATE TABLE IF NOT EXISTS usersActivities (activityType TEXT, userName TEXT, PRIMARY KEY (userName, activityType));`,
          },
        ]).runAll();

        expect(res).toEqual({
          rowsRead: 2,
          rowsWritten: 6,
        });
        return Promise.resolve();
      },
    );
  });

  it("multiple DDL and data inserts", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("data-test");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        makeM(state, [
          {
            idMonotonicInc: 1,
            description: "tbl1",
            sql: `CREATE TABLE users(name TEXT PRIMARY KEY, age INTEGER);`,
          },
          {
            idMonotonicInc: 2,
            description: "data",
            sql: `INSERT INTO users VALUES ('ironman', 100); INSERT INTO users VALUES ('thor', 9000);`,
          },
          {
            idMonotonicInc: 3,
            description: "data2",
            sql: `INSERT INTO users VALUES ('hulk', 5);`,
          },
        ]).runAll();

        let rows = state.storage.sql
          .exec<{ name: string; age: number }>(`SELECT * FROM users;`)
          .toArray();
        expect(rows).toEqual([
          { name: "ironman", age: 100 },
          { name: "thor", age: 9000 },
          { name: "hulk", age: 5 },
        ]);

        makeM(state, [
          // Edge case but even if you do not provide the previous migrations
          // it should run only from where it left off.
          // WARNING: Do not do this in production since running these for a first time
          // on a fresh DO will fail!!!
          {
            idMonotonicInc: 4,
            description: "data3",
            sql: `DELETE FROM users WHERE name = 'thor';`,
          },
        ]).runAll();

        rows = state.storage.sql
          .exec<{ name: string; age: number }>(`SELECT * FROM users;`)
          .toArray();
        expect(rows).toEqual([
          { name: "ironman", age: 100 },
          { name: "hulk", age: 5 },
        ]);
        return Promise.resolve();
      },
    );
  });

  it("run migrations multiple times", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("emptyDO");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        const m1 = [
          {
            idMonotonicInc: 1,
            description: "tbl1",
            sql: `CREATE TABLE IF NOT EXISTS users(name TEXT PRIMARY KEY, age INTEGER);`,
          },
          {
            idMonotonicInc: 2,
            description: "tbl2",
            sql: `CREATE TABLE IF NOT EXISTS usersActivities (activityType TEXT, userName TEXT, PRIMARY KEY (userName, activityType));`,
          },
        ];
        const r1 = makeM(state, m1).runAll();

        const m2 = m1.concat([
          {
            idMonotonicInc: 3,
            description: "round 2",
            sql: `
                    CREATE TABLE IF NOT EXISTS marvel (heroName TEXT);
                    CREATE TABLE IF NOT EXISTS marvelMovies (name TEXT PRIMARY KEY, releaseDateMs INTEGER);
                    `,
          },
        ]);
        const r2_1 = makeM(state, m2).runAll();

        for (let i = 0; i < 5; i++) {
          const r2_n = makeM(state, m2).runAll();
          // Should have no effect running the same migrations.
          expect(r2_n).toEqual({
            rowsRead: 0,
            rowsWritten: 0,
          });
        }

        return Promise.resolve({ r1, r2: r2_1 });
      },
    );
  });

  it("hasMigrationsToRun", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("hasMigrationsToRun");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        expect(makeM(state, []).hasMigrationsToRun()).toEqual(false);

        const migrations = [
          {
            idMonotonicInc: 1,
            description: "tbl1",
            sql: `CREATE TABLE users(name TEXT PRIMARY KEY, age INTEGER);`,
          },
        ];
        const m = makeM(state, migrations);
        expect(m.hasMigrationsToRun()).toEqual(true);
        m.runAll();
        expect(m.hasMigrationsToRun()).toEqual(false);

        const m1 = makeM(state, migrations);
        expect(m1.hasMigrationsToRun()).toEqual(false);

        return Promise.resolve();
      },
    );
  });
});

describe("expected exceptions", () => {
  it("duplicate IDs", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("emptyDO");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        expect(() =>
          makeM(state, [
            {
              idMonotonicInc: 1,
              description: "tbl1",
              sql: `CREATE TABLE users(name TEXT PRIMARY KEY, age INTEGER);`,
            },
            {
              idMonotonicInc: 2,
              description: "tbl1",
              sql: `CREATE TABLE IF NOT EXISTS users(name TEXT PRIMARY KEY, age INTEGER);`,
            },
            {
              idMonotonicInc: 1,
              description: "tbl2",
              sql: `CREATE TABLE IF NOT EXISTS usersActivities (activityType TEXT, userName TEXT, PRIMARY KEY (userName, activityType));`,
            },
          ]).runAll(),
        ).toThrowError("Duplicate migration ID detected: 1");
        return Promise.resolve();
      },
    );
  });

  it("negative IDs", async () => {
    const id = env.SQL_MIGRATIONS_DO.idFromName("emptyDO");
    const stub = env.SQL_MIGRATIONS_DO.get(id);

    await runInDurableObject(
      stub,
      async (instance: SQLMigrationsDO, state: DurableObjectState) => {
        expect(() =>
          makeM(state, [
            {
              idMonotonicInc: -1,
              description: "tbl1",
              sql: `CREATE TABLE users(name TEXT PRIMARY KEY, age INTEGER);`,
            },
          ]).runAll(),
        ).toThrowError("Migration ID cannot be negative: -1");
        return Promise.resolve();
      },
    );
  });
});