Add a create_many_certificates on the repository to allow bulk insert

This speed up certificates insertion on a db with 1800+ certificates by
a factor of 5+ (8s instead of 44s).

Author: Alenar

mithril-aggregator/src/database/provider/certificate.rs

@@ -1,6 +1,7 @@
 use async_trait::async_trait;
 use chrono::{DateTime, Utc};
 use sqlite::{Connection, Value};
+use std::iter::repeat;
 use std::sync::Arc;
 use tokio::sync::Mutex;
 
@@ -378,29 +379,52 @@ impl<'conn> InsertCertificateRecordProvider<'conn> {
     }
 
     fn get_insert_condition(&self, certificate_record: &CertificateRecord) -> WhereCondition {
+        self.get_insert_many_condition(&vec![certificate_record.clone()])
+    }
+
+    fn get_insert_many_condition(
+        &self,
+        certificates_records: &[CertificateRecord],
+    ) -> WhereCondition {
+        let columns = "(certificate_id, parent_certificate_id, message, signature, \
+aggregate_verification_key, epoch, beacon, protocol_version, protocol_parameters, \
+protocol_message, signers, initiated_at, sealed_at)";
+        let values_columns: Vec<&str> =
+            repeat("(?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*)")
+                .take(certificates_records.len())
+                .collect();
+
+        let values: Vec<Value> = certificates_records
+            .iter()
+            .flat_map(|certificate_record| {
+                vec![
+                    Value::String(certificate_record.certificate_id.to_owned()),
+                    match certificate_record.parent_certificate_id.to_owned() {
+                        Some(parent_certificate_id) => Value::String(parent_certificate_id),
+                        None => Value::Null,
+                    },
+                    Value::String(certificate_record.message.to_owned()),
+                    Value::String(certificate_record.signature.to_owned()),
+                    Value::String(certificate_record.aggregate_verification_key.to_owned()),
+                    Value::Integer(i64::try_from(certificate_record.epoch.0).unwrap()),
+                    Value::String(serde_json::to_string(&certificate_record.beacon).unwrap()),
+                    Value::String(certificate_record.protocol_version.to_owned()),
+                    Value::String(
+                        serde_json::to_string(&certificate_record.protocol_parameters).unwrap(),
+                    ),
+                    Value::String(
+                        serde_json::to_string(&certificate_record.protocol_message).unwrap(),
+                    ),
+                    Value::String(serde_json::to_string(&certificate_record.signers).unwrap()),
+                    Value::String(certificate_record.initiated_at.to_rfc3339()),
+                    Value::String(certificate_record.sealed_at.to_rfc3339()),
+                ]
+            })
+            .collect();
+
         WhereCondition::new(
-            "(certificate_id, parent_certificate_id, message, signature, aggregate_verification_key, epoch, beacon, protocol_version, protocol_parameters, protocol_message, signers, initiated_at, sealed_at) values (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13)",
-            vec![
-                Value::String(certificate_record.certificate_id.to_owned()),
-                if let Some(parent_certificate_id) = certificate_record.parent_certificate_id.to_owned() {
-                    Value::String(parent_certificate_id)
-                }else{
-                    Value::Null
-                },
-                Value::String(certificate_record.message.to_owned()),
-                Value::String(certificate_record.signature.to_owned()),
-                Value::String(certificate_record.aggregate_verification_key.to_owned()),
-                Value::Integer(i64::try_from(certificate_record.epoch.0).unwrap()),
-                Value::String(serde_json::to_string(&certificate_record.beacon).unwrap()),
-                Value::String(certificate_record.protocol_version.to_owned()),
-                Value::String(
-                    serde_json::to_string(&certificate_record.protocol_parameters).unwrap(),
-                ),
-                Value::String(serde_json::to_string(&certificate_record.protocol_message).unwrap()),
-                Value::String(serde_json::to_string(&certificate_record.signers).unwrap()),
-                Value::String(certificate_record.initiated_at.to_rfc3339()),
-                Value::String(certificate_record.sealed_at.to_rfc3339()),
-            ],
+            format!("{columns} values {}", values_columns.join(", ")).as_str(),
+            values,
         )
     }
 
@@ -418,6 +442,15 @@ impl<'conn> InsertCertificateRecordProvider<'conn> {
 
         Ok(entity)
     }
+
+    fn persist_many(
+        &self,
+        certificates_records: Vec<CertificateRecord>,
+    ) -> Result<Vec<CertificateRecord>, StdError> {
+        let filters = self.get_insert_many_condition(&certificates_records);
+
+        Ok(self.find(filters)?.collect())
+    }
 }
 
 impl<'conn> Provider<'conn> for InsertCertificateRecordProvider<'conn> {
@@ -587,6 +620,23 @@ impl CertificateRepository {
         Ok(new_certificate.into())
     }
 
+    /// Create many certificates at once in the database.
+    pub async fn create_many_certificates(
+        &self,
+        certificates: Vec<Certificate>,
+    ) -> StdResult<Vec<Certificate>> {
+        let lock = self.connection.lock().await;
+        let provider = InsertCertificateRecordProvider::new(&lock);
+        let records: Vec<CertificateRecord> =
+            certificates.into_iter().map(|cert| cert.into()).collect();
+        let new_certificates = provider.persist_many(records)?;
+
+        Ok(new_certificates
+            .into_iter()
+            .map(|cert| cert.into())
+            .collect::<Vec<_>>())
+    }
+
     /// Delete all the given certificates from the database
     pub async fn delete_certificates(&self, certificates: &[&Certificate]) -> StdResult<()> {
         let ids = certificates
@@ -856,7 +906,7 @@ mod tests {
     }
 
     #[test]
-    fn insert_certificate_record() {
+    fn insert_certificate_condition() {
         let (certificates, _) = setup_certificate_chain(2, 1);
         let certificate_record: CertificateRecord = certificates.first().unwrap().to_owned().into();
         let connection = Connection::open(":memory:").unwrap();
@@ -890,6 +940,57 @@ mod tests {
         );
     }
 
+    #[test]
+    fn insert_many_certificates_condition() {
+        let (certificates, _) = setup_certificate_chain(2, 1);
+        let certificates_records: Vec<CertificateRecord> =
+            certificates.into_iter().map(|c| c.into()).collect();
+        let connection = Connection::open(":memory:").unwrap();
+        let provider = InsertCertificateRecordProvider::new(&connection);
+        let condition = provider.get_insert_many_condition(&certificates_records);
+        let (values, params) = condition.expand();
+
+        assert_eq!(
+            "(certificate_id, parent_certificate_id, message, signature, \
+aggregate_verification_key, epoch, beacon, protocol_version, protocol_parameters, \
+protocol_message, signers, initiated_at, sealed_at) values \
+(?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13), \
+(?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22, ?23, ?24, ?25, ?26)"
+                .to_string(),
+            values
+        );
+        assert_eq!(
+            certificates_records
+                .into_iter()
+                .flat_map(|certificate_record| {
+                    vec![
+                        Value::String(certificate_record.certificate_id),
+                        match certificate_record.parent_certificate_id {
+                            Some(id) => Value::String(id),
+                            None => Value::Null,
+                        },
+                        Value::String(certificate_record.message),
+                        Value::String(certificate_record.signature),
+                        Value::String(certificate_record.aggregate_verification_key),
+                        Value::Integer(i64::try_from(certificate_record.epoch.0).unwrap()),
+                        Value::String(serde_json::to_string(&certificate_record.beacon).unwrap()),
+                        Value::String(certificate_record.protocol_version),
+                        Value::String(
+                            serde_json::to_string(&certificate_record.protocol_parameters).unwrap(),
+                        ),
+                        Value::String(
+                            serde_json::to_string(&certificate_record.protocol_message).unwrap(),
+                        ),
+                        Value::String(serde_json::to_string(&certificate_record.signers).unwrap()),
+                        Value::String(certificate_record.initiated_at.to_rfc3339()),
+                        Value::String(certificate_record.sealed_at.to_rfc3339()),
+                    ]
+                })
+                .collect::<Vec<_>>(),
+            params
+        );
+    }
+
     #[test]
     fn test_get_certificate_records() {
         let (certificates, _) = setup_certificate_chain(20, 7);
@@ -945,6 +1046,23 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_insert_many_certificates_records() {
+        let (certificates, _) = setup_certificate_chain(5, 2);
+        let certificates_records: Vec<CertificateRecord> =
+            certificates.into_iter().map(|cert| cert.into()).collect();
+
+        let connection = Connection::open(":memory:").unwrap();
+        setup_certificate_db(&connection, Vec::new()).unwrap();
+
+        let provider = InsertCertificateRecordProvider::new(&connection);
+        let certificates_records_saved = provider
+            .persist_many(certificates_records.clone())
+            .expect("saving many records should not fail");
+
+        assert_eq!(certificates_records, certificates_records_saved);
+    }
+
     #[tokio::test]
     async fn test_store_adapter() {
         let (certificates, _) = setup_certificate_chain(5, 2);

mithril-aggregator/src/tools/certificates_hash_migrator.rs

@@ -107,17 +107,9 @@ impl CertificatesHashMigrator {
 
         // 2 - Certificates migrated, we can insert them in the db
         debug!("🔧 Certificate Hash Migrator: inserting migrated certificates in the database");
-        for migrated_certificate in migrated_certificates {
-            trace!(
-                "🔧 Certificate Hash Migrator: inserting migrated certificate {:?}",
-                migrated_certificate.beacon;
-                "hash" => &migrated_certificate.hash,
-                "previous_hash" => &migrated_certificate.previous_hash
-            );
-            self.certificate_repository
-                .create_certificate(migrated_certificate)
-                .await?;
-        }
+        self.certificate_repository
+            .create_many_certificates(migrated_certificates)
+            .await?;
 
         Ok((old_certificates, old_and_new_hashes))
     }