refactor: use MigratorTraitSelf to eliminate global migration state

Cargo.toml

@@ -118,6 +118,7 @@ sbom-walker = { version = "0.19.0", default-features = false, features = ["crypt
 scheck = "0.8.0"
 schemars = "1.0"
 sea-orm = { version = "2.0.4", features = ["debug-print"] }
+sea-orm-cli = { version = "2", default-features = false }
 sea-orm-migration = "2"
 sea-query = "1.0.2" # keep aligned with sea-orm
 semver = "1"

common/db/src/embedded.rs

@@ -1,4 +1,5 @@
 use anyhow::Context;
+use migration::Migrator;
 use postgresql_embedded::{PostgreSQL, Settings, VersionReq};
 use std::path::{Path, PathBuf};
 use tracing::{Instrument, info_span};
@@ -19,12 +20,15 @@ pub fn default_settings() -> anyhow::Result<Settings> {
 }
 
 /// Create a new, embedded database instance
-pub async fn create() -> anyhow::Result<(Database, PostgreSQL)> {
-    create_for(default_settings()?, Default::default()).await
+pub async fn create(migrator: &Migrator) -> anyhow::Result<(Database, PostgreSQL)> {
+    create_for(default_settings()?, Default::default(), migrator).await
 }
 
 /// Create a new, embedded database instance in a specific directory
-pub async fn create_in(base: impl AsRef<Path>) -> anyhow::Result<(Database, PostgreSQL)> {
+pub async fn create_in(
+    base: impl AsRef<Path>,
+    migrator: &Migrator,
+) -> anyhow::Result<(Database, PostgreSQL)> {
     let base = base.as_ref();
 
     create_for(
@@ -34,6 +38,7 @@ pub async fn create_in(base: impl AsRef<Path>) -> anyhow::Result<(Database, Post
             ..default_settings()?
         },
         Default::default(),
+        migrator,
     )
     .await
 }
@@ -54,6 +59,7 @@ pub struct Options {
 pub async fn create_for(
     settings: Settings,
     options: Options,
+    migrator: &Migrator,
 ) -> anyhow::Result<(Database, PostgreSQL)> {
     log::info!("creating embedded database - version: {}", settings.version);
 
@@ -75,7 +81,7 @@ pub async fn create_for(
     let config = crate::config::Database::from_port(postgresql.settings().port)?;
 
     let db = match options.source {
-        Source::Bootstrap => super::Database::bootstrap(&config)
+        Source::Bootstrap => super::Database::bootstrap(&config, migrator)
             .await
             .context("Bootstrapping the test database")?,
         Source::Import(path) => {

common/db/src/lib.rs

@@ -2,12 +2,13 @@
 
 pub mod embedded;
 
+pub use migration::Migrator;
+
 use anyhow::{Context, anyhow, ensure};
-use migration::Migrator;
 use migration::data::Runner;
 use postgresql_commands::{CommandBuilder, psql::PsqlBuilder};
 use sea_orm::ConnectionTrait;
-use sea_orm_migration::prelude::{MigrationTrait, MigratorTrait};
+use sea_orm_migration::{MigrationTrait, MigratorTraitSelf};
 use std::process::Stdio;
 use tokio::io::{self, AsyncRead, AsyncWriteExt};
 use tracing::instrument;
@@ -16,24 +17,28 @@ use trustify_common::{config, db};
 pub struct Database<'a>(pub &'a db::Database);
 
 impl<'a> Database<'a> {
-    #[instrument(skip(self), err(level=tracing::Level::INFO))]
-    pub async fn migrate(&self) -> Result<(), anyhow::Error> {
+    #[instrument(skip(self, migrator), err(level=tracing::Level::INFO))]
+    pub async fn migrate(&self, migrator: &Migrator) -> Result<(), anyhow::Error> {
         log::debug!("applying migrations");
-        Migrator::up(self.0, None).await?;
+        migrator.up(self.0, None).await?;
         log::debug!("applied migrations");
 
         Ok(())
     }
 
     /// Apply migrations up to and including the one matching the given name.
-    #[instrument(skip(self), err(level=tracing::Level::INFO))]
-    pub async fn migrate_up_to(&self, name: &str) -> Result<(), anyhow::Error> {
-        let all_migrations = Migrator::migrations();
+    #[instrument(skip(self, migrator), err(level=tracing::Level::INFO))]
+    pub async fn migrate_up_to(
+        &self,
+        migrator: &Migrator,
+        name: &str,
+    ) -> Result<(), anyhow::Error> {
+        let all_migrations = migrator.migrations();
 
         let target_pos = find_migration_position(&all_migrations, name)?;
         let target_name = all_migrations[target_pos].name().to_string();
 
-        let applied = Migrator::get_applied_migrations(self.0).await?;
+        let applied = migrator.get_applied_migrations(self.0).await?;
         let applied_count = applied.len() as u32;
         let target_count = (target_pos as u32) + 1;
 
@@ -44,16 +49,16 @@ impl<'a> Database<'a> {
 
         let steps = target_count - applied_count;
         log::debug!("applying {steps} migration(s) up to '{target_name}'");
-        Migrator::up(self.0, Some(steps)).await?;
+        migrator.up(self.0, Some(steps)).await?;
         log::debug!("applied migrations up to '{target_name}'");
 
         Ok(())
     }
 
-    #[instrument(skip(self), err(level=tracing::Level::INFO))]
-    pub async fn refresh(&self) -> Result<(), anyhow::Error> {
+    #[instrument(skip(self, migrator), err(level=tracing::Level::INFO))]
+    pub async fn refresh(&self, migrator: &Migrator) -> Result<(), anyhow::Error> {
         log::warn!("refreshing database schema...");
-        Migrator::refresh(self.0).await?;
+        migrator.refresh(self.0).await?;
         log::warn!("refreshing database schema... done!");
 
         Ok(())
@@ -94,11 +99,14 @@ impl<'a> Database<'a> {
         Ok(db)
     }
 
-    #[instrument(err(level=tracing::Level::INFO))]
-    pub async fn bootstrap(database: &config::Database) -> Result<db::Database, anyhow::Error> {
+    #[instrument(skip(migrator), err(level=tracing::Level::INFO))]
+    pub async fn bootstrap(
+        database: &config::Database,
+        migrator: &Migrator,
+    ) -> Result<db::Database, anyhow::Error> {
         let db = Self::setup(database).await?;
 
-        Database(&db).migrate().await?;
+        Database(&db).migrate(migrator).await?;
 
         Ok(db)
     }

migration/Cargo.toml

@@ -24,12 +24,12 @@ bytes = { workspace = true }
 clap = { workspace = true, features = ["derive", "env"] }
 csaf-rs = { workspace = true }
 cve = { workspace = true }
-futures = { workspace = true }
 futures-util = { workspace = true }
 humantime = { workspace = true }
 indicatif = { workspace = true, features = ["tokio", "futures"] }
 osv = { workspace = true, features = ["schema"] }
 sea-orm = { workspace = true }
+sea-orm-cli = { workspace = true }
 sea-orm-migration = { workspace = true, features = ["runtime-tokio-rustls", "sqlx-postgres", "with-uuid"] }
 serde-cyclonedx = { workspace = true }
 serde_json = { workspace = true }

migration/src/data/migration.rs

@@ -2,14 +2,11 @@ use crate::{
     async_trait,
     data::{Document, DocumentProcessor, Handler, Options},
 };
-use clap::Parser;
-use futures::executor::block_on;
 use sea_orm::{DatabaseConnection, DbErr};
 use sea_orm_migration::{MigrationName, MigrationTrait, SchemaManager};
-use std::{ffi::OsString, ops::Deref, sync::LazyLock};
-use tokio::task_local;
+use std::ops::Deref;
 use tracing::log;
-use trustify_module_storage::{config::StorageConfig, service::dispatch::DispatchBackend};
+use trustify_module_storage::service::dispatch::DispatchBackend;
 
 /// A migration which also processes data.
 pub struct MigrationWithData {
@@ -18,68 +15,6 @@ pub struct MigrationWithData {
     pub migration: Box<dyn MigrationTraitWithData>,
 }
 
-static STORAGE: LazyLock<DispatchBackend> = LazyLock::new(init_storage);
-static OPTIONS: LazyLock<Options> = LazyLock::new(init_options);
-
-task_local! {
-    static TEST_STORAGE: DispatchBackend;
-    static TEST_OPTIONS: Options;
-}
-
-#[allow(clippy::expect_used)]
-fn init_storage() -> DispatchBackend {
-    // create from env-vars only
-    let config = StorageConfig::parse_from::<_, OsString>(vec![]);
-
-    block_on(config.into_storage(false)).expect("task panicked")
-}
-
-fn init_options() -> Options {
-    // create from env-vars only
-    Options::parse_from::<_, OsString>(vec![])
-}
-
-impl MigrationWithData {
-    /// Wrap a data migration, turning it into a combined schema/data migration.
-    ///
-    /// **NOTE:** This may panic if the storage configuration is missing.
-    pub fn new(migration: Box<dyn MigrationTraitWithData>) -> Self {
-        // if we have a test storage set, use this instead.
-        let storage = TEST_STORAGE
-            .try_with(|s| s.clone())
-            .unwrap_or_else(|_| STORAGE.clone());
-
-        let options = TEST_OPTIONS
-            .try_with(|o| o.clone())
-            .unwrap_or_else(|_| OPTIONS.clone());
-
-        Self {
-            storage,
-            options,
-            migration,
-        }
-    }
-
-    /// Set a storage backend to be used for running tests.
-    ///
-    /// This will, for the duration of the call, initialize the migrator with the provided storage
-    /// backend.
-    pub async fn run_with_test<F>(
-        storage: impl Into<DispatchBackend>,
-        options: impl Into<Options>,
-        f: F,
-    ) -> F::Output
-    where
-        F: Future,
-    {
-        TEST_STORAGE
-            .scope(storage.into(), async {
-                TEST_OPTIONS.scope(options.into(), f).await
-            })
-            .await
-    }
-}
-
 /// A [`SchemaManager`], extended with data migration features.
 pub struct SchemaDataManager<'c> {
     pub manager: &'c SchemaManager<'c>,

migration/src/data/mod.rs

@@ -357,4 +357,23 @@ impl Migrations {
         self.all.push(Migration::Data(Box::new(migration)));
         self
     }
+
+    /// Convert into SeaORM's `Vec<Box<dyn MigrationTrait>>`, wrapping data migrations with the
+    /// provided storage backend and options.
+    pub fn into_migration_vec(
+        self,
+        storage: DispatchBackend,
+        options: Options,
+    ) -> Vec<Box<dyn MigrationTrait>> {
+        self.into_iter()
+            .map(|migration| match migration {
+                Migration::Normal(m) => m,
+                Migration::Data(m) => Box::new(MigrationWithData {
+                    storage: storage.clone(),
+                    options: options.clone(),
+                    migration: m,
+                }) as Box<dyn MigrationTrait>,
+            })
+            .collect()
+    }
 }

migration/src/lib.rs

@@ -1,9 +1,9 @@
 #![recursion_limit = "1024"]
-use crate::data::{
-    Migration, MigrationTraitWithData, MigrationWithData, Migrations, MigratorWithData,
-};
+use crate::data::{MigrationTraitWithData, Migrations, MigratorWithData, Options};
 
+pub use sea_orm_migration::MigratorTraitSelf;
 pub use sea_orm_migration::prelude::*;
+use trustify_module_storage::service::dispatch::DispatchBackend;
 
 pub mod data;
 
@@ -85,19 +85,21 @@ mod m0002380_version_function_parallelism;
 
 pub trait MigratorExt: Send {
     fn build_migrations() -> Migrations;
+}
 
-    fn into_migrations() -> Vec<Box<dyn MigrationTrait>> {
-        Self::build_migrations()
-            .into_iter()
-            .map(|migration| match migration {
-                Migration::Normal(migration) => migration,
-                Migration::Data(migration) => Box::new(MigrationWithData::new(migration)),
-            })
-            .collect()
-    }
+pub struct Migrator {
+    pub storage: DispatchBackend,
+    pub options: Options,
 }
 
-pub struct Migrator;
+impl Migrator {
+    pub fn new(storage: impl Into<DispatchBackend>, options: impl Into<Options>) -> Self {
+        Self {
+            storage: storage.into(),
+            options: options.into(),
+        }
+    }
+}
 
 impl MigratorExt for Migrator {
     fn build_migrations() -> Migrations {
@@ -187,9 +189,9 @@ impl<M: MigratorExt> MigratorWithData for M {
 }
 
 #[async_trait::async_trait]
-impl MigratorTrait for Migrator {
-    fn migrations() -> Vec<Box<dyn MigrationTrait>> {
-        Self::into_migrations()
+impl MigratorTraitSelf for Migrator {
+    fn migrations(&self) -> Vec<Box<dyn MigrationTrait>> {
+        Self::build_migrations().into_migration_vec(self.storage.clone(), self.options.clone())
     }
 }