feat(aggregator): implement persistence of synchronised certificates

Author: Alenar

mithril-aggregator/src/database/query/certificate/conditions.rs

@@ -0,0 +1,65 @@
+//! Shared `WhereCondition` across certificates queries
+
+use sqlite::Value;
+use std::iter::repeat_n;
+
+use mithril_persistence::sqlite::WhereCondition;
+
+use crate::database::record::CertificateRecord;
+
+pub(super) fn insert_many(certificates_records: Vec<CertificateRecord>) -> WhereCondition {
+    let columns = "(\
+        certificate_id, \
+        parent_certificate_id, \
+        message, \
+        signature, \
+        aggregate_verification_key, \
+        epoch, \
+        network, \
+        signed_entity_type_id, \
+        signed_entity_beacon, \
+        protocol_version, \
+        protocol_parameters, \
+        protocol_message, \
+        signers, \
+        initiated_at, \
+        sealed_at)";
+    let values_columns: Vec<&str> = repeat_n(
+        "(?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*)",
+        certificates_records.len(),
+    )
+    .collect();
+
+    let values: Vec<Value> = certificates_records
+        .into_iter()
+        .flat_map(|certificate_record| {
+            vec![
+                Value::String(certificate_record.certificate_id),
+                match certificate_record.parent_certificate_id {
+                    Some(parent_certificate_id) => Value::String(parent_certificate_id),
+                    None => Value::Null,
+                },
+                Value::String(certificate_record.message),
+                Value::String(certificate_record.signature),
+                Value::String(certificate_record.aggregate_verification_key),
+                Value::Integer(certificate_record.epoch.try_into().unwrap()),
+                Value::String(certificate_record.network),
+                Value::Integer(certificate_record.signed_entity_type.index() as i64),
+                Value::String(certificate_record.signed_entity_type.get_json_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();
+
+    WhereCondition::new(
+        format!("{columns} values {}", values_columns.join(", ")).as_str(),
+        values,
+    )
+}

mithril-aggregator/src/database/query/certificate/insert_certificate.rs

@@ -1,11 +1,9 @@
-use std::iter::repeat_n;
-
-use sqlite::Value;
-
 use mithril_persistence::sqlite::{Query, SourceAlias, SqLiteEntity, WhereCondition};
 
 use crate::database::record::CertificateRecord;
 
+use super::conditions;
+
 /// Query to insert [CertificateRecord] in the sqlite database
 pub struct InsertCertificateRecordQuery {
     condition: WhereCondition,
@@ -17,64 +15,9 @@ impl InsertCertificateRecordQuery {
     }
 
     pub fn many(certificates_records: Vec<CertificateRecord>) -> Self {
-        let columns = "(\
-        certificate_id, \
-        parent_certificate_id, \
-        message, \
-        signature, \
-        aggregate_verification_key, \
-        epoch, \
-        network, \
-        signed_entity_type_id, \
-        signed_entity_beacon, \
-        protocol_version, \
-        protocol_parameters, \
-        protocol_message, \
-        signers, \
-        initiated_at, \
-        sealed_at)";
-        let values_columns: Vec<&str> = repeat_n(
-            "(?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*)",
-            certificates_records.len(),
-        )
-        .collect();
-
-        let values: Vec<Value> = certificates_records
-            .into_iter()
-            .flat_map(|certificate_record| {
-                vec![
-                    Value::String(certificate_record.certificate_id),
-                    match certificate_record.parent_certificate_id {
-                        Some(parent_certificate_id) => Value::String(parent_certificate_id),
-                        None => Value::Null,
-                    },
-                    Value::String(certificate_record.message),
-                    Value::String(certificate_record.signature),
-                    Value::String(certificate_record.aggregate_verification_key),
-                    Value::Integer(certificate_record.epoch.try_into().unwrap()),
-                    Value::String(certificate_record.network),
-                    Value::Integer(certificate_record.signed_entity_type.index() as i64),
-                    Value::String(certificate_record.signed_entity_type.get_json_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();
-
-        let condition = WhereCondition::new(
-            format!("{columns} values {}", values_columns.join(", ")).as_str(),
-            values,
-        );
-
-        Self { condition }
+        Self {
+            condition: conditions::insert_many(certificates_records),
+        }
     }
 }
 

mithril-aggregator/src/database/query/certificate/insert_or_replace_certificate.rs

@@ -0,0 +1,169 @@
+use mithril_persistence::sqlite::{Query, SourceAlias, SqLiteEntity, WhereCondition};
+
+use crate::database::record::CertificateRecord;
+
+use super::conditions;
+
+/// Query to insert or replace [CertificateRecord] in the sqlite database
+pub struct InsertOrReplaceCertificateRecordQuery {
+    condition: WhereCondition,
+}
+
+impl InsertOrReplaceCertificateRecordQuery {
+    pub fn many(certificates_records: Vec<CertificateRecord>) -> Self {
+        Self {
+            condition: conditions::insert_many(certificates_records),
+        }
+    }
+}
+
+impl Query for InsertOrReplaceCertificateRecordQuery {
+    type Entity = CertificateRecord;
+
+    fn filters(&self) -> WhereCondition {
+        self.condition.clone()
+    }
+
+    fn get_definition(&self, condition: &str) -> String {
+        // it is important to alias the fields with the same name as the table
+        // since the table cannot be aliased in a RETURNING statement in SQLite.
+        let projection = Self::Entity::get_projection()
+            .expand(SourceAlias::new(&[("{:certificate:}", "certificate")]));
+
+        format!("insert or replace into certificate {condition} returning {projection}")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+
+    use mithril_common::crypto_helper::tests_setup::setup_certificate_chain;
+    use mithril_common::entities::Epoch;
+    use mithril_common::test_utils::fake_data;
+    use mithril_persistence::sqlite::ConnectionExtensions;
+
+    use crate::database::query::{GetCertificateRecordQuery, InsertCertificateRecordQuery};
+    use crate::database::test_helper::main_db_connection;
+
+    use super::*;
+
+    #[test]
+    fn test_insert_many_certificates_records_in_empty_db() {
+        let certificates = setup_certificate_chain(5, 2);
+        let certificates_records: Vec<CertificateRecord> = certificates.into();
+
+        let connection = main_db_connection().unwrap();
+
+        let certificates_records_saved: Vec<CertificateRecord> = connection
+            .fetch_collect(InsertOrReplaceCertificateRecordQuery::many(
+                certificates_records.clone(),
+            ))
+            .expect("saving many records should not fail");
+
+        assert_eq!(certificates_records, certificates_records_saved);
+
+        // Check insertion order
+        let all_records: Vec<CertificateRecord> =
+            connection.fetch_collect(GetCertificateRecordQuery::all()).unwrap();
+        assert_eq!(
+            certificates_records.into_iter().rev().collect::<Vec<_>>(),
+            all_records
+        );
+    }
+
+    #[test]
+    fn test_replace_one_certificate_record() {
+        let certificate_record = CertificateRecord {
+            epoch: Epoch(12),
+            ..fake_data::certificate("hash").into()
+        };
+
+        let connection = main_db_connection().unwrap();
+        let certificate_record_saved = connection
+            .fetch_first(InsertCertificateRecordQuery::one(
+                certificate_record.clone(),
+            ))
+            .unwrap();
+        assert_eq!(Some(Epoch(12)), certificate_record_saved.map(|r| r.epoch));
+
+        let modified_certificate_record = CertificateRecord {
+            epoch: Epoch(23),
+            ..certificate_record
+        };
+        let certificate_record_saved = connection
+            .fetch_first(InsertOrReplaceCertificateRecordQuery::many(vec![
+                modified_certificate_record.clone(),
+            ]))
+            .unwrap();
+        assert_eq!(Some(Epoch(23)), certificate_record_saved.map(|r| r.epoch));
+
+        let all_records_cursor = connection.fetch(GetCertificateRecordQuery::all()).unwrap();
+        assert_eq!(1, all_records_cursor.count());
+    }
+
+    #[test]
+    fn test_insert_and_replace_many_certificate_record() {
+        let tested_records: HashMap<_, CertificateRecord> = HashMap::from([
+            (
+                "cert1-genesis",
+                fake_data::genesis_certificate("genesis").into(),
+            ),
+            ("cert2", fake_data::certificate("cert2").into()),
+            (
+                "cert2-modified",
+                CertificateRecord {
+                    epoch: Epoch(14),
+                    ..fake_data::certificate("cert2").into()
+                },
+            ),
+            ("cert3", fake_data::certificate("cert3").into()),
+            ("cert4", fake_data::certificate("cert4").into()),
+            (
+                "cert4-modified",
+                CertificateRecord {
+                    epoch: Epoch(32),
+                    ..fake_data::certificate("cert4").into()
+                },
+            ),
+            ("cert5", fake_data::certificate("cert5").into()),
+        ]);
+        let connection = main_db_connection().unwrap();
+
+        let cursor = connection
+            .fetch(InsertCertificateRecordQuery::many(vec![
+                tested_records["cert1-genesis"].clone(),
+                tested_records["cert2"].clone(),
+                tested_records["cert3"].clone(),
+                tested_records["cert4"].clone(),
+                tested_records["cert5"].clone(),
+            ]))
+            .unwrap();
+        assert_eq!(5, cursor.count());
+
+        let cursor = connection
+            .fetch(InsertOrReplaceCertificateRecordQuery::many(vec![
+                tested_records["cert1-genesis"].clone(),
+                tested_records["cert2-modified"].clone(),
+                tested_records["cert3"].clone(),
+                tested_records["cert4-modified"].clone(),
+            ]))
+            .unwrap();
+        assert_eq!(4, cursor.count());
+
+        let all_records: Vec<CertificateRecord> =
+            connection.fetch_collect(GetCertificateRecordQuery::all()).unwrap();
+        assert_eq!(5, all_records.len());
+        assert_eq!(
+            all_records,
+            vec![
+                tested_records["cert4-modified"].clone(),
+                tested_records["cert3"].clone(),
+                tested_records["cert2-modified"].clone(),
+                tested_records["cert1-genesis"].clone(),
+                // Since the cert5 was not in the Insert/replace query, it now has a lower rowid and shows first
+                tested_records["cert5"].clone(),
+            ]
+        );
+    }
+}

mithril-aggregator/src/database/query/certificate/mod.rs

@@ -1,9 +1,12 @@
+mod conditions;
 mod delete_certificate;
 mod get_certificate;
 mod get_master_certificate;
 mod insert_certificate;
+mod insert_or_replace_certificate;
 
 pub use delete_certificate::*;
 pub use get_certificate::*;
 pub use get_master_certificate::*;
 pub use insert_certificate::*;
+pub use insert_or_replace_certificate::*;

mithril-aggregator/src/database/repository/certificate_repository.rs

@@ -11,9 +11,10 @@ use mithril_persistence::sqlite::ConnectionExtensions;
 
 use crate::database::query::{
     DeleteCertificateQuery, GetCertificateRecordQuery, InsertCertificateRecordQuery,
-    MasterCertificateQuery,
+    InsertOrReplaceCertificateRecordQuery, MasterCertificateQuery,
 };
 use crate::database::record::CertificateRecord;
+use crate::services::SynchronizedCertificateStorer;
 
 /// Database frontend API for Certificate queries.
 pub struct CertificateRepository {
@@ -105,6 +106,24 @@ impl CertificateRepository {
         Ok(new_certificates.map(|cert| cert.into()).collect())
     }
 
+    /// Create, or replace if they already exist, many certificates at once in the database.
+    pub async fn create_or_replace_many_certificates(
+        &self,
+        certificates: Vec<Certificate>,
+    ) -> StdResult<Vec<Certificate>> {
+        if certificates.is_empty() {
+            return Ok(vec![]);
+        }
+
+        let records: Vec<CertificateRecord> =
+            certificates.into_iter().map(|cert| cert.into()).collect();
+        let new_certificates = self
+            .connection
+            .fetch(InsertOrReplaceCertificateRecordQuery::many(records))?;
+
+        Ok(new_certificates.map(|cert| cert.into()).collect())
+    }
+
     /// Delete all the given certificates from the database
     pub async fn delete_certificates(&self, certificates: &[&Certificate]) -> StdResult<()> {
         let ids = certificates.iter().map(|c| c.hash.as_str()).collect::<Vec<_>>();
@@ -131,6 +150,18 @@ impl CertificateRetriever for CertificateRepository {
     }
 }
 
+#[async_trait]
+impl SynchronizedCertificateStorer for CertificateRepository {
+    async fn insert_or_replace_many(&self, certificates: Vec<Certificate>) -> StdResult<()> {
+        self.create_or_replace_many_certificates(certificates).await?;
+        Ok(())
+    }
+
+    async fn get_latest_genesis(&self) -> StdResult<Option<Certificate>> {
+        self.get_latest_genesis_certificate().await
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use mithril_common::crypto_helper::tests_setup::setup_certificate_chain;