xmtp_db/encrypted_store/
identity_update.rs

1use std::collections::HashMap;
2
3use crate::StorageError;
4use crate::impl_store;
5
6use super::{
7    ConnectionExt,
8    db_connection::DbConnection,
9    schema::identity_updates::{self, dsl},
10};
11use derive_builder::Builder;
12use diesel::{dsl::max, prelude::*};
13
14/// StoredIdentityUpdate holds a serialized IdentityUpdate record
15#[derive(Insertable, Identifiable, Queryable, Debug, Clone, PartialEq, Eq, Builder)]
16#[diesel(table_name = identity_updates)]
17#[diesel(primary_key(inbox_id, sequence_id))]
18#[builder(setter(into), build_fn(error = "StorageError"))]
19pub struct StoredIdentityUpdate {
20    pub inbox_id: String,
21    pub sequence_id: i64,
22    pub server_timestamp_ns: i64,
23    pub payload: Vec<u8>,
24}
25
26impl StoredIdentityUpdate {
27    pub fn build() -> StoredIdentityUpdateBuilder {
28        StoredIdentityUpdateBuilder::default()
29    }
30
31    pub fn new(
32        inbox_id: String,
33        sequence_id: i64,
34        server_timestamp_ns: i64,
35        payload: Vec<u8>,
36    ) -> Self {
37        Self {
38            inbox_id,
39            sequence_id,
40            server_timestamp_ns,
41            payload,
42        }
43    }
44}
45
46impl_store!(StoredIdentityUpdate, identity_updates);
47
48pub trait QueryIdentityUpdates {
49    /// Returns all identity updates for the given inbox ID up to the provided sequence_id.
50    /// Returns updates greater than `from_sequence_id` and less than _or equal to_ `to_sequence_id`
51    fn get_identity_updates<InboxId: AsRef<str>>(
52        &self,
53        inbox_id: InboxId,
54        from_sequence_id: Option<i64>,
55        to_sequence_id: Option<i64>,
56    ) -> Result<Vec<StoredIdentityUpdate>, crate::ConnectionError>;
57
58    /// Batch insert identity updates, ignoring duplicates.
59    fn insert_or_ignore_identity_updates(
60        &self,
61        updates: &[StoredIdentityUpdate],
62    ) -> Result<(), crate::ConnectionError>;
63
64    fn get_latest_sequence_id_for_inbox(
65        &self,
66        inbox_id: &str,
67    ) -> Result<i64, crate::ConnectionError>;
68
69    /// Given a list of inbox_ids return a HashMap of each inbox ID -> highest known sequence ID
70    fn get_latest_sequence_id(
71        &self,
72        inbox_ids: &[&str],
73    ) -> Result<HashMap<String, i64>, crate::ConnectionError>;
74
75    /// Returns the count of identity updates for inbox_ids
76    fn count_inbox_updates(
77        &self,
78        inbox_ids: &[&str],
79    ) -> Result<HashMap<String, i64>, crate::ConnectionError>;
80}
81
82impl<T> QueryIdentityUpdates for &T
83where
84    T: QueryIdentityUpdates,
85{
86    fn get_identity_updates<InboxId: AsRef<str>>(
87        &self,
88        inbox_id: InboxId,
89        from_sequence_id: Option<i64>,
90        to_sequence_id: Option<i64>,
91    ) -> Result<Vec<StoredIdentityUpdate>, crate::ConnectionError> {
92        (**self).get_identity_updates(inbox_id, from_sequence_id, to_sequence_id)
93    }
94
95    fn insert_or_ignore_identity_updates(
96        &self,
97        updates: &[StoredIdentityUpdate],
98    ) -> Result<(), crate::ConnectionError> {
99        (**self).insert_or_ignore_identity_updates(updates)
100    }
101
102    fn get_latest_sequence_id_for_inbox(
103        &self,
104        inbox_id: &str,
105    ) -> Result<i64, crate::ConnectionError> {
106        (**self).get_latest_sequence_id_for_inbox(inbox_id)
107    }
108
109    fn get_latest_sequence_id(
110        &self,
111        inbox_ids: &[&str],
112    ) -> Result<HashMap<String, i64>, crate::ConnectionError> {
113        (**self).get_latest_sequence_id(inbox_ids)
114    }
115
116    fn count_inbox_updates(
117        &self,
118        inbox_ids: &[&str],
119    ) -> Result<HashMap<String, i64>, crate::ConnectionError> {
120        (**self).count_inbox_updates(inbox_ids)
121    }
122}
123
124impl<C: ConnectionExt> QueryIdentityUpdates for DbConnection<C> {
125    /// Returns all identity updates for the given inbox ID up to the provided sequence_id.
126    /// Returns updates greater than `from_sequence_id` and less than _or equal to_ `to_sequence_id`
127    fn get_identity_updates<InboxId: AsRef<str>>(
128        &self,
129        inbox_id: InboxId,
130        from_sequence_id: Option<i64>,
131        to_sequence_id: Option<i64>,
132    ) -> Result<Vec<StoredIdentityUpdate>, crate::ConnectionError> {
133        let mut query = dsl::identity_updates
134            .order(dsl::sequence_id.asc())
135            .filter(dsl::inbox_id.eq(inbox_id.as_ref()))
136            .into_boxed();
137
138        if let Some(sequence_id) = from_sequence_id {
139            query = query.filter(dsl::sequence_id.gt(sequence_id));
140        }
141
142        if let Some(sequence_id) = to_sequence_id {
143            query = query.filter(dsl::sequence_id.le(sequence_id));
144        }
145
146        self.raw_query(|conn| query.load::<StoredIdentityUpdate>(conn))
147    }
148
149    /// Batch insert identity updates, ignoring duplicates.
150    #[tracing::instrument(level = "trace", skip(updates))]
151    fn insert_or_ignore_identity_updates(
152        &self,
153        updates: &[StoredIdentityUpdate],
154    ) -> Result<(), crate::ConnectionError> {
155        self.raw_query(|conn| {
156            diesel::insert_or_ignore_into(dsl::identity_updates)
157                .values(updates)
158                .execute(conn)
159        })?;
160        Ok(())
161    }
162
163    fn get_latest_sequence_id_for_inbox(
164        &self,
165        inbox_id: &str,
166    ) -> Result<i64, crate::ConnectionError> {
167        let query = dsl::identity_updates
168            .select(dsl::sequence_id)
169            .order(dsl::sequence_id.desc())
170            .limit(1)
171            .filter(dsl::inbox_id.eq(inbox_id))
172            .into_boxed();
173
174        self.raw_query(|conn| query.first::<i64>(conn))
175    }
176
177    /// Given a list of inbox_ids return a HashMap of each inbox ID -> highest known sequence ID
178    #[tracing::instrument(level = "trace", skip_all)]
179    fn get_latest_sequence_id(
180        &self,
181        inbox_ids: &[&str],
182    ) -> Result<HashMap<String, i64>, crate::ConnectionError> {
183        // Query IdentityUpdates grouped by inbox_id, getting the max sequence_id
184        let query = dsl::identity_updates
185            .group_by(dsl::inbox_id)
186            .select((dsl::inbox_id, max(dsl::sequence_id)))
187            .filter(dsl::inbox_id.eq_any(inbox_ids));
188
189        // Get the results as a Vec of (inbox_id, sequence_id) tuples
190        let result_tuples: Vec<(String, i64)> = self
191            .raw_query(|conn| query.load::<(String, Option<i64>)>(conn))?
192            .into_iter()
193            // Diesel needs an Option type for aggregations like max(sequence_id), so we
194            // unwrap the option here
195            .filter_map(|(inbox_id, sequence_id_opt)| {
196                sequence_id_opt.map(|sequence_id| (inbox_id, sequence_id))
197            })
198            .collect();
199
200        // Convert the Vec to a HashMap
201        Ok(HashMap::from_iter(result_tuples))
202    }
203
204    fn count_inbox_updates(
205        &self,
206        inbox_ids: &[&str],
207    ) -> Result<HashMap<String, i64>, crate::ConnectionError> {
208        use diesel::dsl::count_star;
209        let query = dsl::identity_updates
210            .group_by(dsl::inbox_id)
211            .select((dsl::inbox_id, count_star()))
212            .filter(dsl::inbox_id.eq_any(inbox_ids));
213        self.raw_query(|conn| {
214            query
215                .load_iter::<(String, i64), _>(conn)?
216                .collect::<Result<HashMap<_, _>, _>>()
217        })
218    }
219}
220
221#[cfg(test)]
222pub(crate) mod tests {
223    use crate::{Store, test_utils::with_connection};
224    use xmtp_common::{rand_time, rand_vec};
225
226    use super::*;
227
228    fn build_update(inbox_id: &str, sequence_id: i64) -> StoredIdentityUpdate {
229        StoredIdentityUpdate::new(
230            inbox_id.to_string(),
231            sequence_id,
232            rand_time(),
233            rand_vec::<24>(),
234        )
235    }
236
237    #[xmtp_common::test]
238    fn insert_and_read() {
239        with_connection(|conn| {
240            let inbox_id = "inbox_1";
241            let update_1 = build_update(inbox_id, 1);
242            let update_1_payload = update_1.payload.clone();
243            let update_2 = build_update(inbox_id, 2);
244            let update_2_payload = update_2.payload.clone();
245
246            update_1.store(conn).expect("should store without error");
247            update_2.store(conn).expect("should store without error");
248
249            let all_updates = conn
250                .get_identity_updates(inbox_id, None, None)
251                .expect("query should work");
252
253            assert_eq!(all_updates.len(), 2);
254            let first_update = all_updates.first().unwrap();
255            assert_eq!(first_update.payload, update_1_payload);
256            let second_update = all_updates.last().unwrap();
257            assert_eq!(second_update.payload, update_2_payload);
258        })
259    }
260
261    #[xmtp_common::test]
262    fn test_filter() {
263        with_connection(|conn| {
264            let inbox_id = "inbox_1";
265            let update_1 = build_update(inbox_id, 1);
266            let update_2 = build_update(inbox_id, 2);
267            let update_3 = build_update(inbox_id, 3);
268
269            conn.insert_or_ignore_identity_updates(&[update_1, update_2, update_3])
270                .expect("insert should succeed");
271
272            let update_1_and_2 = conn
273                .get_identity_updates(inbox_id, None, Some(2))
274                .expect("query should work");
275
276            assert_eq!(update_1_and_2.len(), 2);
277
278            let all_updates = conn
279                .get_identity_updates(inbox_id, None, None)
280                .expect("query should work");
281
282            assert_eq!(all_updates.len(), 3);
283
284            let only_update_2 = conn
285                .get_identity_updates(inbox_id, Some(1), Some(2))
286                .expect("query should work");
287
288            assert_eq!(only_update_2.len(), 1);
289            assert_eq!(only_update_2[0].sequence_id, 2);
290        })
291    }
292
293    #[xmtp_common::test]
294    fn test_get_latest_sequence_id() {
295        with_connection(|conn| {
296            let inbox_1 = "inbox_1";
297            let inbox_2 = "inbox_2";
298            let update_1 = build_update(inbox_1, 1);
299            let update_2 = build_update(inbox_1, 3);
300            let update_3 = build_update(inbox_2, 5);
301            let update_4 = build_update(inbox_2, 6);
302
303            conn.insert_or_ignore_identity_updates(&[update_1, update_2, update_3, update_4])
304                .expect("insert should succeed");
305
306            let latest_sequence_ids = conn
307                .get_latest_sequence_id(&[inbox_1, inbox_2])
308                .expect("query should work");
309
310            assert_eq!(latest_sequence_ids.get(inbox_1), Some(&3));
311            assert_eq!(latest_sequence_ids.get(inbox_2), Some(&6));
312
313            let latest_sequence_ids_with_missing_member = conn
314                .get_latest_sequence_id(&[inbox_1, "missing_inbox"])
315                .expect("should still succeed");
316
317            assert_eq!(
318                latest_sequence_ids_with_missing_member.get(inbox_1),
319                Some(&3)
320            );
321            assert_eq!(
322                latest_sequence_ids_with_missing_member.get("missing_inbox"),
323                None
324            );
325        })
326    }
327
328    #[xmtp_common::test]
329    fn get_single_sequence_id() {
330        with_connection(|conn| {
331            let inbox_id = "inbox_1";
332            let update = build_update(inbox_id, 1);
333            let update_2 = build_update(inbox_id, 2);
334            update.store(conn).expect("should store without error");
335            update_2.store(conn).expect("should store without error");
336
337            let sequence_id = conn
338                .get_latest_sequence_id_for_inbox(inbox_id)
339                .expect("query should work");
340            assert_eq!(sequence_id, 2);
341        })
342    }
343
344    #[xmtp_common::test]
345    fn test_count_inbox_updates() {
346        with_connection(|conn| {
347            let inbox_1 = "inbox_1";
348            let inbox_2 = "inbox_2";
349            conn.insert_or_ignore_identity_updates(&[
350                build_update(inbox_1, 1),
351                build_update(inbox_1, 2),
352                build_update(inbox_2, 1),
353            ])
354            .unwrap();
355            let counts = conn
356                .count_inbox_updates(&[inbox_1, inbox_2, "missing"])
357                .unwrap();
358            assert_eq!(counts.get(inbox_1), Some(&2));
359            assert_eq!(counts.get(inbox_2), Some(&1));
360            assert_eq!(counts.get("missing"), None);
361        })
362    }
363}
xmtp_db/encrypted_store/identity_update.rs

xmtp_db/encrypted_store/
identity_update.rs
