feat: add a way to run data migrations individually

Author: ctron

Cargo.lock

@@ -24,6 +24,8 @@ serde-cyclonedx = { workspace = true }
 serde_json = { workspace = true }
 spdx-rs = { workspace = true }
 tokio = { workspace = true, features = ["full"] }
+tracing = { workspace = true }
+tracing-subscriber = { workspace = true }
 uuid = { workspace = true, features = ["v5"] }
 
 [dev-dependencies]

migration/Cargo.toml

@@ -0,0 +1,153 @@
+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 trustify_module_storage::config::StorageConfig;
+
+#[derive(clap::Parser, Debug, Clone)]
+struct Cli {
+    #[command(subcommand)]
+    command: Command,
+}
+
+#[derive(clap::Subcommand, Debug, Clone)]
+enum Command {
+    /// List all data migrations
+    List,
+    /// Run a list of migrations
+    Run(Run),
+}
+
+#[derive(clap::Args, Debug, Clone)]
+struct Run {
+    /// Migration direction to run
+    #[arg(
+        long,
+        value_enum,
+        default_value_t = Direction::Up,
+        overrides_with = "down"
+    )]
+    direction: Direction,
+
+    /// Shortcut for `--direction down`
+    #[arg(long, action = clap::ArgAction::SetTrue, overrides_with = "direction")]
+    down: bool,
+
+    // from sea_orm
+    #[arg(
+        global = true,
+        short = 's',
+        long,
+        env = "DATABASE_SCHEMA",
+        long_help = "Database schema\n \
+                    - For MySQL and SQLite, this argument is ignored.\n \
+                    - For PostgreSQL, this argument is optional with default value 'public'.\n"
+    )]
+    database_schema: Option<String>,
+
+    // from sea_orm
+    #[arg(
+        global = true,
+        short = 'u',
+        long,
+        env = "DATABASE_URL",
+        help = "Database URL"
+    )]
+    database_url: Option<String>,
+
+    #[arg()]
+    migrations: Vec<String>,
+
+    #[command(flatten)]
+    options: Options,
+
+    #[command(flatten)]
+    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 {
+            Direction::Down
+        } else {
+            self.direction
+        }
+    }
+
+    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?,
+            }
+        }
+
+        Ok(())
+    }
+}
+
+impl Command {
+    pub async fn run(self) -> anyhow::Result<()> {
+        match self {
+            Command::Run(run) => run.run().await,
+            Command::List => {
+                for m in Migrator::data_migrations() {
+                    println!("{}", m.name());
+                }
+                Ok(())
+            }
+        }
+    }
+}
+
+#[tokio::main]
+async fn main() {
+    let cli = Cli::parse();
+
+    tracing_subscriber::fmt::init();
+
+    cli.command.run().await.unwrap();
+}

migration/src/bin/data.rs

@@ -1,22 +1,21 @@
 use crate::{
     async_trait,
-    data::{Document, DocumentProcessor, Handler},
+    data::{Document, DocumentProcessor, Handler, Options},
 };
 use clap::Parser;
 use sea_orm::DbErr;
 use sea_orm_migration::{MigrationName, MigrationTrait, SchemaManager};
 use std::{ffi::OsString, sync::LazyLock};
 use trustify_module_storage::{config::StorageConfig, service::dispatch::DispatchBackend};
 
-pub struct MigrationWithData<M>
-where
-    M: MigrationTraitWithData,
-{
+pub struct MigrationWithData {
     pub storage: DispatchBackend,
-    pub migration: M,
+    pub options: Options,
+    pub migration: Box<dyn MigrationTraitWithData>,
 }
 
 static STORAGE: LazyLock<DispatchBackend> = LazyLock::new(init_storage);
+static OPTIONS: LazyLock<Options> = LazyLock::new(init_options);
 
 #[allow(clippy::expect_used)]
 fn init_storage() -> DispatchBackend {
@@ -33,76 +32,84 @@ fn init_storage() -> DispatchBackend {
     })
 }
 
-impl<M> MigrationWithData<M>
-where
-    M: MigrationTraitWithData,
-{
+fn init_options() -> Options {
+    // create from env-vars only
+    Options::parse_from::<_, OsString>(vec![])
+}
+
+impl MigrationWithData {
     #[allow(clippy::expect_used)]
-    pub fn new(migration: M) -> Self {
+    pub fn new(migration: Box<dyn MigrationTraitWithData>) -> Self {
         Self {
             storage: STORAGE.clone(),
+            options: OPTIONS.clone(),
             migration,
         }
     }
 }
 
-impl<M> From<M> for MigrationWithData<M>
+impl<M> From<M> for MigrationWithData
 where
-    M: MigrationTraitWithData,
+    M: MigrationTraitWithData + 'static,
 {
     fn from(value: M) -> Self {
-        MigrationWithData::new(value)
+        MigrationWithData::new(Box::new(value))
     }
 }
 
 pub struct SchemaDataManager<'c> {
     pub manager: &'c SchemaManager<'c>,
     storage: &'c DispatchBackend,
+    options: &'c Options,
 }
 
 impl<'c> SchemaDataManager<'c> {
-    pub fn new(manager: &'c SchemaManager<'c>, storage: &'c DispatchBackend) -> Self {
-        Self { manager, storage }
+    pub fn new(
+        manager: &'c SchemaManager<'c>,
+        storage: &'c DispatchBackend,
+        options: &'c Options,
+    ) -> Self {
+        Self {
+            manager,
+            storage,
+            options,
+        }
     }
 
     pub async fn process<D>(&self, f: impl Handler<D>) -> Result<(), DbErr>
     where
         D: Document,
     {
-        self.manager
-            .process(self.storage, Default::default(), f)
-            .await
+        self.manager.process(self.storage, self.options, f).await
     }
 }
 
 #[async_trait::async_trait]
-pub trait MigrationTraitWithData {
+pub trait MigrationTraitWithData: MigrationName + Send + Sync {
     async fn up(&self, manager: &SchemaDataManager) -> Result<(), DbErr>;
     async fn down(&self, manager: &SchemaDataManager) -> Result<(), DbErr>;
 }
 
 #[async_trait::async_trait]
-impl<M> MigrationTrait for MigrationWithData<M>
-where
-    M: MigrationTraitWithData + MigrationName + Send + Sync,
-{
+impl MigrationTrait for MigrationWithData {
     async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
-        self.migration
-            .up(&SchemaDataManager::new(manager, &self.storage))
-            .await
+        MigrationTraitWithData::up(
+            &*self.migration,
+            &SchemaDataManager::new(manager, &self.storage, &self.options),
+        )
+        .await
     }
 
     async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
-        self.migration
-            .down(&SchemaDataManager::new(manager, &self.storage))
-            .await
+        MigrationTraitWithData::down(
+            &*self.migration,
+            &SchemaDataManager::new(manager, &self.storage, &self.options),
+        )
+        .await
     }
 }
 
-impl<M> MigrationName for MigrationWithData<M>
-where
-    M: MigrationTraitWithData + MigrationName + Send + Sync,
-{
+impl MigrationName for MigrationWithData {
     fn name(&self) -> &str {
         self.migration.name()
     }

migration/src/data/migration.rs

@@ -90,8 +90,9 @@ where
     ) -> anyhow::Result<()>;
 }
 
-#[derive(Clone, Debug, PartialEq, Eq)]
+#[derive(Clone, Debug, PartialEq, Eq, clap::Parser)]
 pub struct Options {
+    #[arg(long, env = "MIGRATION_DATA_CONCURRENT", default_value = "5")]
     pub concurrent: NonZeroUsize,
 }
 
@@ -107,7 +108,7 @@ pub trait DocumentProcessor {
     async fn process<D>(
         &self,
         storage: &DispatchBackend,
-        options: Options,
+        options: &Options,
         f: impl Handler<D>,
     ) -> anyhow::Result<(), DbErr>
     where
@@ -118,34 +119,37 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
     async fn process<D>(
         &self,
         storage: &DispatchBackend,
-        options: Options,
+        options: &Options,
         f: impl Handler<D>,
     ) -> Result<(), DbErr>
     where
         D: Document,
     {
         let db = self.get_connection();
-        let tx = db.begin().await?;
 
+        let tx = db.begin().await?;
         let all = D::all(&tx).await?;
+        drop(tx);
 
         stream::iter(all)
             .map(async |model| {
+                let tx = db.begin().await?;
+
                 let doc = D::source(&model, storage, &tx).await.map_err(|err| {
                     DbErr::Migration(format!("Failed to load source document: {err}"))
                 })?;
                 f.call(doc, model, &tx).await.map_err(|err| {
                     DbErr::Migration(format!("Failed to process document: {err}"))
                 })?;
 
+                tx.commit().await?;
+
                 Ok::<_, DbErr>(())
             })
             .buffer_unordered(options.concurrent.into())
             .try_collect::<Vec<_>>()
             .await?;
 
-        tx.commit().await?;
-
         Ok(())
     }
 }