Centralize durable schema migration ownership

crates/remora-schema/src/lib.rs

@@ -0,0 +1,148 @@
+use sqlx::{PgPool, migrate::Migrator};
+
+pub const SCHEMA_GENERATION: i64 = 1;
+
+const IDENTITY_HISTORY: &str = "_remora_identity_migrations";
+const CREDENTIALS_HISTORY: &str = "_remora_credentials_migrations";
+const ADMISSION_HISTORY: &str = "_remora_admission_migrations";
+const ATTEMPT_HISTORY: &str = "_remora_attempt_migrations";
+const RESULT_HISTORY: &str = "_remora_result_migrations";
+const DELETION_HISTORY: &str = "_remora_deletion_migrations";
+const API_HISTORY: &str = "_remora_api_migrations";
+const WORKER_HISTORY: &str = "_remora_worker_migrations";
+const CONTROL_HISTORY: &str = "_remora_schema_migrations";
+
+#[derive(Debug, thiserror::Error)]
+pub enum SchemaError {
+    #[error("schema migration failed: {0}")]
+    Migration(#[from] sqlx::migrate::MigrateError),
+    #[error("schema readiness query failed: {0}")]
+    Query(#[from] sqlx::Error),
+}
+
+pub async fn migrate(pool: &PgPool) -> Result<(), SchemaError> {
+    durable_migrator().run(pool).await?;
+    identity_migrator().run(pool).await?;
+    credentials_migrator().run(pool).await?;
+    admission_migrator().run(pool).await?;
+    attempt_migrator().run(pool).await?;
+    result_migrator().run(pool).await?;
+    deletion_migrator().run(pool).await?;
+    api_migrator().run(pool).await?;
+    worker_migrator().run(pool).await?;
+    control_migrator().run(pool).await?;
+    Ok(())
+}
+
+pub async fn generation(pool: &PgPool) -> Result<Option<i64>, sqlx::Error> {
+    sqlx::query_scalar(
+        "SELECT generation
+         FROM remora_schema_state
+         WHERE singleton",
+    )
+    .fetch_optional(pool)
+    .await
+}
+
+pub async fn is_ready(pool: &PgPool) -> Result<bool, sqlx::Error> {
+    generation(pool)
+        .await
+        .map(|generation| generation == Some(SCHEMA_GENERATION))
+}
+
+#[cfg(test)]
+fn component_migrators() -> [(&'static str, Migrator); 9] {
+    [
+        ("durable", durable_migrator()),
+        ("identity", identity_migrator()),
+        ("credentials", credentials_migrator()),
+        ("admission", admission_migrator()),
+        ("attempt", attempt_migrator()),
+        ("result", result_migrator()),
+        ("deletion", deletion_migrator()),
+        ("api", api_migrator()),
+        ("worker", worker_migrator()),
+    ]
+}
+
+fn durable_migrator() -> Migrator {
+    sqlx::migrate!("./migrations/durable")
+}
+
+fn identity_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/identity"), IDENTITY_HISTORY)
+}
+
+fn credentials_migrator() -> Migrator {
+    with_history(
+        sqlx::migrate!("./migrations/credentials"),
+        CREDENTIALS_HISTORY,
+    )
+}
+
+fn admission_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/admission"), ADMISSION_HISTORY)
+}
+
+fn attempt_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/attempt"), ATTEMPT_HISTORY)
+}
+
+fn result_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/result"), RESULT_HISTORY)
+}
+
+fn deletion_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/deletion"), DELETION_HISTORY)
+}
+
+fn api_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/api"), API_HISTORY)
+}
+
+fn worker_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/worker"), WORKER_HISTORY)
+}
+
+fn control_migrator() -> Migrator {
+    with_history(sqlx::migrate!("./migrations/control"), CONTROL_HISTORY)
+}
+
+fn with_history(mut migrator: Migrator, table: &'static str) -> Migrator {
+    migrator.dangerous_set_table_name(table);
+    migrator
+}
+
+#[cfg(test)]
+mod tests {
+    use super::{SCHEMA_GENERATION, component_migrators, control_migrator};
+
+    #[test]
+    fn embeds_every_component_in_dependency_order() {
+        let migrations: Vec<_> = component_migrators()
+            .into_iter()
+            .map(|(name, migrator)| (name, migrator.iter().count()))
+            .collect();
+        assert_eq!(
+            migrations,
+            [
+                ("durable", 1),
+                ("identity", 1),
+                ("credentials", 5),
+                ("admission", 1),
+                ("attempt", 6),
+                ("result", 2),
+                ("deletion", 1),
+                ("api", 2),
+                ("worker", 1),
+            ]
+        );
+        assert_eq!(
+            control_migrator()
+                .iter()
+                .last()
+                .map(|migration| migration.version),
+            Some(SCHEMA_GENERATION)
+        );
+    }
+}