@prairielearn/migrations 1.1.0 → 1.2.1

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

@@ -0,0 +1,215 @@
+import {
+  loadSqlEquiv,
+  queryAsync,
+  queryValidatedOneRow,
+  queryValidatedSingleColumnOneRow,
+  queryValidatedZeroOrOneRow,
+} from '@prairielearn/postgres';
+import { logger } from '@prairielearn/logger';
+import { serializeError } from 'serialize-error';
+import { z } from 'zod';
+
+import {
+  BatchedMigrationStatus,
+  BatchedMigrationRow,
+  updateBatchedMigrationStatus,
+  BatchedMigrationStatusSchema,
+  BatchedMigrationImplementation,
+} from './batched-migration';
+import {
+  BatchedMigrationJobRowSchema,
+  BatchedMigrationJobStatus,
+  BatchedMigrationJobRow,
+} from './batched-migration-job';
+
+const sql = loadSqlEquiv(__filename);
+
+interface BatchedMigrationRunnerOptions {
+  logProgress?: boolean;
+}
+
+export class BatchedMigrationRunner {
+  private options: BatchedMigrationRunnerOptions;
+  private migration: BatchedMigrationRow;
+  private migrationImplementation: BatchedMigrationImplementation;
+  private migrationStatus: BatchedMigrationStatus;
+
+  constructor(
+    migration: BatchedMigrationRow,
+    migrationImplementation: BatchedMigrationImplementation,
+    options: BatchedMigrationRunnerOptions = {}
+  ) {
+    this.options = options;
+    this.migration = migration;
+    this.migrationImplementation = migrationImplementation;
+    this.migrationStatus = migration.status;
+  }
+
+  private log(message: string, ...meta: any[]) {
+    if (this.options.logProgress) {
+      logger.info(`[${this.migration.filename}] ${message}`, ...meta);
+    }
+  }
+
+  private async hasIncompleteJobs(migration: BatchedMigrationRow): Promise<boolean> {
+    return queryValidatedSingleColumnOneRow(
+      sql.batched_migration_has_incomplete_jobs,
+      { batched_migration_id: migration.id },
+      z.boolean()
+    );
+  }
+
+  private async hasFailedJobs(migration: BatchedMigrationRow): Promise<boolean> {
+    return queryValidatedSingleColumnOneRow(
+      sql.batched_migration_has_failed_jobs,
+      { batched_migration_id: migration.id },
+      z.boolean()
+    );
+  }
+
+  private async refreshMigrationStatus(migration: BatchedMigrationRow) {
+    this.migrationStatus = await queryValidatedSingleColumnOneRow(
+      sql.get_migration_status,
+      {
+        id: migration.id,
+      },
+      BatchedMigrationStatusSchema
+    );
+  }
+
+  private async finishRunningMigration(migration: BatchedMigrationRow) {
+    // Safety check: if there are any pending jobs, don't mark this
+    // migration as finished.
+    if (await this.hasIncompleteJobs(migration)) {
+      this.log(`Incomplete jobs found, not marking as finished`);
+      return;
+    }
+
+    const hasFailedJobs = await this.hasFailedJobs(migration);
+    const finalStatus = hasFailedJobs ? 'failed' : 'succeeded';
+    await updateBatchedMigrationStatus(migration.id, finalStatus);
+    this.log(`Finished with status '${finalStatus}'`);
+  }
+
+  private async getNextBatchBounds(
+    migration: BatchedMigrationRow
+  ): Promise<null | [bigint, bigint]> {
+    const lastJob = await queryValidatedZeroOrOneRow(
+      sql.select_last_batched_migration_job,
+      { batched_migration_id: migration.id },
+      BatchedMigrationJobRowSchema
+    );
+
+    const nextMin = lastJob ? lastJob.max_value + 1n : migration.min_value;
+    if (nextMin > migration.max_value) return null;
+
+    let nextMax = nextMin + BigInt(migration.batch_size) - 1n;
+    if (nextMax > migration.max_value) nextMax = migration.max_value;
+
+    return [nextMin, nextMax];
+  }
+
+  private async startJob(job: BatchedMigrationJobRow) {
+    await queryAsync(sql.start_batched_migration_job, { id: job.id });
+    const jobRange = `[${job.min_value}, ${job.max_value}]`;
+    const migrationRange = `[${this.migration.min_value}, ${this.migration.max_value}]`;
+    this.log(`Started job ${job.id} for range ${jobRange} in ${migrationRange}`);
+  }
+
+  private serializeJobData(data: unknown) {
+    if (data == null) return null;
+
+    // Return JSON-stringified data. Convert BigInts to strings.
+    return JSON.stringify(data, (_key, value) => {
+      if (typeof value === 'bigint') return value.toString();
+      return value;
+    });
+  }
+
+  private async finishJob(
+    job: BatchedMigrationJobRow,
+    status: Extract<BatchedMigrationJobStatus, 'failed' | 'succeeded'>,
+    data?: unknown
+  ) {
+    await queryAsync(sql.finish_batched_migration_job, {
+      id: job.id,
+      status,
+      data: this.serializeJobData(data),
+    });
+    this.log(`Job ${job.id} finished with status '${status}'`);
+  }
+
+  private async getOrCreateNextMigrationJob(
+    migration: BatchedMigrationRow
+  ): Promise<BatchedMigrationJobRow | null> {
+    const nextBatchBounds = await this.getNextBatchBounds(migration);
+    if (nextBatchBounds) {
+      return queryValidatedOneRow(
+        sql.insert_batched_migration_job,
+        {
+          batched_migration_id: migration.id,
+          min_value: nextBatchBounds[0],
+          max_value: nextBatchBounds[1],
+        },
+        BatchedMigrationJobRowSchema
+      );
+    } else {
+      // Pick up any old pending jobs from this migration. These will only exist if
+      // an admin manually elected to retry all failed jobs; we'll never automatically
+      // transition failed jobs back to pending.
+      return queryValidatedZeroOrOneRow(
+        sql.select_first_pending_batched_migration_job,
+        { batched_migration_id: migration.id },
+        BatchedMigrationJobRowSchema
+      );
+    }
+  }
+
+  private async runMigrationJob(
+    migration: BatchedMigrationRow,
+    migrationImplementation: BatchedMigrationImplementation
+  ) {
+    const nextJob = await this.getOrCreateNextMigrationJob(migration);
+    if (nextJob) {
+      await this.startJob(nextJob);
+
+      let error = null;
+      try {
+        // We'll only handle errors thrown by the migration itself. If any of
+        // our own execution machinery throws an error, we'll let it bubble up.
+        await migrationImplementation.execute(nextJob.min_value, nextJob.max_value);
+      } catch (err) {
+        error = err;
+      }
+
+      if (error) {
+        await this.finishJob(nextJob, 'failed', { error: serializeError(error) });
+      } else {
+        await this.finishJob(nextJob, 'succeeded');
+      }
+    } else {
+      await this.finishRunningMigration(migration);
+    }
+  }
+
+  async run({
+    signal,
+    iterations,
+    durationMs,
+  }: { signal?: AbortSignal; iterations?: number; durationMs?: number } = {}) {
+    let iterationCount = 0;
+    const endTime = durationMs ? Date.now() + durationMs : null;
+    while (
+      !signal?.aborted &&
+      (iterations ? iterationCount < iterations : true) &&
+      (endTime ? Date.now() < endTime : true) &&
+      (this.migrationStatus === 'running' || this.migrationStatus === 'finalizing')
+    ) {
+      await this.runMigrationJob(this.migration, this.migrationImplementation);
+      iterationCount += 1;
+      // Always refresh the status so we can detect if the migration was marked
+      // as paused by another process.
+      await this.refreshMigrationStatus(this.migration);
+    }
+  }
+}

package/src/batched-migrations/batched-migration.sql

@@ -0,0 +1,95 @@
+-- BLOCK insert_batched_migration
+INSERT INTO
+  batched_migrations (
+    project,
+    filename,
+    timestamp,
+    batch_size,
+    min_value,
+    max_value,
+    status,
+    started_at
+  )
+VALUES
+  (
+    $project,
+    $filename,
+    $timestamp,
+    $batch_size,
+    $min_value,
+    $max_value,
+    $status,
+    -- If the migration is marked as already having succeeded, set `started_at`
+    -- since the migration did technically start.
+    CASE
+      WHEN $status::enum_batched_migration_status = 'succeeded' THEN CURRENT_TIMESTAMP
+      ELSE NULL
+    END
+  )
+ON CONFLICT DO NOTHING
+RETURNING
+  *;
+
+-- BLOCK select_all_batched_migrations
+SELECT
+  *
+FROM
+  batched_migrations
+WHERE
+  project = $project
+ORDER BY
+  id ASC;
+
+-- BLOCK select_batched_migration
+SELECT
+  *
+FROM
+  batched_migrations
+WHERE
+  project = $project
+  AND id = $id;
+
+-- BLOCK select_batched_migration_for_timestamp
+SELECT
+  *
+FROM
+  batched_migrations
+WHERE
+  project = $project
+  AND timestamp = $timestamp;
+
+-- BLOCK update_batched_migration_status
+UPDATE batched_migrations
+SET
+  status = $status,
+  updated_at = CURRENT_TIMESTAMP
+WHERE
+  id = $id
+RETURNING
+  *;
+
+-- BLOCK retry_failed_jobs
+WITH
+  updated_batched_migration AS (
+    UPDATE batched_migrations
+    SET
+      status = 'running',
+      started_at = CURRENT_TIMESTAMP,
+      updated_at = CURRENT_TIMESTAMP
+    WHERE
+      project = $project
+      AND id = $id
+    RETURNING
+      *
+  )
+UPDATE batched_migration_jobs
+SET
+  status = 'pending',
+  started_at = NULL,
+  finished_at = NULL,
+  updated_at = CURRENT_TIMESTAMP
+FROM
+  updated_batched_migration
+WHERE
+  batched_migration_id = updated_batched_migration.id
+  AND batched_migration_jobs.status = 'failed';

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

@@ -0,0 +1,129 @@
+import {
+  loadSqlEquiv,
+  queryAsync,
+  queryValidatedOneRow,
+  queryValidatedRows,
+  queryValidatedZeroOrOneRow,
+} from '@prairielearn/postgres';
+import { z } from 'zod';
+
+const sql = loadSqlEquiv(__filename);
+
+export const BatchedMigrationStatusSchema = z.enum([
+  'pending',
+  'paused',
+  'running',
+  'finalizing',
+  'failed',
+  'succeeded',
+]);
+export type BatchedMigrationStatus = z.infer<typeof BatchedMigrationStatusSchema>;
+
+export const BatchedMigrationRowSchema = z.object({
+  id: z.string(),
+  project: z.string(),
+  filename: z.string(),
+  timestamp: z.string(),
+  batch_size: z.number(),
+  min_value: z.bigint({ coerce: true }),
+  max_value: z.bigint({ coerce: true }),
+  status: BatchedMigrationStatusSchema,
+  created_at: z.date(),
+  updated_at: z.date(),
+  started_at: z.date().nullable(),
+});
+export type BatchedMigrationRow = z.infer<typeof BatchedMigrationRowSchema>;
+
+export interface BatchedMigrationParameters {
+  min?: bigint | null;
+  max: bigint | null;
+  batchSize?: number;
+}
+
+export interface BatchedMigrationImplementation {
+  getParameters(): Promise<BatchedMigrationParameters>;
+  execute(start: bigint, end: bigint): Promise<void>;
+}
+
+/**
+ * Identity function that helps to write correct batched migrations.
+ */
+export function makeBatchedMigration<T extends BatchedMigrationImplementation>(fns: T): T {
+  validateBatchedMigrationImplementation(fns);
+  return fns;
+}
+
+export function validateBatchedMigrationImplementation(
+  fns: BatchedMigrationImplementation
+): asserts fns is BatchedMigrationImplementation {
+  if (typeof fns.getParameters !== 'function') {
+    throw new Error('getParameters() must be a function');
+  }
+  if (typeof fns.execute !== 'function') {
+    throw new Error('execute() must be a function');
+  }
+}
+
+type NewBatchedMigration = Pick<
+  BatchedMigrationRow,
+  'project' | 'filename' | 'timestamp' | 'batch_size' | 'min_value' | 'max_value' | 'status'
+>;
+
+/**
+ * Inserts a new batched migration. If one already exists for the given
+ * project/timestamp pair, returns null, otherwise returns the inserted row.
+ */
+export async function insertBatchedMigration(
+  migration: NewBatchedMigration
+): Promise<BatchedMigrationRow | null> {
+  return queryValidatedZeroOrOneRow(
+    sql.insert_batched_migration,
+    migration,
+    BatchedMigrationRowSchema
+  );
+}
+
+export async function selectAllBatchedMigrations(project: string) {
+  return queryValidatedRows(
+    sql.select_all_batched_migrations,
+    { project },
+    BatchedMigrationRowSchema
+  );
+}
+
+export async function selectBatchedMigration(
+  project: string,
+  id: string
+): Promise<BatchedMigrationRow> {
+  return queryValidatedOneRow(
+    sql.select_batched_migration,
+    { project, id },
+    BatchedMigrationRowSchema
+  );
+}
+
+export async function selectBatchedMigrationForTimestamp(
+  project: string,
+  timestamp: string
+): Promise<BatchedMigrationRow> {
+  return queryValidatedOneRow(
+    sql.select_batched_migration_for_timestamp,
+    { project, timestamp },
+    BatchedMigrationRowSchema
+  );
+}
+
+export async function updateBatchedMigrationStatus(
+  id: string,
+  status: BatchedMigrationStatus
+): Promise<BatchedMigrationRow> {
+  return queryValidatedOneRow(
+    sql.update_batched_migration_status,
+    { id, status },
+    BatchedMigrationRowSchema
+  );
+}
+
+export async function retryFailedBatchedMigrationJobs(project: string, id: string): Promise<void> {
+  await queryAsync(sql.retry_failed_jobs, { project, id });
+}

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

@@ -0,0 +1,35 @@
+-- BLOCK select_running_migration
+SELECT
+  *
+FROM
+  batched_migrations
+WHERE
+  status = 'running'
+  AND project = $project
+ORDER BY
+  id ASC
+LIMIT
+  1;
+
+-- BLOCK start_next_pending_migration
+UPDATE batched_migrations
+SET
+  status = 'running',
+  started_at = CURRENT_TIMESTAMP,
+  updated_at = CURRENT_TIMESTAMP
+WHERE
+  id = (
+    SELECT
+      id
+    FROM
+      batched_migrations
+    WHERE
+      status = 'pending'
+      AND project = $project
+    ORDER BY
+      id ASC
+    LIMIT
+      1
+  )
+RETURNING
+  *;

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

@@ -0,0 +1,111 @@
+import chai, { assert } from 'chai';
+import chaiAsPromised from 'chai-as-promised';
+import path from 'node:path';
+import { makePostgresTestUtils } from '@prairielearn/postgres';
+import * as namedLocks from '@prairielearn/named-locks';
+
+import { SCHEMA_MIGRATIONS_PATH, init } from '../index';
+import { BatchedMigrationsRunner } from './batched-migrations-runner';
+import { selectAllBatchedMigrations } from './batched-migration';
+
+chai.use(chaiAsPromised);
+
+const postgresTestUtils = makePostgresTestUtils({
+  database: 'prairielearn_migrations',
+});
+
+describe('BatchedMigrationsRunner', () => {
+  before(async () => {
+    await postgresTestUtils.createDatabase();
+    await namedLocks.init(postgresTestUtils.getPoolConfig(), (err) => {
+      throw err;
+    });
+    await init([SCHEMA_MIGRATIONS_PATH], 'prairielearn_migrations');
+  });
+
+  afterEach(async () => {
+    await postgresTestUtils.resetDatabase();
+  });
+
+  after(async () => {
+    await namedLocks.close();
+    await postgresTestUtils.dropDatabase();
+  });
+
+  it('enqueues migrations', async () => {
+    const runner = new BatchedMigrationsRunner({
+      project: 'test',
+      directories: [path.join(__dirname, 'fixtures')],
+    });
+
+    await runner.enqueueBatchedMigration('20230406184103_successful_migration');
+    await runner.enqueueBatchedMigration('20230406184107_failing_migration');
+    await runner.enqueueBatchedMigration('20230407230446_no_rows_migration');
+
+    const migrations = await selectAllBatchedMigrations('test');
+
+    assert.lengthOf(migrations, 3);
+    assert.equal(migrations[0].timestamp, '20230406184103');
+    assert.equal(migrations[0].filename, '20230406184103_successful_migration.ts');
+    assert.equal(migrations[0].status, 'pending');
+    assert.equal(migrations[1].timestamp, '20230406184107');
+    assert.equal(migrations[1].filename, '20230406184107_failing_migration.js');
+    assert.equal(migrations[1].status, 'pending');
+    assert.equal(migrations[2].timestamp, '20230407230446');
+    assert.equal(migrations[2].filename, '20230407230446_no_rows_migration.ts');
+    assert.equal(migrations[2].status, 'succeeded');
+  });
+
+  it('safely enqueues migrations multiple times', async () => {
+    const runner = new BatchedMigrationsRunner({
+      project: 'test',
+      directories: [path.join(__dirname, 'fixtures')],
+    });
+
+    await runner.enqueueBatchedMigration('20230406184103_successful_migration');
+    await runner.enqueueBatchedMigration('20230406184103_successful_migration');
+    await runner.enqueueBatchedMigration('20230406184103_successful_migration');
+
+    const migrations = await selectAllBatchedMigrations('test');
+
+    assert.lengthOf(migrations, 1);
+  });
+
+  it('finalizes a successful migration', async () => {
+    const runner = new BatchedMigrationsRunner({
+      project: 'test',
+      directories: [path.join(__dirname, 'fixtures')],
+    });
+
+    await runner.enqueueBatchedMigration('20230406184103_successful_migration');
+    await runner.finalizeBatchedMigration('20230406184103_successful_migration', {
+      logProgress: false,
+    });
+
+    const migrations = await selectAllBatchedMigrations('test');
+    assert.lengthOf(migrations, 1);
+    assert.equal(migrations[0].timestamp, '20230406184103');
+    assert.equal(migrations[0].status, 'succeeded');
+  });
+
+  it('finalizes a failing migration', async () => {
+    const runner = new BatchedMigrationsRunner({
+      project: 'test',
+      directories: [path.join(__dirname, 'fixtures')],
+    });
+
+    await runner.enqueueBatchedMigration('20230406184107_failing_migration');
+
+    await assert.isRejected(
+      runner.finalizeBatchedMigration('20230406184107_failing_migration', {
+        logProgress: false,
+      }),
+      "but it is 'failed'"
+    );
+
+    const migrations = await selectAllBatchedMigrations('test');
+    assert.lengthOf(migrations, 1);
+    assert.equal(migrations[0].timestamp, '20230406184107');
+    assert.equal(migrations[0].status, 'failed');
+  });
+});