@prairielearn/migrations 1.1.0 → 1.2.1

package/src/batched-migrations/batched-migrations-runner.ts

@@ -0,0 +1,327 @@
+import EventEmitter from 'node:events';
+import path from 'node:path';
+import { setTimeout as sleep } from 'node:timers/promises';
+import { loadSqlEquiv, queryValidatedZeroOrOneRow } from '@prairielearn/postgres';
+import { doWithLock, tryWithLock } from '@prairielearn/named-locks';
+
+import { MigrationFile, readAndValidateMigrationsFromDirectories } from '../load-migrations';
+import {
+  BatchedMigrationRowSchema,
+  BatchedMigrationRow,
+  insertBatchedMigration,
+  BatchedMigrationStatus,
+  selectBatchedMigrationForTimestamp,
+  updateBatchedMigrationStatus,
+  BatchedMigrationImplementation,
+  validateBatchedMigrationImplementation,
+} from './batched-migration';
+import { BatchedMigrationRunner } from './batched-migration-runner';
+
+const sql = loadSqlEquiv(__filename);
+
+const DEFAULT_MIN_VALUE = 1n;
+const DEFAULT_BATCH_SIZE = 1_000;
+const DEFAULT_WORK_DURATION_MS = 60_000;
+const DEFAULT_SLEEP_DURATION_MS = 30_000;
+const EXTENSIONS = ['.js', '.ts', '.mjs', '.mts'];
+
+interface BatchedMigrationRunnerOptions {
+  project: string;
+  directories: string[];
+}
+
+interface BatchedMigrationStartOptions {
+  workDurationMs?: number;
+  sleepDurationMs?: number;
+}
+
+interface BatchedMigrationFinalizeOptions {
+  logProgress?: boolean;
+}
+
+export class BatchedMigrationsRunner extends EventEmitter {
+  private readonly options: BatchedMigrationRunnerOptions;
+  private readonly lockName: string;
+  private running = false;
+  private migrationFiles: MigrationFile[] | null = null;
+  private abortController = new AbortController();
+
+  constructor(options: BatchedMigrationRunnerOptions) {
+    super();
+    this.options = options;
+    this.lockName = `batched-migrations:${this.options.project}`;
+  }
+
+  private lockNameForTimestamp(timestamp: string) {
+    return `${this.lockName}:${timestamp}`;
+  }
+
+  private getMigrationFiles = async () => {
+    if (!this.migrationFiles) {
+      this.migrationFiles = await readAndValidateMigrationsFromDirectories(
+        this.options.directories,
+        EXTENSIONS
+      );
+    }
+    return this.migrationFiles;
+  };
+
+  private async getMigrationForIdentifier(identifier: string): Promise<MigrationFile | null> {
+    const timestamp = identifier.split('_')[0];
+
+    const migrationFiles = await this.getMigrationFiles();
+    const migrationFile = migrationFiles.find((m) => m.timestamp === timestamp);
+    return migrationFile ?? null;
+  }
+
+  /**
+   * Loads the implementation for the migration with the given identifier. The identifier
+   * must start with a 14-character timestamp. It may optionally be followed by
+   * an underscore with additional characters, which are ignored. These should
+   * typically be used to provide a human-readable name for the migration.
+   */
+  private async loadMigrationImplementation(migrationFile: MigrationFile) {
+    // We use dynamic imports to handle both CJS and ESM modules.
+    const migrationModulePath = path.join(migrationFile.directory, migrationFile.filename);
+    const migrationModule = await import(migrationModulePath);
+
+    const migrationImplementation = migrationModule.default as BatchedMigrationImplementation;
+    validateBatchedMigrationImplementation(migrationImplementation);
+    return migrationImplementation;
+  }
+
+  async enqueueBatchedMigration(identifier: string) {
+    const migrationFile = await this.getMigrationForIdentifier(identifier);
+    if (!migrationFile) {
+      throw new Error(`No migration found for identifier ${identifier}`);
+    }
+
+    const migrationImplementation = await this.loadMigrationImplementation(migrationFile);
+    const migrationParameters = await migrationImplementation.getParameters();
+
+    // If `max` is null, that implies that there are no rows to process, so
+    // we can immediately mark the migration as finished.
+    const status: BatchedMigrationStatus =
+      migrationParameters.max === null ? 'succeeded' : 'pending';
+
+    const minValue = BigInt(migrationParameters.min ?? DEFAULT_MIN_VALUE);
+    const maxValue = BigInt(migrationParameters.max ?? minValue);
+    const batchSize = migrationParameters.batchSize ?? DEFAULT_BATCH_SIZE;
+
+    await insertBatchedMigration({
+      project: this.options.project,
+      filename: migrationFile.filename,
+      timestamp: migrationFile.timestamp,
+      batch_size: batchSize,
+      min_value: minValue,
+      max_value: maxValue,
+      status,
+    });
+  }
+
+  async finalizeBatchedMigration(identifier: string, options?: BatchedMigrationFinalizeOptions) {
+    const timestamp = identifier.split('_')[0];
+
+    let migration = await selectBatchedMigrationForTimestamp(this.options.project, timestamp);
+
+    if (migration.status === 'succeeded') return;
+
+    // If the migration isn't already in the finalizing state, mark it as such.
+    if (migration.status !== 'finalizing') {
+      migration = await updateBatchedMigrationStatus(migration.id, 'finalizing');
+    }
+
+    await doWithLock(this.lockNameForTimestamp(timestamp), { autoRenew: true }, async () => {
+      const migrationFile = await this.getMigrationForIdentifier(identifier);
+      if (!migrationFile) {
+        throw new Error(`No migration found for identifier ${identifier}`);
+      }
+      const migrationImplementation = await this.loadMigrationImplementation(migrationFile);
+
+      const runner = new BatchedMigrationRunner(migration, migrationImplementation, {
+        // Always log progress unless explicitly disabled.
+        logProgress: options?.logProgress ?? true,
+      });
+
+      // Because we don't give any arguments to `run()`, it will run until it
+      // has attempted every job.
+      await runner.run();
+    });
+
+    migration = await selectBatchedMigrationForTimestamp(this.options.project, timestamp);
+
+    if (migration.status === 'succeeded') return;
+
+    throw new Error(
+      `Expected batched migration with identifier ${identifier} to be marked as 'succeeded', but it is '${migration.status}'.`
+    );
+  }
+
+  start(options: BatchedMigrationStartOptions = {}) {
+    if (this.running) {
+      throw new Error('BatchedMigrationsRunner is already running');
+    }
+
+    this.loop(options);
+  }
+
+  async loop({ workDurationMs, sleepDurationMs }: BatchedMigrationStartOptions) {
+    workDurationMs ??= DEFAULT_WORK_DURATION_MS;
+    sleepDurationMs ??= DEFAULT_SLEEP_DURATION_MS;
+
+    this.running = true;
+    while (this.running) {
+      if (this.abortController.signal.aborted) {
+        // We assign this here so that `stop()` can tell when this loop is done
+        // processing jobs.
+        this.running = false;
+        return;
+      }
+
+      let didWork = false;
+      try {
+        didWork = await this.maybePerformWork(workDurationMs);
+      } catch (err) {
+        this.emit('error', err);
+      }
+
+      // If we did work, we'll immediately try again since there's probably more
+      // work to be done. If not, we'll sleep for a while - maybe some more work
+      // will become available!
+      if (!didWork) {
+        // We provide the signal here so that we can more quickly stop things
+        // when we're shutting down.
+        try {
+          await sleep(sleepDurationMs, null, { ref: false, signal: this.abortController.signal });
+        } catch (err) {
+          // We don't care about errors here, they should only ever occur when
+          // the AbortController is aborted. Continue to the next iteration of
+          // the loop so we can shut down.
+          continue;
+        }
+      }
+    }
+  }
+
+  private async getOrStartMigration(): Promise<BatchedMigrationRow | null> {
+    return doWithLock(this.lockName, {}, async () => {
+      let migration = await queryValidatedZeroOrOneRow(
+        sql.select_running_migration,
+        { project: this.options.project },
+        BatchedMigrationRowSchema
+      );
+
+      if (!migration) {
+        migration = await queryValidatedZeroOrOneRow(
+          sql.start_next_pending_migration,
+          { project: this.options.project },
+          BatchedMigrationRowSchema
+        );
+      }
+
+      return migration;
+    });
+  }
+
+  async maybePerformWork(durationMs: number): Promise<boolean> {
+    const migration = await this.getOrStartMigration();
+    if (!migration) {
+      // No work to do. Handle this case.
+      return false;
+    }
+
+    // This server may not yet know about the current running migration. If
+    // that's the case, we'll just skip it for now.
+    const migrationFile = await this.getMigrationForIdentifier(migration.timestamp);
+    if (!migrationFile) {
+      return false;
+    }
+
+    let didWork = false;
+    await tryWithLock(
+      this.lockNameForTimestamp(migrationFile.timestamp),
+      { autoRenew: true },
+      async () => {
+        didWork = true;
+        const migrationImplementation = await this.loadMigrationImplementation(migrationFile);
+
+        const runner = new BatchedMigrationRunner(migration, migrationImplementation);
+
+        try {
+          await runner.run({ signal: this.abortController.signal, durationMs });
+        } catch (err) {
+          this.emit('error', err);
+        }
+      }
+    );
+
+    return didWork;
+  }
+
+  async stop() {
+    this.abortController.abort();
+
+    // Spin until we're no longer running.
+    while (this.running) {
+      await sleep(1000);
+    }
+  }
+}
+
+let runner: BatchedMigrationsRunner | null = null;
+
+function assertRunner(
+  runner: BatchedMigrationsRunner | null
+): asserts runner is BatchedMigrationsRunner {
+  if (!runner) throw new Error('Batched migrations not initialized');
+}
+
+export function initBatchedMigrations(options: BatchedMigrationRunnerOptions) {
+  if (runner) throw new Error('Batched migrations already initialized');
+  runner = new BatchedMigrationsRunner(options);
+  return runner;
+}
+
+export function startBatchedMigrations(options: BatchedMigrationStartOptions = {}) {
+  assertRunner(runner);
+  runner.start(options);
+  return runner;
+}
+
+export async function stopBatchedMigrations() {
+  assertRunner(runner);
+  await runner.stop();
+  runner = null;
+}
+
+/**
+ * Given a batched migration identifier like `20230406184103_migration`,
+ * enqueues it for execution by creating a row in the `batched_migrations`
+ * table.
+ *
+ * Despite taking a full identifier, only the timestamp is used to uniquely
+ * identify the batched migration. The remaining part is just used to make
+ * calls more human-readable.
+ *
+ * @param identifier The identifier of the batched migration to enqueue.
+ */
+export async function enqueueBatchedMigration(identifier: string) {
+  assertRunner(runner);
+  await runner.enqueueBatchedMigration(identifier);
+}
+
+/**
+ * Given a batched migration identifier like `20230406184103_migration`,
+ * synchronously runs it to completion. An error will be thrown if the final
+ * status of the migration is not `succeeded`.
+ *
+ * @param identifier The identifier of the batched migration to finalize.
+ * @param options Options for finalizing the batched migration.
+ */
+export async function finalizeBatchedMigration(
+  identifier: string,
+  options?: BatchedMigrationFinalizeOptions
+) {
+  assertRunner(runner);
+  await runner.finalizeBatchedMigration(identifier, options);
+}

package/src/batched-migrations/fixtures/20230406184103_successful_migration.ts

@@ -0,0 +1,13 @@
+import { makeBatchedMigration } from '../batched-migration';
+
+export default makeBatchedMigration({
+  async getParameters() {
+    return {
+      min: 1n,
+      max: 100n,
+      batchSize: 10,
+    };
+  },
+
+  async execute(_min: bigint, _max: bigint) {},
+});

package/src/batched-migrations/fixtures/20230406184107_failing_migration.js

@@ -0,0 +1,16 @@
+// @ts-check
+const { makeBatchedMigration } = require('../batched-migration');
+
+module.exports = makeBatchedMigration({
+  async getParameters() {
+    return {
+      min: 2n,
+      max: 200n,
+      batchSize: 20,
+    };
+  },
+
+  async execute(_min, _max) {
+    throw new Error('Testing failure');
+  },
+});

package/src/batched-migrations/fixtures/20230407230446_no_rows_migration.ts

@@ -0,0 +1,15 @@
+import { makeBatchedMigration } from '../batched-migration';
+
+export default makeBatchedMigration({
+  async getParameters() {
+    return {
+      // Simulates the case where there are no rows to process. A null
+      // max value is what we would get for some query like
+      // `SELECT MAX(id) FROM table;`.
+      max: null,
+      batchSize: 10,
+    };
+  },
+
+  async execute(_min: bigint, _max: bigint) {},
+});

package/src/batched-migrations/index.ts

@@ -0,0 +1,21 @@
+export {
+  BatchedMigrationRow,
+  BatchedMigrationStatus,
+  makeBatchedMigration,
+  selectAllBatchedMigrations,
+  selectBatchedMigration,
+  selectBatchedMigrationForTimestamp,
+  retryFailedBatchedMigrationJobs,
+} from './batched-migration';
+export {
+  BatchedMigrationJobRow,
+  BatchedMigrationJobStatus,
+  selectRecentJobsWithStatus,
+} from './batched-migration-job';
+export {
+  initBatchedMigrations,
+  startBatchedMigrations,
+  stopBatchedMigrations,
+  enqueueBatchedMigration,
+  finalizeBatchedMigration,
+} from './batched-migrations-runner';

package/src/index.ts

@@ -1,203 +1,22 @@
-import fs from 'fs-extra';
 import path from 'path';
 
-import * as namedLocks from '@prairielearn/named-locks';
-import { logger } from '@prairielearn/logger';
-import * as sqldb from '@prairielearn/postgres';
-import * as error from '@prairielearn/error';
-
-const sql = sqldb.loadSqlEquiv(__filename);
-
-export async function init(migrationDir: string, project: string) {
-  const lockName = 'migrations';
-  logger.verbose(`Waiting for lock ${lockName}`);
-  await namedLocks.doWithLock(
-    lockName,
-    {
-      // Migrations *might* take a long time to run, so we'll enable automatic
-      // lock renewal so that our lock doesn't get killed by the Postgres
-      // idle session timeout.
-      //
-      // That said, we should generally try to keep migrations executing as
-      // quickly as possible. A long-running migration likely means that
-      // Postgres is locking a whole table, which is unacceptable in production.
-      autoRenew: true,
-    },
-    async () => {
-      logger.verbose(`Acquired lock ${lockName}`);
-      await initWithLock(migrationDir, project);
-    }
-  );
-  logger.verbose(`Released lock ${lockName}`);
-}
-
-/**
- * Timestamp prefixes will be of the form `YYYYMMDDHHMMSS`, which will have 14 digits.
- * If this code is still around in the year 10000... good luck.
- */
-const MIGRATION_FILENAME_REGEX = /^([0-9]{14})_.+\.[a-z]+$/;
-
-interface MigrationFile {
-  filename: string;
-  timestamp: string;
-}
-
-export async function readAndValidateMigrationsFromDirectory(
-  dir: string,
-  extensions: string[]
-): Promise<MigrationFile[]> {
-  const migrationFiles = (await fs.readdir(dir)).filter((m) =>
-    extensions.some((e) => m.endsWith(e))
-  );
-
-  const migrations = migrationFiles.map((mf) => {
-    const match = mf.match(MIGRATION_FILENAME_REGEX);
-
-    if (!match) {
-      throw new Error(`Invalid migration filename: ${mf}`);
-    }
-
-    const timestamp = match[1] ?? null;
-
-    if (timestamp === null) {
-      throw new Error(`Migration ${mf} does not have a timestamp`);
-    }
-
-    return {
-      filename: mf,
-      timestamp,
-    };
-  });
-
-  // First pass: validate that all migrations have a unique timestamp prefix.
-  // This will avoid data loss and conflicts in unexpected scenarios.
-  const seenTimestamps = new Set();
-  for (const migration of migrations) {
-    const { filename, timestamp } = migration;
-
-    if (timestamp !== null) {
-      if (seenTimestamps.has(timestamp)) {
-        throw new Error(`Duplicate migration timestamp: ${timestamp} (${filename})`);
-      }
-      seenTimestamps.add(timestamp);
-    }
-  }
-
-  return migrations;
-}
-
-export function sortMigrationFiles(migrationFiles: MigrationFile[]): MigrationFile[] {
-  return migrationFiles.sort((a, b) => {
-    return a.timestamp.localeCompare(b.timestamp);
-  });
-}
-
-export function getMigrationsToExecute(
-  migrationFiles: MigrationFile[],
-  executedMigrations: { timestamp: string | null }[]
-): MigrationFile[] {
-  // If no migrations have ever been run, run them all.
-  if (executedMigrations.length === 0) {
-    return migrationFiles;
-  }
-
-  const executedMigrationTimestamps = new Set(executedMigrations.map((m) => m.timestamp));
-  return migrationFiles.filter((m) => !executedMigrationTimestamps.has(m.timestamp));
-}
-
-export async function initWithLock(migrationDir: string, project: string) {
-  logger.verbose('Starting DB schema migration');
-
-  // Create the migrations table if needed
-  await sqldb.queryAsync(sql.create_migrations_table, {});
-
-  // Apply necessary changes to the migrations table as needed.
-  try {
-    await sqldb.queryAsync('SELECT project FROM migrations;', {});
-  } catch (err: any) {
-    if (err.routine === 'errorMissingColumn') {
-      logger.info('Altering migrations table');
-      await sqldb.queryAsync(sql.add_projects_column, {});
-    } else {
-      throw err;
-    }
-  }
-  try {
-    await sqldb.queryAsync('SELECT timestamp FROM migrations;', {});
-  } catch (err: any) {
-    if (err.routine === 'errorMissingColumn') {
-      logger.info('Altering migrations table again');
-      await sqldb.queryAsync(sql.add_timestamp_column, {});
-    } else {
-      throw err;
-    }
-  }
-
-  let allMigrations = await sqldb.queryAsync(sql.get_migrations, { project });
-
-  const migrationFiles = await readAndValidateMigrationsFromDirectory(migrationDir, [
-    '.sql',
-    '.js',
-    '.ts',
-    '.mjs',
-  ]);
-
-  // Validation: if we not all previously-executed migrations have timestamps,
-  // prompt the user to deploy an earlier version that includes both indexes
-  // and timestamps.
-  const migrationsMissingTimestamps = allMigrations.rows.filter((m) => !m.timestamp);
-  if (migrationsMissingTimestamps.length > 0) {
-    throw new Error(
-      [
-        'The following migrations are missing timestamps:',
-        migrationsMissingTimestamps.map((m) => `  ${m.filename}`),
-        // This revision was the most recent commit to `master` before the
-        // code handling indexes was removed.
-        'You must deploy revision 1aa43c7348fa24cf636413d720d06a2fa9e38ef2 first.',
-      ].join('\n')
-    );
-  }
-
-  // Refetch the list of migrations from the database.
-  allMigrations = await sqldb.queryAsync(sql.get_migrations, { project });
-
-  // Sort the migration files into execution order.
-  const sortedMigrationFiles = sortMigrationFiles(migrationFiles);
-
-  // Figure out which migrations have to be applied.
-  const migrationsToExecute = getMigrationsToExecute(sortedMigrationFiles, allMigrations.rows);
-
-  for (const { filename, timestamp } of migrationsToExecute) {
-    if (allMigrations.rows.length === 0) {
-      // if we are running all the migrations then log at a lower level
-      logger.verbose(`Running migration ${filename}`);
-    } else {
-      logger.info(`Running migration ${filename}`);
-    }
-
-    const migrationPath = path.join(migrationDir, filename);
-    if (filename.endsWith('.sql')) {
-      const migrationSql = await fs.readFile(migrationPath, 'utf8');
-      try {
-        await sqldb.queryAsync(migrationSql, {});
-      } catch (err) {
-        error.addData(err, { sqlFile: filename });
-        throw err;
-      }
-    } else {
-      const migrationModule = await import(migrationPath);
-      const implementation = migrationModule.default;
-      if (typeof implementation !== 'function') {
-        throw new Error(`Migration ${filename} does not export a default function`);
-      }
-      await implementation();
-    }
-
-    // Record the migration.
-    await sqldb.queryAsync(sql.insert_migration, {
-      filename: filename,
-      timestamp,
-      project,
-    });
-  }
-}
+export { init } from './migrations';
+
+export {
+  BatchedMigrationRow,
+  BatchedMigrationStatus,
+  BatchedMigrationJobRow,
+  BatchedMigrationJobStatus,
+  makeBatchedMigration,
+  initBatchedMigrations,
+  startBatchedMigrations,
+  stopBatchedMigrations,
+  enqueueBatchedMigration,
+  finalizeBatchedMigration,
+  selectAllBatchedMigrations,
+  selectBatchedMigration,
+  selectBatchedMigrationForTimestamp,
+  selectRecentJobsWithStatus,
+} from './batched-migrations';
+
+export const SCHEMA_MIGRATIONS_PATH = path.resolve(__dirname, '..', 'schema-migrations');