fix(rivetkit): run db migrations before lifecycle hooks

Author: NathanFlurry

rivetkit-typescript/packages/rivetkit-napi/src/napi_actor_events.rs

@@ -198,6 +198,17 @@ async fn run_preamble(
 ) -> Result<RunHandlerSlot> {
 	let is_new = snapshot.is_none();
 
+	// Run database migrations before any user lifecycle hook so `c.db` is
+	// usable from createState, onCreate, and createVars.
+	if let Some(callback) = &bindings.on_migrate {
+		with_timeout(
+			"onMigrate",
+			config.on_migrate_timeout,
+			call_on_migrate(callback, ctx, is_new),
+		)
+		.await?;
+	}
+
 	if is_new {
 		if let Some(callback) = &bindings.create_state {
 			let bytes = with_timeout(
@@ -240,15 +251,6 @@ async fn run_preamble(
 		.await?;
 	}
 
-	if let Some(callback) = &bindings.on_migrate {
-		with_timeout(
-			"onMigrate",
-			config.on_migrate_timeout,
-			call_on_migrate(callback, ctx, is_new),
-		)
-		.await?;
-	}
-
 	ctx.init_alarms().await?;
 	ctx.mark_ready_internal();
 

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/db-init-order.ts

@@ -0,0 +1,92 @@
+import { actor } from "rivetkit";
+import { db } from "@/common/database/mod";
+
+// Verifies that onMigrate runs before createState/onCreate/createVars so the
+// schema is queryable from those lifecycle hooks. The runtime should make
+// `c.db` usable as soon as user code can read it.
+
+export const dbInitOrderCreateStateActor = actor({
+	createState: async (c, _input) => {
+		const rows = await c.db.execute<{ count: number }>(
+			"SELECT COUNT(*) as count FROM init_order_items",
+		);
+		return { count: rows[0]?.count ?? -1 };
+	},
+	db: db({
+		onMigrate: async (db) => {
+			await db.execute(`
+				CREATE TABLE IF NOT EXISTS init_order_items (
+					id INTEGER PRIMARY KEY AUTOINCREMENT,
+					name TEXT NOT NULL
+				)
+			`);
+		},
+	}),
+	actions: {
+		getInitialCount: (c) => (c.state as { count: number }).count,
+		insert: async (c, name: string) => {
+			await c.db.execute(
+				"INSERT INTO init_order_items (name) VALUES (?)",
+				name,
+			);
+		},
+	},
+	options: {
+		actionTimeout: 120_000,
+		sleepTimeout: 100,
+	},
+});
+
+export const dbInitOrderOnCreateActor = actor({
+	state: { initialCount: -1 },
+	onCreate: async (c, _input) => {
+		const rows = await c.db.execute<{ count: number }>(
+			"SELECT COUNT(*) as count FROM init_order_items",
+		);
+		c.state.initialCount = rows[0]?.count ?? -1;
+	},
+	db: db({
+		onMigrate: async (db) => {
+			await db.execute(`
+				CREATE TABLE IF NOT EXISTS init_order_items (
+					id INTEGER PRIMARY KEY AUTOINCREMENT,
+					name TEXT NOT NULL
+				)
+			`);
+		},
+	}),
+	actions: {
+		getInitialCount: (c) => c.state.initialCount,
+	},
+	options: {
+		actionTimeout: 120_000,
+		sleepTimeout: 100,
+	},
+});
+
+export const dbInitOrderCreateVarsActor = actor({
+	state: {},
+	createVars: async (c) => {
+		const rows = await c.db.execute<{ count: number }>(
+			"SELECT COUNT(*) as count FROM init_order_items",
+		);
+		return { initialCount: rows[0]?.count ?? -1 };
+	},
+	db: db({
+		onMigrate: async (db) => {
+			await db.execute(`
+				CREATE TABLE IF NOT EXISTS init_order_items (
+					id INTEGER PRIMARY KEY AUTOINCREMENT,
+					name TEXT NOT NULL
+				)
+			`);
+		},
+	}),
+	actions: {
+		getInitialCount: (c) => (c.vars as { initialCount: number }).initialCount,
+	},
+	options: {
+		actionTimeout: 120_000,
+		sleepTimeout: 100,
+	},
+});

rivetkit-typescript/packages/rivetkit/fixtures/driver-test-suite/registry-static.ts

@@ -22,6 +22,11 @@ import {
 import { dbActorRaw } from "./actor-db-raw";
 import { onStateChangeActor } from "./actor-onstatechange";
 import { connErrorSerializationActor } from "./conn-error-serialization";
+import {
+	dbInitOrderCreateStateActor,
+	dbInitOrderCreateVarsActor,
+	dbInitOrderOnCreateActor,
+} from "./db-init-order";
 import { dbPragmaMigrationActor } from "./db-pragma-migration";
 import { counterWithParams } from "./conn-params";
 import { connStateActor } from "./conn-state";
@@ -315,6 +320,10 @@ export const registry = setup({
 		connErrorSerializationActor,
 		// From db-pragma-migration.ts
 		dbPragmaMigrationActor,
+		// From db-init-order.ts
+		dbInitOrderCreateStateActor,
+		dbInitOrderOnCreateActor,
+		dbInitOrderCreateVarsActor,
 		// From state-zod-coercion.ts
 		stateZodCoercionActor,
 		// From lifecycle-hooks.ts

rivetkit-typescript/packages/rivetkit/tests/driver/actor-db-init-order.test.ts

@@ -0,0 +1,62 @@
+import { describe, expect, test } from "vitest";
+import { describeDriverMatrix } from "./shared-matrix";
+import { setupDriverTest } from "./shared-utils";
+
+const REAL_TIMER_DB_TIMEOUT_MS = 180_000;
+
+describeDriverMatrix("Actor Db Init Order", (driverTestConfig) => {
+	const dbTestTimeout = driverTestConfig.useRealTimers
+		? REAL_TIMER_DB_TIMEOUT_MS
+		: undefined;
+
+	describe("onMigrate runs before lifecycle hooks", () => {
+		test(
+			"createState can read schema created by onMigrate",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const key = `db-init-order-create-state-${crypto.randomUUID()}`;
+				const actor = client.dbInitOrderCreateStateActor.getOrCreate([
+					key,
+				]);
+
+				const initialCount = await actor.getInitialCount();
+				expect(initialCount).toBe(0);
+
+				await actor.insert("alpha");
+				const nextActor = client.dbInitOrderCreateStateActor.getOrCreate(
+					[key],
+				);
+				expect(typeof (await nextActor.getInitialCount())).toBe("number");
+			},
+			dbTestTimeout,
+		);
+
+		test(
+			"onCreate can read schema created by onMigrate",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const key = `db-init-order-on-create-${crypto.randomUUID()}`;
+				const actor = client.dbInitOrderOnCreateActor.getOrCreate([key]);
+
+				const initialCount = await actor.getInitialCount();
+				expect(initialCount).toBe(0);
+			},
+			dbTestTimeout,
+		);
+
+		test(
+			"createVars can read schema created by onMigrate",
+			async (c) => {
+				const { client } = await setupDriverTest(c, driverTestConfig);
+				const key = `db-init-order-create-vars-${crypto.randomUUID()}`;
+				const actor = client.dbInitOrderCreateVarsActor.getOrCreate([
+					key,
+				]);
+
+				const initialCount = await actor.getInitialCount();
+				expect(initialCount).toBe(0);
+			},
+			dbTestTimeout,
+		);
+	});
+});