feat: allow running data migrations as part of migrations

ctron · ctron · commit ef04e5da0504 · 2025-09-10T13:03:08.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,107 @@
+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::task::block_in_place(|| {
+        tokio::runtime::Handle::current().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/server/src/profile/api.rs b/server/src/profile/api.rs
@@ -6,11 +6,10 @@ use actix_web::middleware;
 
 use crate::{endpoints, profile::spawn_db_check, sample_data};
 use actix_web::web;
-use anyhow::Context;
 use bytesize::ByteSize;
 use futures::FutureExt;
 use humantime::parse_duration;
-use std::{env, fs::create_dir_all, path::PathBuf, process::ExitCode, sync::Arc};
+use std::{env, process::ExitCode, sync::Arc};
 use tokio_schedule::{Job, every};
 use trustify_auth::{
     auth::AuthConfigArguments,
@@ -32,10 +31,7 @@ use trustify_infrastructure::{
 use trustify_module_analysis::{config::AnalysisConfig, service::AnalysisService};
 use trustify_module_fundamental::purl::service::PurlService;
 use trustify_module_ingestor::graph::Graph;
-use trustify_module_storage::{
-    config::{StorageConfig, StorageStrategy},
-    service::{dispatch::DispatchBackend, fs::FileSystemBackend, s3::S3Backend},
-};
+use trustify_module_storage::{config::StorageConfig, service::dispatch::DispatchBackend};
 use trustify_module_ui::{UI, endpoints::UiResources};
 use utoipa::openapi::{Info, License};
 
@@ -337,7 +333,7 @@ impl InitData {
                     let _ = embedded_oidc.0.await?;
                     Ok::<_, anyhow::Error>(())
                 }
-                    .boxed_local(),
+                .boxed_local(),
             );
         }
 
@@ -406,11 +402,11 @@ pub(crate) struct Config {
 pub(crate) fn configure(svc: &mut utoipa_actix_web::service_config::ServiceConfig, config: Config) {
     let Config {
         config:
-        ModuleConfig {
-            ingestor,
-            fundamental,
-            ui,
-        },
+            ModuleConfig {
+                ingestor,
+                fundamental,
+                ui,
+            },
         db,
         storage,
         auth,
@@ -437,7 +433,7 @@ pub(crate) fn configure(svc: &mut utoipa_actix_web::service_config::ServiceConfi
                     svc.wrap(middleware::NormalizePath::new(
                         middleware::TrailingSlash::Always,
                     ))
-                        .wrap(new_auth(auth.clone()))
+                    .wrap(new_auth(auth.clone()))
                 })
                 .configure(|svc| {
                     trustify_module_graphql::endpoints::configure(svc, db.clone());
@@ -559,7 +555,7 @@ mod test {
                 .apply_openapi(None, None)
                 .configure(|svc| post_configure(svc, PostConfig { ui })),
         )
-            .await;
+        .await;
 
         // main UI
 
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`	`}`