oku_fs/database/
dht.rs

1use super::core::*;
2use miette::IntoDiagnostic;
3use native_db::*;
4use native_model::{native_model, Model};
5use serde::{Deserialize, Serialize};
6
7#[derive(Serialize, Deserialize, Debug, Clone)]
8#[native_model(id = 3, version = 1)]
9#[native_db(
10    primary_key(primary_key -> (Vec<u8>, Vec<u8>))
11)]
12/// A record of a replica announcement on the DHT.
13pub struct ReplicaAnnouncement {
14    /// The public key of the announcement.
15    #[primary_key]
16    pub key: Vec<u8>,
17    /// The signature of the announcement.
18    pub signature: Vec<u8>,
19}
20
21impl OkuDatabase {
22    /// Insert or update a replica announcement record.
23    ///
24    /// # Arguments
25    ///
26    /// * `announcement` - A replica announcement record to upsert.
27    ///
28    /// # Returns
29    ///
30    /// The previous record of the announcement, if one existed.
31    pub fn upsert_announcement(
32        &self,
33        announcement: &ReplicaAnnouncement,
34    ) -> miette::Result<Option<ReplicaAnnouncement>> {
35        let rw = self.database.rw_transaction().into_diagnostic()?;
36        let old_value: Option<ReplicaAnnouncement> =
37            rw.upsert(announcement.to_owned()).into_diagnostic()?;
38        rw.commit().into_diagnostic()?;
39        Ok(old_value)
40    }
41
42    /// Insert or update multiple replica announcement records.
43    ///
44    /// # Arguments
45    ///
46    /// * `announcements` - A list of replica announcement records to upsert.
47    ///
48    /// # Returns
49    ///
50    /// A list containing the previous record of each announcement, if one existed.
51    pub fn upsert_announcements(
52        &self,
53        announcements: &[ReplicaAnnouncement],
54    ) -> miette::Result<Vec<Option<ReplicaAnnouncement>>> {
55        let rw = self.database.rw_transaction().into_diagnostic()?;
56        let old_announcements: Vec<_> = announcements
57            .iter()
58            .cloned()
59            .filter_map(|announcement| rw.upsert(announcement).ok())
60            .collect();
61        rw.commit().into_diagnostic()?;
62        Ok(old_announcements)
63    }
64
65    /// Delete a replica announcement record.
66    ///
67    /// # Arguments
68    ///
69    /// * `announcement` - A replica announcement record to delete.
70    ///
71    /// # Returns
72    ///
73    /// The deleted replica announcement record.
74    pub fn delete_announcement(
75        &self,
76        announcement: &ReplicaAnnouncement,
77    ) -> miette::Result<ReplicaAnnouncement> {
78        let rw = self.database.rw_transaction().into_diagnostic()?;
79        let removed_announcement = rw.remove(announcement.to_owned()).into_diagnostic()?;
80        rw.commit().into_diagnostic()?;
81        Ok(removed_announcement)
82    }
83
84    /// Delete multiple replica announcement records.
85    ///
86    /// # Arguments
87    ///
88    /// * `announcements` - A list of replica announcement records to delete.
89    ///
90    /// # Returns
91    ///
92    /// A list containing the deleted replica announcement records.
93    pub fn delete_announcements(
94        &self,
95        announcements: &[ReplicaAnnouncement],
96    ) -> miette::Result<Vec<ReplicaAnnouncement>> {
97        let rw = self.database.rw_transaction().into_diagnostic()?;
98        let removed_announcements = announcements
99            .iter()
100            .filter_map(|announcement| rw.remove(announcement.to_owned()).ok())
101            .collect();
102        rw.commit().into_diagnostic()?;
103        Ok(removed_announcements)
104    }
105
106    /// Gets the replica announcements recorded by this node.
107    ///
108    /// # Returns
109    ///
110    /// The replica announcements recorded by this node.
111    pub fn get_announcements(&self) -> miette::Result<Vec<ReplicaAnnouncement>> {
112        let r = self.database.r_transaction().into_diagnostic()?;
113        r.scan()
114            .primary()
115            .into_diagnostic()?
116            .all()
117            .into_diagnostic()?
118            .collect::<Result<Vec<_>, _>>()
119            .into_diagnostic()
120    }
121
122    /// Gets a replica announcement record by its public key.
123    ///
124    /// # Arguments
125    ///
126    /// * `key` - The public key of the DHT announcement.
127    ///
128    /// # Returns
129    ///
130    /// A replica announcement record.
131    pub fn get_announcement(&self, key: &Vec<u8>) -> miette::Result<Option<ReplicaAnnouncement>> {
132        let r = self.database.r_transaction().into_diagnostic()?;
133        r.get().primary(key.to_owned()).into_diagnostic()
134    }
135}