ankurah_storage_postgres/
dump.rs

1//! Cursor-backed logical dumps for PostgreSQL storage.
2
3use std::{
4    collections::{BTreeMap, VecDeque},
5    pin::Pin,
6};
7
8use ankurah_core::{
9    error::RetrievalError,
10    storage::{StorageDump, StorageDumpItem},
11};
12use ankurah_proto::{
13    Attestation, AttestationSet, Attested, Clock, CollectionId, EntityId, EntityState, Event, EventId, OperationSet, State, StateBuffers,
14};
15use async_trait::async_trait;
16use bb8_postgres::{tokio_postgres::NoTls, PostgresConnectionManager};
17use futures_util::{stream, Stream};
18
19use crate::Postgres;
20
21const PAGE_SIZE: i64 = 512;
22
23type Pool = bb8::Pool<PostgresConnectionManager<NoTls>>;
24type BoxDumpStream = Pin<Box<dyn Stream<Item = Result<StorageDumpItem, RetrievalError>> + Send + 'static>>;
25
26#[async_trait]
27impl StorageDump for Postgres {
28    type DumpStream = BoxDumpStream;
29
30    async fn dump(&self) -> Result<Self::DumpStream, RetrievalError> {
31        let client = self.pool.get().await.map_err(RetrievalError::storage)?;
32        let (event_collections, state_collections) = dump_collections(&client).await.map_err(RetrievalError::from)?;
33        drop(client);
34
35        let cursor = PostgresDumpCursor {
36            pool: self.pool.clone(),
37            phase: DumpPhase::Events,
38            event_collections,
39            state_collections,
40            collection_index: 0,
41            after_event: None,
42            after_state: None,
43            pending: VecDeque::new(),
44        };
45        Ok(Box::pin(stream::try_unfold(cursor, |mut cursor| async move { Ok(cursor.next().await?.map(|item| (item, cursor))) })))
46    }
47}
48
49#[derive(Clone, Copy)]
50enum DumpPhase {
51    Events,
52    States,
53    Done,
54}
55
56struct PostgresDumpCursor {
57    pool: Pool,
58    phase: DumpPhase,
59    event_collections: Vec<CollectionId>,
60    state_collections: Vec<CollectionId>,
61    collection_index: usize,
62    after_event: Option<EventId>,
63    after_state: Option<EntityId>,
64    pending: VecDeque<StorageDumpItem>,
65}
66
67impl PostgresDumpCursor {
68    async fn next(&mut self) -> Result<Option<StorageDumpItem>, RetrievalError> {
69        loop {
70            if let Some(item) = self.pending.pop_front() {
71                return Ok(Some(item));
72            }
73            match self.phase {
74                DumpPhase::Events => {
75                    let Some(collection) = self.event_collections.get(self.collection_index).cloned() else {
76                        self.phase = DumpPhase::States;
77                        self.collection_index = 0;
78                        continue;
79                    };
80                    let client = self.pool.get().await.map_err(RetrievalError::storage)?;
81                    let page = event_page(&client, &collection, self.after_event.as_ref()).await.map_err(RetrievalError::from)?;
82                    if page.is_empty() {
83                        self.collection_index += 1;
84                        self.after_event = None;
85                        continue;
86                    }
87                    self.after_event = page.last().map(|(id, _)| id.clone());
88                    self.pending.extend(page.into_iter().map(|(_, event)| StorageDumpItem::Event(event)));
89                }
90                DumpPhase::States => {
91                    let Some(collection) = self.state_collections.get(self.collection_index).cloned() else {
92                        self.phase = DumpPhase::Done;
93                        continue;
94                    };
95                    let client = self.pool.get().await.map_err(RetrievalError::storage)?;
96                    let page = state_page(&client, &collection, self.after_state.as_ref()).await.map_err(RetrievalError::from)?;
97                    if page.is_empty() {
98                        self.collection_index += 1;
99                        self.after_state = None;
100                        continue;
101                    }
102                    self.after_state = page.last().map(|(id, _)| *id);
103                    self.pending.extend(page.into_iter().map(|(_, state)| StorageDumpItem::State(state)));
104                }
105                DumpPhase::Done => return Ok(None),
106            }
107        }
108    }
109}
110
111async fn event_page(
112    client: &tokio_postgres::Client,
113    collection: &CollectionId,
114    after: Option<&EventId>,
115) -> anyhow::Result<Vec<(EventId, Attested<Event>)>> {
116    let table = quote_identifier(&format!("{collection}_event"));
117    let rows = if let Some(after) = after {
118        client
119            .query(
120                &format!("SELECT id, entity_id, operations, parent, attestations FROM {table} WHERE id > $1 ORDER BY id LIMIT $2"),
121                &[after, &PAGE_SIZE],
122            )
123            .await?
124    } else {
125        client
126            .query(&format!("SELECT id, entity_id, operations, parent, attestations FROM {table} ORDER BY id LIMIT $1"), &[&PAGE_SIZE])
127            .await?
128    };
129
130    rows.into_iter()
131        .map(|row| {
132            let id: EventId = row.try_get("id")?;
133            let entity_id: EntityId = row.try_get("entity_id")?;
134            let operations: Vec<u8> = row.try_get("operations")?;
135            let parent: Clock = row.try_get("parent")?;
136            let attestations: Vec<u8> = row.try_get("attestations")?;
137            let event = Attested {
138                payload: Event {
139                    collection: collection.clone(),
140                    entity_id,
141                    operations: bincode::deserialize::<OperationSet>(&operations)?,
142                    parent,
143                },
144                attestations: bincode::deserialize::<AttestationSet>(&attestations)?,
145            };
146            anyhow::ensure!(event.payload.id() == id, "stored event id does not match payload for {collection}/{id}");
147            Ok((id, event))
148        })
149        .collect()
150}
151
152async fn state_page(
153    client: &tokio_postgres::Client,
154    collection: &CollectionId,
155    after: Option<&EntityId>,
156) -> anyhow::Result<Vec<(EntityId, Attested<EntityState>)>> {
157    let table = quote_identifier(collection.as_str());
158    let rows = if let Some(after) = after {
159        client
160            .query(
161                &format!("SELECT id, state_buffer, head, attestations FROM {table} WHERE id > $1 ORDER BY id LIMIT $2"),
162                &[after, &PAGE_SIZE],
163            )
164            .await?
165    } else {
166        client.query(&format!("SELECT id, state_buffer, head, attestations FROM {table} ORDER BY id LIMIT $1"), &[&PAGE_SIZE]).await?
167    };
168
169    rows.into_iter()
170        .map(|row| {
171            let entity_id: EntityId = row.try_get("id")?;
172            let state_buffers: Vec<u8> = row.try_get("state_buffer")?;
173            let head: Clock = row.try_get("head")?;
174            let attestations: Vec<Vec<u8>> = row.try_get("attestations")?;
175            let attestations =
176                attestations.into_iter().map(|bytes| bincode::deserialize::<Attestation>(&bytes)).collect::<Result<Vec<_>, _>>()?;
177            Ok((
178                entity_id,
179                Attested {
180                    payload: EntityState {
181                        entity_id,
182                        collection: collection.clone(),
183                        state: State {
184                            state_buffers: StateBuffers(bincode::deserialize::<BTreeMap<String, Vec<u8>>>(&state_buffers)?),
185                            head,
186                        },
187                    },
188                    attestations: AttestationSet(attestations),
189                },
190            ))
191        })
192        .collect()
193}
194
195async fn dump_collections(client: &tokio_postgres::Client) -> anyhow::Result<(Vec<CollectionId>, Vec<CollectionId>)> {
196    let columns = table_columns(client).await?;
197    let state_columns = ["id", "state_buffer", "head", "attestations"];
198    let event_columns = ["id", "entity_id", "operations", "parent", "attestations"];
199    let mut events = Vec::new();
200    let mut states = Vec::new();
201    for (name, present) in columns {
202        if state_columns.iter().all(|column| present.contains(*column)) {
203            states.push(CollectionId::from(name.clone()));
204        }
205        if event_columns.iter().all(|column| present.contains(*column)) {
206            if let Some(collection) = name.strip_suffix("_event") {
207                events.push(CollectionId::from(collection));
208            }
209        }
210    }
211    Ok((events, states))
212}
213
214async fn table_columns(client: &tokio_postgres::Client) -> anyhow::Result<BTreeMap<String, std::collections::BTreeSet<String>>> {
215    let rows = client
216        .query(
217            "SELECT table_name, column_name
218             FROM information_schema.columns
219             WHERE table_schema = current_schema()
220             ORDER BY table_name, ordinal_position",
221            &[],
222        )
223        .await?;
224    let mut columns = BTreeMap::<String, std::collections::BTreeSet<String>>::new();
225    for row in rows {
226        columns.entry(row.try_get("table_name")?).or_default().insert(row.try_get("column_name")?);
227    }
228    Ok(columns)
229}
230
231fn quote_identifier(identifier: &str) -> String { format!("\"{}\"", identifier.replace('"', "\"\"")) }
ankurah_storage_postgres/dump.rs

ankurah_storage_postgres/
dump.rs
