feat(aggregator): implement persistence of openmessage created post sync

Author: Alenar

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

@@ -0,0 +1,32 @@
+//! Shared `WhereCondition` across open message queries
+
+use sqlite::Value;
+
+use mithril_common::StdResult;
+use mithril_persistence::sqlite::WhereCondition;
+
+use crate::database::record::OpenMessageRecord;
+
+pub(crate) fn insert_one(record: OpenMessageRecord) -> StdResult<WhereCondition> {
+    let expression = "(\
+        open_message_id, epoch_setting_id, beacon, signed_entity_type_id, protocol_message, \
+        expires_at, created_at, is_certified, is_expired\
+    ) values (?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*, ?*)";
+    let beacon_str = record.signed_entity_type.get_json_beacon()?;
+    let parameters = vec![
+        Value::String(record.open_message_id.to_string()),
+        Value::Integer(record.epoch.try_into()?),
+        Value::String(beacon_str),
+        Value::Integer(record.signed_entity_type.index() as i64),
+        Value::String(serde_json::to_string(&record.protocol_message)?),
+        record
+            .expires_at
+            .map(|t| Value::String(t.to_rfc3339()))
+            .unwrap_or(Value::Null),
+        Value::String(record.created_at.to_rfc3339()),
+        Value::Integer(record.is_certified as i64),
+        Value::Integer(record.is_expired as i64),
+    ];
+
+    Ok(WhereCondition::new(expression, parameters))
+}

mithril-aggregator/src/database/query/open_message/get_open_message.rs

@@ -15,6 +15,13 @@ pub struct GetOpenMessageQuery {
 }
 
 impl GetOpenMessageQuery {
+    #[cfg(test)]
+    pub fn all() -> Self {
+        Self {
+            condition: WhereCondition::default(),
+        }
+    }
+
     pub fn by_epoch_and_signed_entity_type(
         epoch: Epoch,
         signed_entity_type: &SignedEntityType,

mithril-aggregator/src/database/query/open_message/insert_open_message.rs

@@ -1,11 +1,10 @@
 use chrono::Utc;
-use sqlite::Value;
-use uuid::Uuid;
 
 use mithril_common::StdResult;
 use mithril_common::entities::{Epoch, ProtocolMessage, SignedEntityType};
 use mithril_persistence::sqlite::{Query, SourceAlias, SqLiteEntity, WhereCondition};
 
+use crate::database::query::open_message::conditions;
 use crate::database::record::OpenMessageRecord;
 
 /// Query to insert [OpenMessageRecord] in the sqlite database
@@ -19,23 +18,20 @@ impl InsertOpenMessageQuery {
         signed_entity_type: &SignedEntityType,
         protocol_message: &ProtocolMessage,
     ) -> StdResult<Self> {
-        let expression = "(open_message_id, epoch_setting_id, beacon, signed_entity_type_id, protocol_message, expires_at, created_at) values (?*, ?*, ?*, ?*, ?*, ?*, ?*)";
-        let beacon_str = signed_entity_type.get_json_beacon()?;
-        let parameters = vec![
-            Value::String(Uuid::new_v4().to_string()),
-            Value::Integer(epoch.try_into()?),
-            Value::String(beacon_str),
-            Value::Integer(signed_entity_type.index() as i64),
-            Value::String(serde_json::to_string(protocol_message)?),
-            signed_entity_type
-                .get_open_message_timeout()
-                .map(|t| Value::String((Utc::now() + t).to_rfc3339()))
-                .unwrap_or(Value::Null),
-            Value::String(Utc::now().to_rfc3339()),
-        ];
+        let now = Utc::now();
+        let record = OpenMessageRecord {
+            open_message_id: OpenMessageRecord::new_id(),
+            epoch,
+            signed_entity_type: signed_entity_type.clone(),
+            protocol_message: protocol_message.clone(),
+            is_certified: false,
+            is_expired: false,
+            created_at: now,
+            expires_at: signed_entity_type.get_open_message_timeout().map(|t| now + t),
+        };
 
         Ok(Self {
-            condition: WhereCondition::new(expression, parameters),
+            condition: conditions::insert_one(record)?,
         })
     }
 }
@@ -54,3 +50,69 @@ impl Query for InsertOpenMessageQuery {
         format!("insert into open_message {condition} returning {projection}")
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use mithril_common::entities::ProtocolMessagePartKey;
+    use mithril_persistence::sqlite::ConnectionExtensions;
+
+    use crate::database::query::GetOpenMessageQuery;
+    use crate::database::test_helper::main_db_connection;
+
+    use super::*;
+
+    #[test]
+    fn test_insert_one() {
+        let connection = main_db_connection().unwrap();
+        let epoch = Epoch(5);
+        let signed_entity_type = SignedEntityType::CardanoStakeDistribution(Epoch(10));
+        let mut protocol_message = ProtocolMessage::new();
+        protocol_message.set_message_part(
+            ProtocolMessagePartKey::CardanoStakeDistributionEpoch,
+            "value".to_string(),
+        );
+
+        connection
+            .fetch_first(
+                InsertOpenMessageQuery::one(epoch, &signed_entity_type, &protocol_message).unwrap(),
+            )
+            .unwrap();
+        let records: Vec<OpenMessageRecord> =
+            connection.fetch_collect(GetOpenMessageQuery::all()).unwrap();
+
+        assert_eq!(1, records.len());
+        assert_eq!(
+            OpenMessageRecord {
+                open_message_id: records[0].open_message_id,
+                epoch,
+                signed_entity_type,
+                protocol_message,
+                is_certified: false,
+                is_expired: false,
+                created_at: records[0].created_at,
+                expires_at: records[0].expires_at,
+            },
+            records[0]
+        );
+    }
+
+    #[should_panic]
+    #[test]
+    fn test_insert_two_entries_with_same_signed_entity_violate_unique_constraint() {
+        let connection = main_db_connection().unwrap();
+        let epoch = Epoch(5);
+        let signed_entity_type = SignedEntityType::MithrilStakeDistribution(Epoch(10));
+
+        connection
+            .fetch_first(
+                InsertOpenMessageQuery::one(epoch, &signed_entity_type, &ProtocolMessage::new())
+                    .unwrap(),
+            )
+            .unwrap();
+
+        let _ = connection.fetch_first(
+            InsertOpenMessageQuery::one(epoch + 10, &signed_entity_type, &ProtocolMessage::new())
+                .unwrap(),
+        );
+    }
+}

mithril-aggregator/src/database/query/open_message/insert_or_replace_open_message.rs

@@ -0,0 +1,92 @@
+use mithril_common::StdResult;
+use mithril_persistence::sqlite::{Query, SourceAlias, SqLiteEntity, WhereCondition};
+
+use crate::database::query::open_message::conditions;
+use crate::database::record::OpenMessageRecord;
+
+/// Query to insert [OpenMessageRecord] in the sqlite database
+pub struct InsertOrReplaceOpenMessageQuery {
+    condition: WhereCondition,
+}
+
+impl InsertOrReplaceOpenMessageQuery {
+    pub fn one(record: OpenMessageRecord) -> StdResult<Self> {
+        Ok(Self {
+            condition: conditions::insert_one(record)?,
+        })
+    }
+}
+
+impl Query for InsertOrReplaceOpenMessageQuery {
+    type Entity = OpenMessageRecord;
+
+    fn filters(&self) -> WhereCondition {
+        self.condition.clone()
+    }
+
+    fn get_definition(&self, condition: &str) -> String {
+        let aliases = SourceAlias::new(&[("{:open_message:}", "open_message")]);
+        let projection = Self::Entity::get_projection().expand(aliases);
+
+        format!("insert or replace into open_message {condition} returning {projection}")
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use mithril_persistence::sqlite::ConnectionExtensions;
+
+    use crate::database::query::GetOpenMessageQuery;
+    use crate::database::test_helper::main_db_connection;
+
+    use super::*;
+
+    #[test]
+    fn test_insert_one() {
+        let connection = main_db_connection().unwrap();
+        let record = OpenMessageRecord::dummy();
+
+        connection
+            .fetch_first(InsertOrReplaceOpenMessageQuery::one(record.clone()).unwrap())
+            .unwrap();
+        let records: Vec<OpenMessageRecord> = connection
+            .fetch_collect(
+                GetOpenMessageQuery::by_epoch_and_signed_entity_type(
+                    record.epoch,
+                    &record.signed_entity_type,
+                )
+                .unwrap(),
+            )
+            .unwrap();
+
+        assert_eq!(1, records.len());
+        assert_eq!(record, records[0]);
+    }
+
+    #[test]
+    fn test_insert_record_for_existing_signed_entity_type_replaces_it() {
+        let connection = main_db_connection().unwrap();
+        let record = OpenMessageRecord {
+            is_expired: false,
+            ..OpenMessageRecord::dummy()
+        };
+
+        connection
+            .fetch_first(InsertOrReplaceOpenMessageQuery::one(record.clone()).unwrap())
+            .unwrap();
+
+        let replaced_record = connection
+            .fetch_first(
+                InsertOrReplaceOpenMessageQuery::one(OpenMessageRecord {
+                    is_expired: true,
+                    ..record.clone()
+                })
+                .unwrap(),
+            )
+            .unwrap();
+        let count = connection.fetch(GetOpenMessageQuery::all()).unwrap().count();
+
+        assert_eq!(1, count);
+        assert_eq!(Some(true), replaced_record.map(|r| r.is_expired));
+    }
+}

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

@@ -1,11 +1,14 @@
+mod conditions;
 mod delete_open_message;
 mod get_open_message;
 mod get_open_message_with_single_signatures;
 mod insert_open_message;
+mod insert_or_replace_open_message;
 mod update_open_message;
 
 pub use delete_open_message::*;
 pub use get_open_message::*;
 pub use get_open_message_with_single_signatures::*;
 pub use insert_open_message::*;
+pub use insert_or_replace_open_message::*;
 pub use update_open_message::*;

mithril-aggregator/src/database/record/open_message.rs

@@ -39,6 +39,11 @@ pub struct OpenMessageRecord {
 }
 
 impl OpenMessageRecord {
+    /// Creates a new random id that can be used for a new record
+    pub fn new_id() -> Uuid {
+        Uuid::new_v4()
+    }
+
     #[cfg(test)]
     /// Create a dumb OpenMessage instance mainly for test purposes
     pub fn dummy() -> Self {

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

@@ -1,16 +1,18 @@
-use std::sync::Arc;
-
+use async_trait::async_trait;
 use chrono::Utc;
+use std::sync::Arc;
 
 use mithril_common::StdResult;
 use mithril_common::entities::{Epoch, ProtocolMessage, SignedEntityType};
 use mithril_persistence::sqlite::{ConnectionExtensions, SqliteConnection};
 
 use crate::database::query::{
     DeleteOpenMessageQuery, GetOpenMessageQuery, GetOpenMessageWithSingleSignaturesQuery,
-    InsertOpenMessageQuery, UpdateOpenMessageQuery,
+    InsertOpenMessageQuery, InsertOrReplaceOpenMessageQuery, UpdateOpenMessageQuery,
 };
 use crate::database::record::{OpenMessageRecord, OpenMessageWithSingleSignaturesRecord};
+use crate::entities::OpenMessage;
+use crate::services::OpenMessageStorer;
 
 /// ## Open message repository
 ///
@@ -79,6 +81,21 @@ impl OpenMessageRepository {
         message.ok_or_else(|| panic!("Inserting an open_message should not return nothing."))
     }
 
+    /// Create, or replace if one with the same [SignedEntityType] they already exist, a
+    /// [OpenMessageRecord] in the database.
+    pub async fn create_or_replace_open_message(
+        &self,
+        record: OpenMessageRecord,
+    ) -> StdResult<OpenMessageRecord> {
+        let message = self
+            .connection
+            .fetch_first(InsertOrReplaceOpenMessageQuery::one(record)?)?;
+
+        message.ok_or_else(|| {
+            panic!("Inserting or replacing an open_message should not return nothing.")
+        })
+    }
+
     /// Updates an [OpenMessageRecord] in the database.
     pub async fn update_open_message(
         &self,
@@ -102,6 +119,24 @@ impl OpenMessageRepository {
     }
 }
 
+#[async_trait]
+impl OpenMessageStorer for OpenMessageRepository {
+    async fn insert_or_replace_open_message(&self, open_message: OpenMessage) -> StdResult<()> {
+        let record = OpenMessageRecord {
+            open_message_id: OpenMessageRecord::new_id(),
+            epoch: open_message.epoch,
+            signed_entity_type: open_message.signed_entity_type,
+            protocol_message: open_message.protocol_message,
+            is_certified: open_message.is_certified,
+            is_expired: open_message.is_expired,
+            created_at: open_message.created_at,
+            expires_at: open_message.expires_at,
+        };
+        self.create_or_replace_open_message(record).await?;
+        Ok(())
+    }
+}
+
 #[cfg(test)]
 mod tests {
     use mithril_common::entities::{BlockNumber, CardanoDbBeacon};
@@ -261,6 +296,28 @@ mod tests {
         assert_eq!(open_message.epoch, message.epoch);
     }
 
+    #[tokio::test]
+    async fn repository_create_or_replace_open_message() {
+        let connection = get_connection().await;
+        let repository = OpenMessageRepository::new(connection.clone());
+        let mut inserted_record = repository
+            .create_or_replace_open_message(OpenMessageRecord {
+                epoch: Epoch(5),
+                signed_entity_type: SignedEntityType::MithrilStakeDistribution(Epoch(6)),
+                ..OpenMessageRecord::dummy()
+            })
+            .await
+            .unwrap();
+        assert_eq!(Epoch(5), inserted_record.epoch);
+
+        inserted_record.epoch = Epoch(32);
+        let replaced_record = repository
+            .create_or_replace_open_message(inserted_record)
+            .await
+            .unwrap();
+        assert_eq!(Epoch(32), replaced_record.epoch);
+    }
+
     #[tokio::test]
     async fn repository_update_open_message() {
         let connection = get_connection().await;

mithril-aggregator/src/entities/open_message.rs

@@ -120,7 +120,7 @@ mod test {
     fn test_from_record() {
         let created_at = Utc::now();
         let record = OpenMessageRecord {
-            open_message_id: Uuid::new_v4(),
+            open_message_id: OpenMessageRecord::new_id(),
             epoch: Epoch(1),
             signed_entity_type: SignedEntityType::dummy(),
             protocol_message: ProtocolMessage::default(),