Merge pull request #694 from TNG-dev/zkldi/issue-692

feat: Migration Engine

server/src/external/mongo/db.ts

@@ -40,7 +40,7 @@ import {
 	UserSettings,
 } from "tachi-common";
 import { GetMillisecondsSince } from "utils/misc";
-import { PrivateUserInfoDocument } from "utils/types";
+import { MigrationDocument, PrivateUserInfoDocument } from "utils/types";
 const logger = CreateLogCtx(__filename);
 
 let dbName = ServerConfig.MONGO_DATABASE_NAME;
@@ -192,6 +192,7 @@ const db = {
 		monkDB.get<{ userID: integer; code: string; email: string }>("verify-email-codes"),
 	"recent-folder-views": monkDB.get<RecentlyViewedFolderDocument>("recent-folder-views"),
 	"milestone-sets": monkDB.get<MilestoneSetDocument>("milestone-sets"),
+	migrations: monkDB.get<MigrationDocument>("migrations"),
 };
 
 export type StaticDatabases =
@@ -229,7 +230,8 @@ export type StaticDatabases =
 	| "verify-email-codes"
 	| "class-achievements"
 	| "recent-folder-views"
-	| "milestone-sets";
+	| "milestone-sets"
+	| "migrations";
 
 export type Databases = StaticDatabases | `songs-${Game}` | `charts-${Game}`;
 

server/src/external/mongo/indexes.ts

@@ -118,6 +118,7 @@ const staticIndexes: Partial<Record<Databases, Index[]>> = {
 	"import-locks": [index({ userID: 1 }, UNIQUE)],
 	"arc-saved-profiles": [index({ userID: 1, forImportType: 1 }, UNIQUE)],
 	"score-blacklist": [index({ scoreID: 1 }, UNIQUE)],
+	migrations: [index({ migrationID: 1 }, UNIQUE)],
 };
 
 const indexes: Partial<Record<Databases, Index[]>> = staticIndexes;

server/src/external/mongo/schemas.ts

@@ -725,4 +725,18 @@ export const DatabaseSchemas: Record<
 		playtype: isValidPlaytype,
 		milestones: ["string"],
 	}),
+	migrations: prSchemaify({
+		migrationID: "string",
+		status: p.isIn("applied", "pending"),
+		appliedOn: (self, parent) => {
+			if (parent.status === "pending") {
+				return (
+					self === undefined ||
+					`Expected appliedOn to be undefined since status was pending.`
+				);
+			}
+
+			return p.isPositiveInteger(self);
+		},
+	}),
 };

server/src/lib/migration/migrations.test.ts

@@ -0,0 +1,125 @@
+import db from "external/mongo/db";
+import t from "tap";
+import ResetDBState from "test-utils/resets";
+import { FAKE_MIGRATION } from "test-utils/test-data";
+import { ApplyMigration, FindUnappliedMigrations } from "./migrations";
+
+t.test("#FindUnappliedMigrations", (t) => {
+	t.beforeEach(ResetDBState);
+
+	t.test("Should find all unapplied migrations", async (t) => {
+		const res = await FindUnappliedMigrations();
+
+		t.equal(
+			res.includes("fake-migration"),
+			true,
+			"Should contain fake-migration since it isn't applied."
+		);
+
+		t.end();
+	});
+
+	t.test("Shouldn't return pending migrations", async (t) => {
+		await db.migrations.insert({
+			migrationID: "fake-migration",
+			status: "pending",
+		});
+
+		const res = await FindUnappliedMigrations();
+
+		t.equal(res.includes("fake-migration"), false);
+
+		t.end();
+	});
+
+	t.test("Shouldn't return applied migrations", async (t) => {
+		await db.migrations.insert({
+			migrationID: "fake-migration",
+			status: "applied",
+			appliedOn: 1000,
+		});
+
+		const res = await FindUnappliedMigrations();
+
+		t.equal(res.includes("fake-migration"), false);
+
+		t.end();
+	});
+
+	t.end();
+});
+
+t.test("#ApplyMigration", (t) => {
+	t.beforeEach(ResetDBState);
+
+	t.test("Should apply a migration.", async (t) => {
+		await ApplyMigration(FAKE_MIGRATION);
+
+		const dbRes = await db.users.findOne({
+			id: 1,
+		});
+
+		t.equal(
+			// @ts-expect-error Our migration is a fake one that adds __wasMigrated
+			// to all user documents. There's nothing special about this property name.
+			dbRes?.__wasMigrated,
+			true,
+			"Migration should be applied on the users."
+		);
+
+		t.end();
+	});
+
+	t.test("Should panic and reset if something goes wrong.", async (t) => {
+		try {
+			await ApplyMigration({
+				id: "bad-migration",
+				up: () => {
+					throw new Error("Failed to migrate.");
+				},
+				down: () =>
+					db.users.update({}, { $set: { __wasUnmigrated: true } }, { multi: true }),
+			});
+		} catch (e) {
+			const err = e as Error;
+			t.equal(err?.message, "Was going to exit with statusCode 1, but we're in testing.");
+		}
+
+		const dbRes = await db.users.findOne({
+			id: 1,
+		});
+
+		t.equal(
+			// @ts-expect-error Our un-migration adds this property to check if it worked.
+			dbRes?.__wasUnmigrated,
+			true,
+			"Migration should be applied on the users."
+		);
+
+		t.end();
+	});
+
+	t.test(
+		"Should panic and panic again if un-migration failed, with a scarier log message.",
+		async (t) => {
+			try {
+				await ApplyMigration({
+					id: "bad-migration",
+					up: () => {
+						throw new Error("Failed to migrate.");
+					},
+					down: () => {
+						throw new Error("Failed to undo migration.");
+					},
+				});
+			} catch (e) {
+				const err = e as Error;
+				t.equal(err?.message, "Was going to exit with statusCode 1, but we're in testing.");
+			}
+
+			t.end();
+		}
+	);
+
+	t.end();
+});

server/src/lib/migration/migrations.ts

@@ -0,0 +1,191 @@
+import db from "external/mongo/db";
+import CreateLogCtx from "lib/logger/logger";
+import { Environment } from "lib/setup/config";
+import { FAKE_MIGRATION } from "test-utils/test-data";
+import { Migration } from "utils/types";
+import UGPTRivalsMigration from "./migrations/add-rivals-to-ugpt";
+
+const logger = CreateLogCtx(__filename);
+
+// Migrations are stored in an array because they have some concept of order
+// That is, migrations should ideally be applied in a fixed order just to avoid
+// any potential unsavoury interactions.
+const REGISTERED_MIGRATIONS: Migration[] =
+	// If we're testing, we should pull fake migrations instead to ensure the tests
+	// stay consistent
+	Environment.nodeEnv === "test" ? [FAKE_MIGRATION] : [UGPTRivalsMigration];
+
+function CreateMigrationLookupMap(migrations: Migration[]) {
+	const map = new Map<string, Migration>();
+
+	for (const mig of migrations) {
+		if (map.get(mig.id)) {
+			logger.crit(
+				`Multiple migrations are registered for ${mig.id}. Cannot safely apply migrations.`
+			);
+
+			// note, we want to exit in testing here, this is fine.
+			process.exit(1);
+		}
+
+		map.set(mig.id, mig);
+	}
+
+	return map;
+}
+
+// That said, we still make a lookup map so we can efficiently find registered
+// migrations.
+const MIGRATION_LOOKUP = CreateMigrationLookupMap(REGISTERED_MIGRATIONS);
+
+/**
+ * Return the migrationID of all migrations that have not been applied.
+ */
+export async function FindUnappliedMigrations() {
+	const migrations = await db.migrations.find(
+		{},
+		{
+			projection: {
+				migrationID: 1,
+			},
+		}
+	);
+
+	const appliedIDs = new Set(migrations.map((e) => e.migrationID));
+
+	const notApplied = [];
+
+	for (const migration of REGISTERED_MIGRATIONS) {
+		if (!appliedIDs.has(migration.id)) {
+			notApplied.push(migration.id);
+		}
+	}
+
+	return notApplied;
+}
+
+export async function ApplyUnappliedMigrations() {
+	const unapplied = await FindUnappliedMigrations();
+
+	if (unapplied.length === 0) {
+		logger.verbose(`Found no migrations to apply. Not bothering.`);
+		return;
+	}
+
+	logger.info(
+		`Found ${
+			unapplied.length
+		} migration(s) that have not been applied yet. These are '${unapplied.join("', '")}'.`,
+		{ unapplied }
+	);
+
+	for (const migrationID of unapplied) {
+		// These things need to apply in lockstep.
+		// Note that if any migration fails, this will exit at CRIT level
+		// and not continue.
+		// eslint-disable-next-line no-await-in-loop
+		await ApplyMigrationByID(migrationID);
+	}
+}
+
+export function ApplyMigrationByID(migrationID: string) {
+	const migration = MIGRATION_LOOKUP.get(migrationID);
+
+	if (!migration) {
+		logger.error(
+			`Attempted to apply migration ${migrationID}, but that migration doesn't exist?`
+		);
+		throw new Error(
+			`Attempted to apply migration ${migrationID}, but that migration doesn't exist?`
+		);
+	}
+
+	return ApplyMigration(migration);
+}
+
+/**
+ * Applies a migration. Will **EXIT TACHI** if a migration fails to apply, as this implies
+ * database level inconsistencies which could cause severe problems.
+ */
+export async function ApplyMigration(migration: Migration) {
+	const migrationID = migration.id;
+
+	logger.info(`Recieved request to apply migration '${migrationID}'.`);
+
+	// Lock the migration here:
+	// If the migration is already applied (i.e. migrationID is set and pending/applied)
+	// don't bother applying again
+	// If the migration isn't, set it as pending and continue. This is race-condition
+	// safe.
+	const wasApplied = await db.migrations.findOneAndUpdate(
+		{ migrationID },
+		{
+			$setOnInsert: {
+				migrationID,
+				status: "pending",
+			},
+		},
+		{
+			returnOriginal: true,
+			upsert: true,
+		}
+	);
+
+	if (wasApplied) {
+		logger.error(`Tried to apply migration ${migrationID}, but it was already applied.`);
+		throw new Error(`Tried to apply migration ${migrationID}, but it was already applied.`);
+	}
+
+	try {
+		await migration.up();
+
+		logger.info(`Successfully applied migration '${migrationID}'.`);
+
+		await db.migrations.update(
+			{
+				migrationID,
+			},
+			{
+				$set: {
+					status: "applied",
+					appliedOn: Date.now(),
+				},
+			}
+		);
+
+		logger.info(`Successfully set migration state for '${migrationID}' to applied.`);
+	} catch (err) {
+		logger.severe(
+			`Failed to apply migration '${migrationID}'. Attempting to revert migration. Tachi will exit after attempting to revert this migration.`,
+			{ err }
+		);
+
+		try {
+			await migration.down();
+
+			logger.info(
+				`Successfully reverted migration. Database state is okay, but Tachi will not boot.`
+			);
+
+			logger.crit(
+				`Failed to apply migration '${migrationID}'. Although successfully reverted, Tachi will not boot if a migration fails.`
+			);
+		} catch (err) {
+			logger.crit(
+				`FAILED TO REVERT PARTIAL MIGRATION '${migrationID}'. THE DATABASE IS ALMOST CERTAINLY TAINTED. MANUAL INTERVENTION/BACKUP RESTORE NECESSARY.`,
+				{
+					err,
+				}
+			);
+		}
+
+		// remove the stale migration so it can be re-ran in the future.
+		await db.migrations.findOneAndDelete({ migrationID });
+
+		if (Environment.nodeEnv === "test") {
+			throw new Error("Was going to exit with statusCode 1, but we're in testing.");
+		}
+
+		process.exit(1);
+	}
+}

server/src/lib/migration/migrations/add-rivals-to-ugpt.ts

@@ -0,0 +1,28 @@
+import db from "external/mongo/db";
+import { Migration } from "utils/types";
+
+const migration: Migration = {
+	id: "add-rivals-to-ugpt",
+	up: async () => {
+		await db["game-settings"].update(
+			{},
+			{
+				$set: {
+					rivals: [],
+				},
+			},
+			{ multi: true }
+		);
+	},
+	down: async () => {
+		await db["game-settings"].update(
+			{},
+			{
+				$unset: { rivals: 1 },
+			},
+			{ multi: true }
+		);
+	},
+};
+
+export default migration;

server/src/main.ts

@@ -13,6 +13,7 @@ import { LoadDefaultClients } from "lib/builtin-clients/builtin-clients";
 import { VERSION_PRETTY } from "lib/constants/version";
 import { HandleSIGTERMGracefully } from "lib/handlers/sigterm";
 import CreateLogCtx from "lib/logger/logger";
+import { ApplyUnappliedMigrations } from "lib/migration/migrations";
 import { Environment, ServerConfig, TachiConfig } from "lib/setup/config";
 import path from "path";
 import server from "server/server";
@@ -32,6 +33,8 @@ async function RunOnInit() {
 	await InitSequenceDocs();
 	await SetIndexesIfNoneSet();
 
+	await ApplyUnappliedMigrations();
+
 	await db["folder-chart-lookup"].findOne().then((r) => {
 		// If there are no folder chart lookups, initialise them.
 		if (!r) {

server/src/scripts/migrations/migrate-ugpt-rivals.ts

@@ -1,15 +0,0 @@
-import db from "external/mongo/db";
-
-if (require.main === module) {
-	(async () => {
-		// give everyone empty rivals.
-		await db["game-settings"].update(
-			{},
-			{
-				$set: {
-					rivals: [],
-				},
-			}
-		);
-	})();
-}

server/src/test-utils/mock-db/migrations.json

@@ -0,0 +1 @@
+[]

server/src/test-utils/test-data.ts

@@ -28,6 +28,7 @@ import {
 	UGPTSettings,
 } from "tachi-common";
 import { ApplyNTimes, RFA } from "utils/misc";
+import { Migration } from "utils/types";
 
 const file = (name: string) => path.join(__dirname, "/test-data", name);
 
@@ -637,3 +638,9 @@ export const FakeGameSettings: UGPTSettings = {
 	rivals: [],
 	userID: 1,
 };
+
+export const FAKE_MIGRATION: Migration = {
+	id: "fake-migration",
+	up: () => db.users.update({}, { $set: { __wasMigrated: true } }, { multi: true }),
+	down: () => db.users.update({}, { $unset: { __wasMigrated: 1 } }, { multi: true }),
+};

server/src/utils/misc.ts

@@ -2,7 +2,7 @@ import { exec } from "child_process";
 import crypto from "crypto";
 import { ONE_HOUR } from "lib/constants/time";
 import { TachiConfig } from "lib/setup/config";
-import { Game, GamePTConfig, GetGameConfig, integer, Playtype, Playtypes } from "tachi-common";
+import { Game, GamePTConfig, GetGameConfig, integer, Playtypes } from "tachi-common";
 import { URL } from "url";
 
 // https://github.com/sindresorhus/escape-string-regexp/blob/main/index.js

server/src/utils/types.ts

@@ -84,3 +84,21 @@ export interface PrivateUserInfoDocument {
 	password: string;
 	email: string;
 }
+
+export interface Migration {
+	id: string;
+	up: () => Promise<unknown>;
+	down: () => Promise<unknown>;
+}
+
+export type MigrationDocument = {
+	migrationID: string;
+} & (
+	| {
+			status: "applied";
+			appliedOn: integer;
+	  }
+	| {
+			status: "pending";
+	  }
+);