feat: allow running data migrations as part of migrations

ctron · ctron · commit 391b790a6692 · 2025-09-10T10:47:47.000+02:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/migration/Cargo.toml b/migration/Cargo.toml
@@ -14,9 +14,10 @@ trustify-common = { workspace = true }
 trustify-entity = { workspace = true }
 trustify-module-storage = { workspace = true }
 
+anyhow = { workspace = true }
 bytes = { workspace = true }
+clap = { workspace = true, features = ["derive", "env"] }
 futures-util = { workspace = true }
-anyhow = { workspace = true }
 sea-orm = { workspace = true }
 sea-orm-migration = { workspace = true, features = ["runtime-tokio-rustls", "sqlx-postgres", "with-uuid"] }
 serde-cyclonedx = { workspace = true }
diff --git a/migration/src/data/migration.rs b/migration/src/data/migration.rs
@@ -0,0 +1,109 @@
+use crate::{
+    async_trait,
+    data::{Document, DocumentProcessor, Handler},
+};
+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 storage: DispatchBackend,
+    pub migration: M,
+}
+
+static STORAGE: LazyLock<DispatchBackend> = LazyLock::new(init_storage);
+
+#[allow(clippy::expect_used)]
+fn init_storage() -> DispatchBackend {
+    // create from env-vars only
+    let config = StorageConfig::parse_from::<_, OsString>(vec![]);
+
+    tokio::runtime::Builder::new_current_thread()
+        .enable_all()
+        .build()
+        .expect("failed to build runtime")
+        .block_on(async {
+            config
+                .into_storage(false)
+                .await
+                .expect("Failed to create storage")
+        })
+}
+
+impl<M> MigrationWithData<M>
+where
+    M: MigrationTraitWithData,
+{
+    #[allow(clippy::expect_used)]
+    pub fn new(migration: M) -> Self {
+        Self {
+            storage: STORAGE.clone(),
+            migration,
+        }
+    }
+}
+
+impl<M> From<M> for MigrationWithData<M>
+where
+    M: MigrationTraitWithData,
+{
+    fn from(value: M) -> Self {
+        MigrationWithData::new(value)
+    }
+}
+
+pub struct SchemaDataManager<'c> {
+    pub manager: &'c SchemaManager<'c>,
+    storage: &'c DispatchBackend,
+}
+
+impl<'c> SchemaDataManager<'c> {
+    pub fn new(manager: &'c SchemaManager<'c>, storage: &'c DispatchBackend) -> Self {
+        Self { manager, storage }
+    }
+
+    pub async fn process<D>(&self, f: impl Handler<D>) -> Result<(), DbErr>
+    where
+        D: Document,
+    {
+        self.manager.process(self.storage, f).await
+    }
+}
+
+#[async_trait::async_trait]
+pub trait MigrationTraitWithData {
+    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,
+{
+    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        self.migration
+            .up(&SchemaDataManager::new(manager, &self.storage))
+            .await
+    }
+
+    async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
+        self.migration
+            .down(&SchemaDataManager::new(manager, &self.storage))
+            .await
+    }
+}
+
+impl<M> MigrationName for MigrationWithData<M>
+where
+    M: MigrationTraitWithData + MigrationName + Send + Sync,
+{
+    fn name(&self) -> &str {
+        self.migration.name()
+    }
+}
diff --git a/migration/src/data/mod.rs b/migration/src/data/mod.rs
@@ -1,7 +1,13 @@
+mod migration;
+
+pub use migration::*;
+
 use anyhow::{anyhow, bail};
 use bytes::BytesMut;
-use futures_util::stream::TryStreamExt;
-use futures_util::{StreamExt, stream};
+use futures_util::{
+    StreamExt,
+    stream::{self, TryStreamExt},
+};
 use sea_orm::{
     ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, ModelTrait, TransactionTrait,
 };
@@ -16,6 +22,7 @@ pub enum Sbom {
     Spdx(spdx_rs::models::SPDX),
 }
 
+#[allow(async_fn_in_trait)]
 pub trait Document: Sized + Send + Sync {
     type Model: Send;
 
@@ -70,6 +77,7 @@ impl Document for Sbom {
     }
 }
 
+#[allow(async_fn_in_trait)]
 pub trait Handler<D>: Send
 where
     D: Document,
@@ -93,11 +101,7 @@ pub trait DocumentProcessor {
 }
 
 impl<'c> DocumentProcessor for SchemaManager<'c> {
-    async fn process<D>(
-        &self,
-        storage: &DispatchBackend,
-        f: impl Handler<D>,
-    ) -> anyhow::Result<(), DbErr>
+    async fn process<D>(&self, storage: &DispatchBackend, f: impl Handler<D>) -> Result<(), DbErr>
     where
         D: Document,
     {
diff --git a/migration/src/lib.rs b/migration/src/lib.rs
@@ -1,6 +1,7 @@
 pub use sea_orm_migration::prelude::*;
 
 mod data;
+use crate::data::MigrationWithData;
 
 mod m0000010_init;
 mod m0000020_add_sbom_group;
@@ -60,7 +61,9 @@ impl MigratorTrait for Migrator {
             Box::new(m0001110_sbom_node_checksum_indexes::Migration),
             Box::new(m0001120_sbom_external_node_indexes::Migration),
             Box::new(m0001130_gover_cmp::Migration),
-            Box::new(m0001140_example_data_migration::Migration),
+            Box::new(MigrationWithData::new(
+                m0001140_example_data_migration::Migration,
+            )),
         ]
     }
 }
diff --git a/migration/src/m0001140_example_data_migration.rs b/migration/src/m0001140_example_data_migration.rs
@@ -1,52 +1,40 @@
 use crate::{
-    data::{DocumentProcessor, Sbom},
+    data::{MigrationTraitWithData, Sbom, SchemaDataManager},
     sbom,
     sea_orm::{ActiveModelTrait, IntoActiveModel, Set},
 };
 use sea_orm_migration::prelude::*;
-use trustify_module_storage::service::{dispatch::DispatchBackend, fs::FileSystemBackend};
 
 #[derive(DeriveMigrationName)]
 pub struct Migration;
 
 #[async_trait::async_trait]
-impl MigrationTrait for Migration {
-    async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
-        // TODO: make this configurable
-        let (storage, _tmp) = FileSystemBackend::for_test()
-            .await
-            .map_err(|err| DbErr::Migration(format!("failed to create storage backend: {err}")))?;
-        let storage = DispatchBackend::Filesystem(storage);
-
-        // process data
-
+impl MigrationTraitWithData for Migration {
+    async fn up(&self, manager: &SchemaDataManager) -> Result<(), DbErr> {
         manager
-            .process(
-                &storage,
-                sbom!(async |sbom, model, tx| {
-                    let mut model = model.into_active_model();
-                    match sbom {
-                        Sbom::CycloneDx(_sbom) => {
-                            // TODO: just an example
-                            model.authors = Set(vec![]);
-                        }
-                        Sbom::Spdx(_sbom) => {
-                            // TODO: just an example
-                            model.authors = Set(vec![]);
-                        }
+            .process(sbom!(async |sbom, model, tx| {
+                let mut model = model.into_active_model();
+                match sbom {
+                    Sbom::CycloneDx(_sbom) => {
+                        // TODO: just an example
+                        model.authors = Set(vec![]);
+                    }
+                    Sbom::Spdx(_sbom) => {
+                        // TODO: just an example
+                        model.authors = Set(vec![]);
                     }
+                }
 
-                    model.save(tx).await?;
+                model.save(tx).await?;
 
-                    Ok(())
-                }),
-            )
+                Ok(())
+            }))
             .await?;
 
         Ok(())
     }
 
-    async fn down(&self, _manager: &SchemaManager) -> Result<(), DbErr> {
+    async fn down(&self, _manager: &SchemaDataManager) -> Result<(), DbErr> {
         Ok(())
     }
 }
diff --git a/trustd/src/db.rs b/trustd/src/db.rs
@@ -1,11 +1,6 @@
 use postgresql_embedded::{PostgreSQL, VersionReq};
-use std::collections::HashMap;
-use std::env;
-use std::fs::create_dir_all;
-use std::process::ExitCode;
-use std::time::Duration;
-use trustify_common::config::Database;
-use trustify_common::db;
+use std::{collections::HashMap, env, fs::create_dir_all, process::ExitCode, time::Duration};
+use trustify_common::{config::Database, db};
 use trustify_infrastructure::otel::{Tracing, init_tracing};
 
 #[derive(clap::Args, Debug)]
@@ -40,6 +35,7 @@ impl Run {
             Err(e) => Err(e),
         }
     }
+
     async fn refresh(self) -> anyhow::Result<ExitCode> {
         match db::Database::new(&self.database).await {
             Ok(db) => {
@@ -49,6 +45,7 @@ impl Run {
             Err(e) => Err(e),
         }
     }
+
     async fn migrate(self) -> anyhow::Result<ExitCode> {
         match db::Database::new(&self.database).await {
             Ok(db) => {

Original file line number	Diff line number	Diff line change
`@@ -1,6 +1,7 @@`
`1`	`1`	`pub use sea_orm_migration::prelude::*;`
`2`	`2`
`3`	`3`	`mod data;`
	`4`	`+use crate::data::MigrationWithData;`
`4`	`5`
`5`	`6`	`mod m0000010_init;`
`6`	`7`	`mod m0000020_add_sbom_group;`
`@@ -60,7 +61,9 @@ impl MigratorTrait for Migrator {`
`60`	`61`	`Box::new(m0001110_sbom_node_checksum_indexes::Migration),`
`61`	`62`	`Box::new(m0001120_sbom_external_node_indexes::Migration),`
`62`	`63`	`Box::new(m0001130_gover_cmp::Migration),`
`63`		`- Box::new(m0001140_example_data_migration::Migration),`
	`64`	`+ Box::new(MigrationWithData::new(`
	`65`	`+ m0001140_example_data_migration::Migration,`
	`66`	`+ )),`
`64`	`67`	`]`
`65`	`68`	`}`
`66`	`69`	`}`