xmtp_mls/groups/message_list/
page.rs

1use super::*;
2use xmtp_db::{
3    TransactionalKeyStore, XmtpMlsStorageProvider,
4    delivery::{AppVisibleMessageRow, QueryDelivery},
5    diesel::{Connection, SqliteConnection},
6    group_message::RecoveryQueryArgs,
7};
8
9pub struct EnrichedMessagePage<P> {
10    pub messages: Vec<EnrichedStoredMessage>,
11    pub first_position: Option<P>,
12    pub last_position: Option<P>,
13    pub has_more: bool,
14}
15
16pub type EnrichedHistoryPage = EnrichedMessagePage<xmtp_db::delivery::HistoryPosition>;
17pub type EnrichedRecoveryPage = EnrichedMessagePage<xmtp_db::group_message::RecoveryPosition>;
18
19#[cfg(test)]
20thread_local! {
21    static BEFORE_ENRICHMENT: std::cell::RefCell<Option<Box<dyn FnOnce()>>> = const { std::cell::RefCell::new(None) };
22}
23
24fn with_page_snapshot<T>(
25    conn: &impl DbQuery,
26    read: impl FnOnce(&mut SqliteConnection) -> Result<T, EnrichMessageError>,
27) -> Result<T, EnrichMessageError> {
28    // A deferred transaction allows a WAL writer while this page keeps its snapshot.
29    conn.raw_query(|sqlite| Ok(sqlite.transaction(read)))?
30}
31
32fn visible_recovery_query(mut query: RecoveryQueryArgs) -> RecoveryQueryArgs {
33    query.messages = filter_out_hidden_message_types_from_query(&query.messages);
34    query
35}
36
37fn enrich_page_rows(
38    conn: &impl DbQuery,
39    rows: Vec<AppVisibleMessageRow>,
40) -> Result<Vec<EnrichedStoredMessage>, EnrichMessageError> {
41    #[cfg(test)]
42    BEFORE_ENRICHMENT.with_borrow_mut(|hook| {
43        if let Some(hook) = hook.take() {
44            hook();
45        }
46    });
47    let positions: std::collections::HashMap<_, _> = rows
48        .iter()
49        .enumerate()
50        .map(|(index, row)| (row.stored.id.clone(), (index, row.cursor)))
51        .collect();
52    let mut sources: std::collections::HashMap<_, Vec<_>> = std::collections::HashMap::new();
53    for row in rows {
54        sources
55            .entry(row.stored.group_id)
56            .or_default()
57            .push(row.stored);
58    }
59    let mut messages = Vec::with_capacity(positions.len());
60    for (group_id, rows) in sources {
61        messages.extend(enrich_messages_with_stored(conn, &group_id, rows)?);
62    }
63    // Relations use the physical source. Display order uses the selected raw keys.
64    messages.sort_by_key(|message| positions[&message.stored.id].0);
65    for message in &mut messages {
66        message.delivery_cursor = positions[&message.stored.id].1;
67    }
68    Ok(messages)
69}
70
71fn history_page_with_stored(
72    conn: &impl DbQuery,
73    group_id: &xmtp_proto::types::GroupId,
74    query: &MsgQueryArgs,
75) -> Result<EnrichedHistoryPage, EnrichMessageError> {
76    with_page_snapshot(conn, |sqlite| {
77        let store = sqlite.key_store();
78        let conn = store.db();
79        let page =
80            conn.history_page_rows(group_id, &filter_out_hidden_message_types_from_query(query))?;
81        Ok(EnrichedHistoryPage {
82            messages: enrich_page_rows(&conn, page.rows)?,
83            first_position: page.first_position,
84            last_position: page.last_position,
85            has_more: page.has_more,
86        })
87    })
88}
89
90fn recovery_page_with_stored(
91    conn: &impl DbQuery,
92    group_id: &xmtp_proto::types::GroupId,
93    query: RecoveryQueryArgs,
94) -> Result<EnrichedRecoveryPage, EnrichMessageError> {
95    let query = visible_recovery_query(query);
96    with_page_snapshot(conn, |sqlite| {
97        let store = sqlite.key_store();
98        let conn = store.db();
99        let page = conn.recovery_page_rows(group_id, &query)?;
100        Ok(EnrichedRecoveryPage {
101            messages: enrich_page_rows(&conn, page.rows)?,
102            first_position: page.first_position,
103            last_position: page.last_position,
104            has_more: page.has_more,
105        })
106    })
107}
108
109impl<Context: XmtpSharedContext> MlsGroup<Context> {
110    /// Read base rows, raw coverage and relations in one database snapshot.
111    pub fn find_history_page_with_stored(
112        &self,
113        query: &MsgQueryArgs,
114    ) -> Result<EnrichedHistoryPage, EnrichMessageError> {
115        history_page_with_stored(&self.context.db(), &self.group_id, query)
116    }
117
118    /// Read Failed and Unpublished rows and relations in one database snapshot.
119    // implements: DMS-009
120    pub fn find_recovery_page_with_stored(
121        &self,
122        query: RecoveryQueryArgs,
123    ) -> Result<EnrichedRecoveryPage, EnrichMessageError> {
124        recovery_page_with_stored(&self.context.db(), &self.group_id, query)
125    }
126}
127
128#[cfg(test)]
129mod tests;
xmtp_mls/groups/message_list/page.rs

xmtp_mls/groups/message_list/
page.rs
