fix: use a dedicated DB pool only for out-of-band data migrations

In case of in-band migrations, the sea orm migration manager will
execute all migrations inside a single transaction. This leads to the
problem, that concurrent operations are not possible. However, there
can be the strategy to run data migrations upfront, using concurrency.
And then re-run the migrations, but skipping the data migration.

Author: ctron

Cargo.lock

@@ -96,6 +96,7 @@ mime = "0.3.17"
 moka = "0.12.10"
 native-tls = "0.2"
 num-traits = "0.2"
+num_cpus = "1"
 oci-client = "0.16.0"
 openid = "0.22.0"
 openssl = "0.10"

Cargo.toml

@@ -28,6 +28,7 @@ futures = { workspace = true }
 futures-util = { workspace = true }
 humantime = { workspace = true }
 indicatif = { workspace = true, features = ["tokio", "futures"] }
+num_cpus = { workspace = true }
 osv = { workspace = true, features = ["schema"] }
 sea-orm = { workspace = true }
 sea-orm-migration = { workspace = true, features = ["runtime-tokio-rustls", "sqlx-postgres", "with-uuid"] }

migration/Cargo.toml

@@ -4,7 +4,7 @@ use crate::{
 };
 use clap::Parser;
 use futures::executor::block_on;
-use sea_orm::DbErr;
+use sea_orm::{DatabaseConnection, DbErr};
 use sea_orm_migration::{MigrationName, MigrationTrait, SchemaManager};
 use std::{ffi::OsString, ops::Deref, sync::LazyLock};
 use tokio::task_local;
@@ -79,18 +79,10 @@ impl MigrationWithData {
     }
 }
 
-impl<M> From<M> for MigrationWithData
-where
-    M: MigrationTraitWithData + 'static,
-{
-    fn from(value: M) -> Self {
-        MigrationWithData::new(Box::new(value))
-    }
-}
-
 /// A [`SchemaManager`], extended with data migration features.
 pub struct SchemaDataManager<'c> {
     pub manager: &'c SchemaManager<'c>,
+    pub db: Option<&'c DatabaseConnection>,
     storage: &'c DispatchBackend,
     options: &'c Options,
 }
@@ -106,11 +98,13 @@ impl<'a> Deref for SchemaDataManager<'a> {
 impl<'c> SchemaDataManager<'c> {
     pub fn new(
         manager: &'c SchemaManager<'c>,
+        db: Option<&'c DatabaseConnection>,
         storage: &'c DispatchBackend,
         options: &'c Options,
     ) -> Self {
         Self {
             manager,
+            db,
             storage,
             options,
         }
@@ -126,7 +120,18 @@ impl<'c> SchemaDataManager<'c> {
             return Ok(());
         }
 
-        self.manager.process(self.storage, self.options, f).await
+        match self.db {
+            Some(db) => {
+                self.manager
+                    .process(db, self.storage, self.options, f)
+                    .await
+            }
+            None => {
+                self.manager
+                    .process(self.manager.get_connection(), self.storage, self.options, f)
+                    .await
+            }
+        }
     }
 }
 
@@ -141,7 +146,7 @@ impl MigrationTrait for MigrationWithData {
     async fn up(&self, manager: &SchemaManager) -> Result<(), DbErr> {
         MigrationTraitWithData::up(
             &*self.migration,
-            &SchemaDataManager::new(manager, &self.storage, &self.options),
+            &SchemaDataManager::new(manager, None, &self.storage, &self.options),
         )
         .await
         .inspect_err(|err| tracing::warn!("Migration failed: {err}"))
@@ -150,7 +155,7 @@ impl MigrationTrait for MigrationWithData {
     async fn down(&self, manager: &SchemaManager) -> Result<(), DbErr> {
         MigrationTraitWithData::down(
             &*self.migration,
-            &SchemaDataManager::new(manager, &self.storage, &self.options),
+            &SchemaDataManager::new(manager, None, &self.storage, &self.options),
         )
         .await
     }

migration/src/data/migration.rs

@@ -13,12 +13,10 @@ use futures_util::{
     stream::{self, TryStreamExt},
 };
 use indicatif::{ProgressBar, ProgressStyle};
-use sea_orm::{DatabaseTransaction, DbErr, TransactionTrait};
+use sea_orm::{ConnectionTrait, DatabaseTransaction, DbErr, TransactionTrait};
 use sea_orm_migration::{MigrationTrait, SchemaManager};
-use std::{
-    num::{NonZeroU64, NonZeroUsize},
-    sync::Arc,
-};
+use std::{num::NonZeroU64, sync::Arc};
+use tracing::log;
 use trustify_module_storage::service::dispatch::DispatchBackend;
 
 /// A handler for processing a [`Document`] data migration.
@@ -43,8 +41,10 @@ where
 #[derive(Clone, Debug, PartialEq, Eq, clap::Parser)]
 pub struct Options {
     /// Number of concurrent documents being processes
-    #[arg(long, env = "MIGRATION_DATA_CONCURRENT", default_value = "5")]
-    pub concurrent: NonZeroUsize,
+    ///
+    /// If the value is zero, use the number of logical CPUs
+    #[arg(long, env = "MIGRATION_DATA_CONCURRENT", default_value = "0")]
+    pub concurrent: usize,
 
     /// The instance number of the current runner (zero based)
     #[arg(long, env = "MIGRATION_DATA_CURRENT_RUNNER", default_value = "0")]
@@ -70,7 +70,7 @@ pub struct Options {
 impl Default for Options {
     fn default() -> Self {
         Self {
-            concurrent: unsafe { NonZeroUsize::new_unchecked(5) },
+            concurrent: 5,
             current: 0,
             total: unsafe { NonZeroU64::new_unchecked(1) },
             skip_all: false,
@@ -118,6 +118,7 @@ impl From<&Options> for Partition {
 pub trait DocumentProcessor {
     fn process<D>(
         &self,
+        db: &(impl ConnectionTrait + TransactionTrait),
         storage: &DispatchBackend,
         options: &Options,
         f: impl Handler<D>,
@@ -165,6 +166,7 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
     /// actual system is still running from the read-only clone of the original data.
     async fn process<D>(
         &self,
+        db: &(impl ConnectionTrait + TransactionTrait),
         storage: &DispatchBackend,
         options: &Options,
         f: impl Handler<D>,
@@ -173,7 +175,6 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
         D: Document,
     {
         let partition: Partition = options.into();
-        let db = self.get_connection();
 
         let tx = db.begin().await?;
         let all: Vec<_> = D::all(&tx)
@@ -189,12 +190,21 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
             ProgressStyle::with_template(
                 "{spinner:.green} [{elapsed_precise}] [{wide_bar:.cyan/blue}] {human_pos}/{human_len} ({per_sec}) ({eta})",
             )
-            .map_err(|err| DbErr::Migration(err.to_string()))?
-            .progress_chars("##-"),
+                .map_err(|err| DbErr::Migration(err.to_string()))?
+                .progress_chars("#>-"),
         );
 
         let pb = Some(pb);
 
+        // get concurrency value
+        let mut concurrent = options.concurrent;
+        if concurrent == 0 {
+            // if zero, use number of logical CPUs
+            concurrent = num_cpus::get();
+        }
+
+        log::info!("Running {concurrent} parallel operations");
+
         stream::iter(all)
             .map(async |model| {
                 let tx = db.begin().await?;
@@ -220,7 +230,7 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
 
                 Ok::<_, DbErr>(())
             })
-            .buffer_unordered(options.concurrent.into())
+            .buffer_unordered(concurrent)
             .try_collect::<Vec<_>>()
             .await?;
 

migration/src/data/mod.rs

@@ -1,6 +1,6 @@
 use crate::data::{MigratorWithData, Options, SchemaDataManager};
 use anyhow::bail;
-use sea_orm::ConnectOptions;
+use sea_orm::{ConnectOptions, DatabaseConnection, DbErr};
 use sea_orm_migration::{IntoSchemaManagerConnection, SchemaManager};
 use std::{collections::HashMap, time::SystemTime};
 use trustify_module_storage::service::dispatch::DispatchBackend;
@@ -20,9 +20,39 @@ pub struct Runner {
     pub options: Options,
 }
 
+#[derive(Clone)]
 pub enum Database {
     Config { url: String, schema: Option<String> },
-    Provided(sea_orm::DatabaseConnection),
+    Provided(DatabaseConnection),
+}
+
+impl From<DatabaseConnection> for Database {
+    fn from(value: DatabaseConnection) -> Self {
+        Self::Provided(value)
+    }
+}
+
+impl From<trustify_common::db::Database> for Database {
+    fn from(value: trustify_common::db::Database) -> Self {
+        Self::Provided(value.into_connection())
+    }
+}
+
+impl Database {
+    pub async fn try_into_connection(self) -> Result<DatabaseConnection, DbErr> {
+        Ok(match self {
+            Self::Config { url, schema } => {
+                let schema = schema.clone().unwrap_or_else(|| "public".to_owned());
+
+                let connect_options = ConnectOptions::new(url)
+                    .set_schema_search_path(schema)
+                    .to_owned();
+
+                sea_orm::Database::connect(connect_options).await?
+            }
+            Self::Provided(database) => database.clone(),
+        })
+    }
 }
 
 impl Runner {
@@ -41,21 +71,11 @@ impl Runner {
             running.push(migration);
         }
 
-        let database = match self.database {
-            Database::Config { url, schema } => {
-                let schema = schema.unwrap_or_else(|| "public".to_owned());
-
-                let connect_options = ConnectOptions::new(url)
-                    .set_schema_search_path(schema)
-                    .to_owned();
-
-                sea_orm::Database::connect(connect_options).await?
-            }
-            Database::Provided(database) => database,
-        };
+        let database = self.database.clone().try_into_connection().await?;
 
         let manager = SchemaManager::new(database.into_schema_manager_connection());
-        let manager = SchemaDataManager::new(&manager, &self.storage, &self.options);
+        let manager =
+            SchemaDataManager::new(&manager, Some(&database), &self.storage, &self.options);
 
         for run in running {
             tracing::info!(name = run.name(), "Running data migration");

migration/src/data/run.rs

@@ -4,9 +4,9 @@ use migration::{
     data::{Database, Direction, MigrationWithData, Options, Runner},
 };
 use sea_orm_migration::MigratorTrait;
-use std::num::NonZeroUsize;
 use test_context::test_context;
 use test_log::test;
+use tracing::log;
 use trustify_test_context::{TrustifyMigrationContext, commit, dump};
 
 commit!(Commit("6d3ea814b4b44fe16ea8f21724dda5abb0fc7932"));
@@ -66,13 +66,36 @@ dump!(
     ignore = "enable with: cargo test --features long_running"
 )]
 async fn performance(ctx: &TrustifyMigrationContext<Ds4>) -> Result<(), anyhow::Error> {
+    let migrations = vec![
+        "m0002000_add_sbom_properties".into(),
+        "m0002010_add_advisory_scores".into(),
+    ];
+
+    // we simulate running the migrations out-of-band
+
+    log::info!("Running data migrations out-of-band");
+
+    Runner {
+        database: ctx.db.clone().into(),
+        storage: ctx.storage.clone().into(),
+        direction: Default::default(),
+        migrations: migrations.clone(),
+        options: Default::default(),
+    }
+    .run::<Migrator>()
+    .await?;
+
+    // now run the standard migration, skipping the out-of-band ones
+
+    log::info!("Running migrations");
+
     MigrationWithData::run_with_test(
         ctx.storage.clone(),
         Options {
-            concurrent: NonZeroUsize::new(32).unwrap(),
+            skip: migrations,
             ..Options::default()
         },
-        async { MigratorTest::up(&ctx.db, None).await },
+        async { Migrator::up(&ctx.db, None).await },
     )
     .await?;