feat: add way to run data migrations from main binary

Author: ctron

Cargo.lock

@@ -2,6 +2,7 @@ pub mod embedded;
 
 use anyhow::ensure;
 use migration::Migrator;
+use migration::data::Runner;
 use sea_orm::{ConnectionTrait, Statement};
 use sea_orm_migration::prelude::MigratorTrait;
 use tracing::instrument;
@@ -35,7 +36,7 @@ impl<'a> Database<'a> {
             "Unable to bootstrap database with '--db-url'"
         );
 
-        let url = crate::config::Database {
+        let url = config::Database {
             name: "postgres".into(),
             ..database.clone()
         }
@@ -64,4 +65,8 @@ impl<'a> Database<'a> {
 
         Ok(db)
     }
+
+    pub async fn data_migrate(&self, runner: Runner) -> Result<(), anyhow::Error> {
+        runner.run::<Migrator>().await
+    }
 }

common/db/src/lib.rs

@@ -17,6 +17,7 @@ trustify-module-storage = { workspace = true }
 anyhow = { workspace = true }
 bytes = { workspace = true }
 clap = { workspace = true, features = ["derive", "env"] }
+futures = { workspace = true }
 futures-util = { workspace = true }
 sea-orm = { workspace = true }
 sea-orm-migration = { workspace = true, features = ["runtime-tokio-rustls", "sqlx-postgres", "with-uuid"] }

migration/Cargo.toml

@@ -1,9 +1,8 @@
-use anyhow::bail;
 use clap::Parser;
-use migration::{Migrator, Options, SchemaDataManager};
-use sea_orm::{ConnectOptions, Database};
-use sea_orm_migration::{IntoSchemaManagerConnection, SchemaManager};
-use std::collections::HashMap;
+use migration::{
+    Migrator,
+    data::{Direction, MigratorWithData, Options, Runner},
+};
 use trustify_module_storage::config::StorageConfig;
 
 #[derive(clap::Parser, Debug, Clone)]
@@ -13,6 +12,7 @@ struct Cli {
 }
 
 #[derive(clap::Subcommand, Debug, Clone)]
+#[allow(clippy::large_enum_variant)]
 enum Command {
     /// List all data migrations
     List,
@@ -67,13 +67,6 @@ struct Run {
     storage: StorageConfig,
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)]
-pub enum Direction {
-    #[default]
-    Up,
-    Down,
-}
-
 impl Run {
     fn direction(&self) -> Direction {
         if self.down {
@@ -83,47 +76,23 @@ impl Run {
         }
     }
 
+    #[allow(clippy::expect_used)]
     pub async fn run(self) -> anyhow::Result<()> {
         let direction = self.direction();
-
-        let migrations = Migrator::data_migrations()
-            .into_iter()
-            .map(|migration| (migration.name().to_string(), migration))
-            .collect::<HashMap<_, _>>();
-
-        let mut running = vec![];
-
-        for migration in self.migrations {
-            let Some(migration) = migrations.get(&migration) else {
-                bail!("Migration {migration} not found");
-            };
-            running.push(migration);
-        }
-
         let storage = self.storage.into_storage(false).await?;
 
-        let url = self
-            .database_url
-            .expect("Environment variable 'DATABASE_URL' not set");
-        let schema = self.database_schema.unwrap_or_else(|| "public".to_owned());
-
-        let connect_options = ConnectOptions::new(url)
-            .set_schema_search_path(schema)
-            .to_owned();
-
-        let db = Database::connect(connect_options).await?;
-
-        let manager = SchemaManager::new(db.into_schema_manager_connection());
-        let manager = SchemaDataManager::new(&manager, &storage, &self.options);
-
-        for run in running {
-            tracing::info!("Running data migration: {}", run.name());
-
-            match direction {
-                Direction::Up => run.up(&manager).await?,
-                Direction::Down => run.down(&manager).await?,
-            }
+        Runner {
+            direction,
+            storage,
+            migrations: self.migrations,
+            database_url: self
+                .database_url
+                .expect("Environment variable 'DATABASE_URL' not set"),
+            database_schema: self.database_schema,
+            options: self.options,
         }
+        .run::<Migrator>()
+        .await?;
 
         Ok(())
     }
@@ -143,6 +112,7 @@ impl Command {
     }
 }
 
+#[allow(clippy::unwrap_used)]
 #[tokio::main]
 async fn main() {
     let cli = Cli::parse();

migration/src/bin/data.rs

@@ -3,6 +3,7 @@ use crate::{
     data::{Document, DocumentProcessor, Handler, Options},
 };
 use clap::Parser;
+use futures::executor::block_on;
 use sea_orm::DbErr;
 use sea_orm_migration::{MigrationName, MigrationTrait, SchemaManager};
 use std::{ffi::OsString, sync::LazyLock};
@@ -22,14 +23,7 @@ fn init_storage() -> DispatchBackend {
     // create from env-vars only
     let config = StorageConfig::parse_from::<_, OsString>(vec![]);
 
-    tokio::task::block_in_place(|| {
-        tokio::runtime::Handle::current().block_on(async {
-            config
-                .into_storage(false)
-                .await
-                .expect("Failed to create storage")
-        })
-    })
+    block_on(config.into_storage(false)).expect("task panicked")
 }
 
 fn init_options() -> Options {

migration/src/data/migration.rs

@@ -1,5 +1,8 @@
 mod migration;
+mod run;
+
 pub use migration::*;
+pub use run::*;
 
 use anyhow::{anyhow, bail};
 use bytes::BytesMut;
@@ -10,7 +13,7 @@ use futures_util::{
 use sea_orm::{
     ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, ModelTrait, TransactionTrait,
 };
-use sea_orm_migration::SchemaManager;
+use sea_orm_migration::{MigrationTrait, SchemaManager};
 use std::num::NonZeroUsize;
 use trustify_common::id::Id;
 use trustify_entity::{sbom, source_document};
@@ -105,12 +108,12 @@ impl Default for Options {
 }
 
 pub trait DocumentProcessor {
-    async fn process<D>(
+    fn process<D>(
         &self,
         storage: &DispatchBackend,
         options: &Options,
         f: impl Handler<D>,
-    ) -> anyhow::Result<(), DbErr>
+    ) -> impl Future<Output = anyhow::Result<(), DbErr>>
     where
         D: Document;
 }
@@ -180,3 +183,42 @@ macro_rules! sbom {
         $crate::handler!(async |$doc: $crate::data::Sbom, $model, $tx| $body)
     };
 }
+
+pub trait MigratorWithData {
+    fn data_migrations() -> Vec<Box<dyn MigrationTraitWithData>>;
+}
+
+#[derive(Default)]
+pub struct Migrations {
+    all: Vec<Migration>,
+}
+
+impl IntoIterator for Migrations {
+    type Item = Migration;
+    type IntoIter = std::vec::IntoIter<Self::Item>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.all.into_iter()
+    }
+}
+
+pub enum Migration {
+    Normal(Box<dyn MigrationTrait>),
+    Data(Box<dyn MigrationTraitWithData>),
+}
+
+impl Migrations {
+    pub fn new() -> Self {
+        Self::default()
+    }
+
+    pub fn normal(mut self, migration: impl MigrationTrait + 'static) -> Self {
+        self.all.push(Migration::Normal(Box::new(migration)));
+        self
+    }
+
+    pub fn data(mut self, migration: impl MigrationTraitWithData + 'static) -> Self {
+        self.all.push(Migration::Data(Box::new(migration)));
+        self
+    }
+}

migration/src/data/mod.rs

@@ -0,0 +1,62 @@
+use crate::data::{MigratorWithData, Options, SchemaDataManager};
+use anyhow::bail;
+use sea_orm::{ConnectOptions, Database};
+use sea_orm_migration::{IntoSchemaManagerConnection, SchemaManager};
+use std::collections::HashMap;
+use trustify_module_storage::service::dispatch::DispatchBackend;
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, clap::ValueEnum)]
+pub enum Direction {
+    #[default]
+    Up,
+    Down,
+}
+
+pub struct Runner {
+    pub database_url: String,
+    pub database_schema: Option<String>,
+    pub storage: DispatchBackend,
+    pub direction: Direction,
+    pub migrations: Vec<String>,
+    pub options: Options,
+}
+
+impl Runner {
+    pub async fn run<M: MigratorWithData>(self) -> anyhow::Result<()> {
+        let migrations = M::data_migrations()
+            .into_iter()
+            .map(|migration| (migration.name().to_string(), migration))
+            .collect::<HashMap<_, _>>();
+
+        let mut running = vec![];
+
+        for migration in self.migrations {
+            let Some(migration) = migrations.get(&migration) else {
+                bail!("Migration '{migration}' not found");
+            };
+            running.push(migration);
+        }
+
+        let schema = self.database_schema.unwrap_or_else(|| "public".to_owned());
+
+        let connect_options = ConnectOptions::new(self.database_url)
+            .set_schema_search_path(schema)
+            .to_owned();
+
+        let db = Database::connect(connect_options).await?;
+
+        let manager = SchemaManager::new(db.into_schema_manager_connection());
+        let manager = SchemaDataManager::new(&manager, &self.storage, &self.options);
+
+        for run in running {
+            tracing::info!(name = run.name(), "Running data migration");
+
+            match self.direction {
+                Direction::Up => run.up(&manager).await?,
+                Direction::Down => run.down(&manager).await?,
+            }
+        }
+
+        Ok(())
+    }
+}

migration/src/data/run.rs

@@ -1,7 +1,9 @@
+use crate::data::{
+    Migration, MigrationTraitWithData, MigrationWithData, Migrations, MigratorWithData,
+};
 pub use sea_orm_migration::prelude::*;
 
-mod data;
-pub use crate::data::{MigrationTraitWithData, MigrationWithData, Options, SchemaDataManager};
+pub mod data;
 
 mod m0000010_init;
 mod m0000020_add_sbom_group;
@@ -30,41 +32,6 @@ mod m0001120_sbom_external_node_indexes;
 mod m0001130_gover_cmp;
 mod m0001140_example_data_migration;
 
-#[derive(Default)]
-pub struct Migrations {
-    all: Vec<Migration>,
-}
-
-impl IntoIterator for Migrations {
-    type Item = Migration;
-    type IntoIter = std::vec::IntoIter<Self::Item>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.all.into_iter()
-    }
-}
-
-pub enum Migration {
-    Normal(Box<dyn MigrationTrait>),
-    Data(Box<dyn MigrationTraitWithData>),
-}
-
-impl Migrations {
-    pub fn new() -> Self {
-        Self::default()
-    }
-
-    pub fn normal(mut self, migration: impl MigrationTrait + 'static) -> Self {
-        self.all.push(Migration::Normal(Box::new(migration)));
-        self
-    }
-
-    pub fn data(mut self, migration: impl MigrationTraitWithData + 'static) -> Self {
-        self.all.push(Migration::Data(Box::new(migration)));
-        self
-    }
-}
-
 pub struct Migrator;
 
 impl Migrator {
@@ -97,8 +64,10 @@ impl Migrator {
             .normal(m0001130_gover_cmp::Migration)
             .data(m0001140_example_data_migration::Migration)
     }
+}
 
-    pub fn data_migrations() -> Vec<Box<dyn MigrationTraitWithData>> {
+impl MigratorWithData for Migrator {
+    fn data_migrations() -> Vec<Box<dyn MigrationTraitWithData>> {
         Self::migrations()
             .into_iter()
             .filter_map(|migration| match migration {
@@ -112,6 +81,7 @@ impl Migrator {
 #[async_trait::async_trait]
 impl MigratorTrait for Migrator {
     fn migrations() -> Vec<Box<dyn MigrationTrait>> {
+        // Get all migrations, wrap data migrations. This will initialize the storage config.
         Self::migrations()
             .into_iter()
             .map(|migration| match migration {