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

xmtp_mls/groups/message_list/
page.rs
