ankurah_storage_postgres/
lib.rs

1use std::{
2    collections::{hash_map::DefaultHasher, BTreeMap},
3    hash::{Hash, Hasher},
4    sync::{Arc, RwLock},
5    time::Duration,
6};
7
8use ankurah_core::{
9    error::{MutationError, RetrievalError, StateError},
10    property::backend::backend_from_string,
11    storage::{StorageCollection, StorageEngine},
12};
13use ankurah_proto::{Attestation, AttestationSet, Attested, EntityState, EventId, OperationSet, State, StateBuffers};
14
15use futures_util::{pin_mut, TryStreamExt};
16
17pub mod sql_builder;
18pub mod value;
19
20use value::PGValue;
21
22use ankurah_proto::{Clock, CollectionId, EntityId, Event};
23use async_trait::async_trait;
24use bb8_postgres::{tokio_postgres::NoTls, PostgresConnectionManager};
25use tokio_postgres::{error::SqlState, types::ToSql};
26use tracing::{debug, error, info, warn};
27
28/// Default connection pool size for `Postgres::open()`.
29/// Production applications should configure their own pool via `Postgres::new()`.
30pub const DEFAULT_POOL_SIZE: u32 = 15;
31
32/// Default connection timeout in seconds
33pub const DEFAULT_CONNECTION_TIMEOUT_SECS: u64 = 30;
34
35pub struct Postgres {
36    pool: bb8::Pool<PostgresConnectionManager<NoTls>>,
37}
38
39impl Postgres {
40    pub fn new(pool: bb8::Pool<PostgresConnectionManager<NoTls>>) -> anyhow::Result<Self> { Ok(Self { pool }) }
41
42    pub async fn open(uri: &str) -> anyhow::Result<Self> {
43        let manager = PostgresConnectionManager::new_from_stringlike(uri, NoTls)?;
44        let pool = bb8::Pool::builder()
45            .max_size(DEFAULT_POOL_SIZE)
46            .connection_timeout(Duration::from_secs(DEFAULT_CONNECTION_TIMEOUT_SECS))
47            .build(manager)
48            .await?;
49        Self::new(pool)
50    }
51
52    // TODO: newtype this to `BucketName(&str)` with a constructor that
53    // only accepts a subset of characters.
54    pub fn sane_name(collection: &str) -> bool {
55        for char in collection.chars() {
56            match char {
57                char if char.is_alphanumeric() => {}
58                char if char.is_numeric() => {}
59                '_' | '.' | ':' => {}
60                _ => return false,
61            }
62        }
63
64        true
65    }
66}
67
68/// Compute advisory lock key from a string identifier
69fn advisory_lock_key(identifier: &str) -> i64 {
70    let mut hasher = DefaultHasher::new();
71    identifier.hash(&mut hasher);
72    hasher.finish() as i64
73}
74
75/// Acquire a PostgreSQL advisory lock for DDL operations on a collection
76async fn acquire_ddl_lock(client: &tokio_postgres::Client, collection_id: &str) -> Result<i64, StateError> {
77    let lock_key = advisory_lock_key(&format!("ankurah_ddl:{}", collection_id));
78    debug!("Acquiring advisory lock {} for collection {}", lock_key, collection_id);
79    client.execute("SELECT pg_advisory_lock($1)", &[&lock_key]).await.map_err(|err| {
80        error!("Failed to acquire advisory lock for {}: {:?}", collection_id, err);
81        StateError::DDLError(Box::new(err))
82    })?;
83    Ok(lock_key)
84}
85
86/// Release a PostgreSQL advisory lock
87async fn release_ddl_lock(client: &tokio_postgres::Client, lock_key: i64) -> Result<(), StateError> {
88    debug!("Releasing advisory lock {}", lock_key);
89    client.execute("SELECT pg_advisory_unlock($1)", &[&lock_key]).await.map_err(|err| {
90        error!("Failed to release advisory lock {}: {:?}", lock_key, err);
91        StateError::DDLError(Box::new(err))
92    })?;
93    Ok(())
94}
95
96#[async_trait]
97impl StorageEngine for Postgres {
98    type Value = PGValue;
99
100    async fn collection(&self, collection_id: &CollectionId) -> Result<std::sync::Arc<dyn StorageCollection>, RetrievalError> {
101        if !Postgres::sane_name(collection_id.as_str()) {
102            return Err(RetrievalError::InvalidBucketName);
103        }
104
105        let mut client = self.pool.get().await.map_err(RetrievalError::storage)?;
106
107        // get the current schema from the database
108        let schema = client.query_one("SELECT current_database()", &[]).await.map_err(RetrievalError::storage)?;
109        let schema = schema.get("current_database");
110
111        let bucket = PostgresBucket {
112            pool: self.pool.clone(),
113            schema,
114            collection_id: collection_id.clone(),
115            columns: Arc::new(RwLock::new(Vec::new())),
116        };
117
118        // Acquire advisory lock to serialize DDL operations for this collection
119        let lock_key = acquire_ddl_lock(&client, collection_id.as_str()).await?;
120
121        // Create tables if they don't exist (protected by advisory lock)
122        let result = async {
123            bucket.create_state_table(&mut client).await?;
124            bucket.create_event_table(&mut client).await?;
125            bucket.rebuild_columns_cache(&mut client).await?;
126            Ok::<_, StateError>(())
127        }
128        .await;
129
130        // Always release the lock, even if DDL failed
131        release_ddl_lock(&client, lock_key).await?;
132
133        result?;
134        Ok(Arc::new(bucket))
135    }
136
137    async fn delete_all_collections(&self) -> Result<bool, MutationError> {
138        let mut client = self.pool.get().await.map_err(|err| MutationError::General(Box::new(err)))?;
139
140        // Get all tables in the public schema
141        let query = r#"
142            SELECT table_name 
143            FROM information_schema.tables 
144            WHERE table_schema = 'public'
145        "#;
146
147        let rows = client.query(query, &[]).await.map_err(|err| MutationError::General(Box::new(err)))?;
148        if rows.is_empty() {
149            return Ok(false);
150        }
151
152        // Start a transaction to drop all tables atomically
153        let transaction = client.transaction().await.map_err(|err| MutationError::General(Box::new(err)))?;
154
155        // Drop each table
156        for row in rows {
157            let table_name: String = row.get("table_name");
158            let drop_query = format!(r#"DROP TABLE IF EXISTS "{}""#, table_name);
159            transaction.execute(&drop_query, &[]).await.map_err(|err| MutationError::General(Box::new(err)))?;
160        }
161
162        // Commit the transaction
163        transaction.commit().await.map_err(|err| MutationError::General(Box::new(err)))?;
164
165        Ok(true)
166    }
167}
168
169#[derive(Clone, Debug)]
170pub struct PostgresColumn {
171    pub name: String,
172    pub is_nullable: bool,
173    pub data_type: String,
174}
175
176pub struct PostgresBucket {
177    pool: bb8::Pool<PostgresConnectionManager<NoTls>>,
178    collection_id: CollectionId,
179    schema: String,
180    columns: Arc<RwLock<Vec<PostgresColumn>>>,
181}
182
183impl PostgresBucket {
184    fn state_table(&self) -> String { self.collection_id.as_str().to_string() }
185
186    pub fn event_table(&self) -> String { format!("{}_event", self.collection_id.as_str()) }
187
188    /// Rebuild the cache of columns in the table.
189    pub async fn rebuild_columns_cache(&self, client: &mut tokio_postgres::Client) -> Result<(), StateError> {
190        debug!("PostgresBucket({}).rebuild_columns_cache", self.collection_id);
191        let column_query =
192            r#"SELECT column_name, is_nullable, data_type FROM information_schema.columns WHERE table_catalog = $1 AND table_name = $2;"#
193                .to_string();
194        let mut new_columns = Vec::new();
195        debug!("Querying existing columns: {:?}, [{:?}, {:?}]", column_query, &self.schema, &self.collection_id.as_str());
196        let rows = client
197            .query(&column_query, &[&self.schema, &self.collection_id.as_str()])
198            .await
199            .map_err(|err| StateError::DDLError(Box::new(err)))?;
200        for row in rows {
201            let is_nullable: String = row.get("is_nullable");
202            new_columns.push(PostgresColumn {
203                name: row.get("column_name"),
204                is_nullable: is_nullable.eq("YES"),
205                data_type: row.get("data_type"),
206            })
207        }
208
209        let mut columns = self.columns.write().unwrap();
210        *columns = new_columns;
211        drop(columns);
212
213        Ok(())
214    }
215
216    pub fn existing_columns(&self) -> Vec<String> {
217        let columns = self.columns.read().unwrap();
218        columns.iter().map(|column| column.name.clone()).collect()
219    }
220
221    pub fn column(&self, column_name: &String) -> Option<PostgresColumn> {
222        let columns = self.columns.read().unwrap();
223        columns.iter().find(|column| column.name == *column_name).cloned()
224    }
225
226    pub fn has_column(&self, column_name: &String) -> bool { self.column(column_name).is_some() }
227
228    pub async fn create_event_table(&self, client: &mut tokio_postgres::Client) -> Result<(), StateError> {
229        let create_query = format!(
230            r#"CREATE TABLE IF NOT EXISTS "{}"(
231                "id" character(43) PRIMARY KEY,
232                "entity_id" character(22),
233                "operations" bytea,
234                "parent" character(43)[],
235                "attestations" bytea
236            )"#,
237            self.event_table()
238        );
239
240        debug!("{create_query}");
241        client.execute(&create_query, &[]).await.map_err(|e| StateError::DDLError(Box::new(e)))?;
242        Ok(())
243    }
244
245    pub async fn create_state_table(&self, client: &mut tokio_postgres::Client) -> Result<(), StateError> {
246        let create_query = format!(
247            r#"CREATE TABLE IF NOT EXISTS "{}"(
248                "id" character(22) PRIMARY KEY,
249                "state_buffer" BYTEA,
250                "head" character(43)[],
251                "attestations" BYTEA[]
252            )"#,
253            self.state_table()
254        );
255
256        debug!("{create_query}");
257        match client.execute(&create_query, &[]).await {
258            Ok(_) => Ok(()),
259            Err(err) => {
260                // Log full error details for debugging
261                if let Some(db_err) = err.as_db_error() {
262                    error!("PostgresBucket({}).create_state_table error: {} (code: {:?})", self.collection_id, db_err, db_err.code());
263                } else {
264                    error!("PostgresBucket({}).create_state_table error: {:?}", self.collection_id, err);
265                }
266                Err(StateError::DDLError(Box::new(err)))
267            }
268        }
269    }
270
271    pub async fn add_missing_columns(
272        &self,
273        client: &mut tokio_postgres::Client,
274        missing: Vec<(String, &'static str)>, // column name, datatype
275    ) -> Result<(), StateError> {
276        if missing.is_empty() {
277            return Ok(());
278        }
279
280        // Acquire advisory lock to serialize DDL operations for this collection
281        let lock_key = acquire_ddl_lock(client, self.collection_id.as_str()).await?;
282
283        let result = async {
284            // Re-check columns after acquiring lock (another session may have added them)
285            self.rebuild_columns_cache(client).await?;
286
287            for (column, datatype) in missing {
288                if Postgres::sane_name(&column) && !self.has_column(&column) {
289                    let alter_query = format!(r#"ALTER TABLE "{}" ADD COLUMN "{}" {}"#, self.state_table(), column, datatype);
290                    info!("PostgresBucket({}).add_missing_columns: {}", self.collection_id, alter_query);
291                    match client.execute(&alter_query, &[]).await {
292                        Ok(_) => {}
293                        Err(err) => {
294                            // Log full error details for debugging
295                            if let Some(db_err) = err.as_db_error() {
296                                warn!(
297                                    "Error adding column {} to table {}: {} (code: {:?})",
298                                    column,
299                                    self.state_table(),
300                                    db_err,
301                                    db_err.code()
302                                );
303                            } else {
304                                warn!("Error adding column {} to table {}: {:?}", column, self.state_table(), err);
305                            }
306                            self.rebuild_columns_cache(client).await?;
307                            return Err(StateError::DDLError(Box::new(err)));
308                        }
309                    }
310                }
311            }
312
313            self.rebuild_columns_cache(client).await?;
314            Ok(())
315        }
316        .await;
317
318        // Always release the lock
319        release_ddl_lock(client, lock_key).await?;
320
321        result
322    }
323}
324
325#[async_trait]
326impl StorageCollection for PostgresBucket {
327    async fn set_state(&self, state: Attested<EntityState>) -> Result<bool, MutationError> {
328        let state_buffers = bincode::serialize(&state.payload.state.state_buffers)?;
329        let attestations: Vec<Vec<u8>> = state.attestations.iter().map(bincode::serialize).collect::<Result<Vec<_>, _>>()?;
330        let id = state.payload.entity_id;
331
332        // Ensure head is not empty for new records
333        if state.payload.state.head.is_empty() {
334            warn!("Warning: Empty head detected for entity {}", id);
335        }
336
337        let mut client = self.pool.get().await.map_err(|err| MutationError::General(err.into()))?;
338
339        let mut columns: Vec<String> = vec!["id".to_owned(), "state_buffer".to_owned(), "head".to_owned(), "attestations".to_owned()];
340        let mut params: Vec<&(dyn ToSql + Sync)> = Vec::new();
341        params.push(&id);
342        params.push(&state_buffers);
343        params.push(&state.payload.state.head);
344        params.push(&attestations);
345
346        let mut materialized: Vec<(String, Option<PGValue>)> = Vec::new();
347        let mut seen_properties = std::collections::HashSet::new();
348
349        // Process property values directly from state buffers
350        for (name, state_buffer) in state.payload.state.state_buffers.iter() {
351            let backend = backend_from_string(name, Some(state_buffer))?;
352            for (column, value) in backend.property_values() {
353                if !seen_properties.insert(column.clone()) {
354                    // Skip if property already seen in another backend
355                    // TODO: this should cause all (or subsequent?) fields with the same name
356                    // to be suffixed with the property id when we have property ids
357                    // requires some thought (and field metadata) on how to do this right
358                    continue;
359                }
360
361                let pg_value: Option<PGValue> = value.map(|value| value.into());
362                if !self.has_column(&column) {
363                    // We don't have the column yet and we know the type.
364                    if let Some(ref pg_value) = pg_value {
365                        self.add_missing_columns(&mut client, vec![(column.clone(), pg_value.postgres_type())]).await?;
366                    } else {
367                        // The column doesn't exist yet and we don't have a value.
368                        // This means the entire column is already null/none so we
369                        // don't need to set anything.
370                        continue;
371                    }
372                }
373
374                materialized.push((column.clone(), pg_value));
375            }
376        }
377
378        for (name, parameter) in &materialized {
379            columns.push(name.clone());
380
381            match &parameter {
382                Some(value) => match value {
383                    PGValue::CharacterVarying(string) => params.push(string),
384                    PGValue::SmallInt(number) => params.push(number),
385                    PGValue::Integer(number) => params.push(number),
386                    PGValue::BigInt(number) => params.push(number),
387                    PGValue::DoublePrecision(float) => params.push(float),
388                    PGValue::Bytea(bytes) => params.push(bytes),
389                    PGValue::Boolean(bool) => params.push(bool),
390                },
391                None => params.push(&UntypedNull),
392            }
393        }
394
395        let columns_str = columns.iter().map(|name| format!("\"{}\"", name)).collect::<Vec<String>>().join(", ");
396        let values_str = params.iter().enumerate().map(|(index, _)| format!("${}", index + 1)).collect::<Vec<String>>().join(", ");
397        let columns_update_str = columns
398            .iter()
399            .enumerate()
400            .skip(1) // Skip "id"
401            .map(|(index, name)| format!("\"{}\" = ${}", name, index + 1))
402            .collect::<Vec<String>>()
403            .join(", ");
404
405        // be careful with sql injection via bucket name
406        let query = format!(
407            r#"WITH old_state AS (
408                SELECT "head" FROM "{0}" WHERE "id" = $1
409            )
410            INSERT INTO "{0}"({1}) VALUES({2})
411            ON CONFLICT("id") DO UPDATE SET {3}
412            RETURNING (SELECT "head" FROM old_state) as old_head"#,
413            self.state_table(),
414            columns_str,
415            values_str,
416            columns_update_str
417        );
418
419        debug!("PostgresBucket({}).set_state: {}", self.collection_id, query);
420        let row = match client.query_one(&query, params.as_slice()).await {
421            Ok(row) => row,
422            Err(err) => {
423                let kind = error_kind(&err);
424                if let ErrorKind::UndefinedTable { table } = kind {
425                    if table == self.state_table() {
426                        self.create_state_table(&mut client).await?;
427                        return self.set_state(state).await; // retry
428                    }
429                }
430
431                return Err(StateError::DDLError(Box::new(err)).into());
432            }
433        };
434
435        // If this is a new entity (no old_head), or if the heads are different, return true
436        let old_head: Option<Clock> = row.get("old_head");
437        let changed = match old_head {
438            None => true, // New entity
439            Some(old_head) => old_head != state.payload.state.head,
440        };
441
442        debug!("PostgresBucket({}).set_state: Changed: {}", self.collection_id, changed);
443        Ok(changed)
444    }
445
446    async fn get_state(&self, id: EntityId) -> Result<Attested<EntityState>, RetrievalError> {
447        // be careful with sql injection via bucket name
448        let query = format!(r#"SELECT "id", "state_buffer", "head", "attestations" FROM "{}" WHERE "id" = $1"#, self.state_table());
449
450        let mut client = match self.pool.get().await {
451            Ok(client) => client,
452            Err(err) => {
453                return Err(RetrievalError::StorageError(err.into()));
454            }
455        };
456
457        debug!("PostgresBucket({}).get_state: {}", self.collection_id, query);
458        let rows = match client.query(&query, &[&id]).await {
459            Ok(rows) => rows,
460            Err(err) => {
461                let kind = error_kind(&err);
462                if let ErrorKind::UndefinedTable { table } = kind {
463                    if table == self.state_table() {
464                        self.create_state_table(&mut client).await.map_err(|e| RetrievalError::StorageError(e.into()))?;
465                        return Err(RetrievalError::EntityNotFound(id));
466                    }
467                }
468                return Err(RetrievalError::StorageError(err.into()));
469            }
470        };
471
472        let row = match rows.into_iter().next() {
473            Some(row) => row,
474            None => return Err(RetrievalError::EntityNotFound(id)),
475        };
476
477        debug!("PostgresBucket({}).get_state: Row: {:?}", self.collection_id, row);
478        let row_id: EntityId = row.try_get("id").map_err(RetrievalError::storage)?;
479        assert_eq!(row_id, id);
480
481        let serialized_buffers: Vec<u8> = row.try_get("state_buffer").map_err(RetrievalError::storage)?;
482        let state_buffers: BTreeMap<String, Vec<u8>> = bincode::deserialize(&serialized_buffers).map_err(RetrievalError::storage)?;
483        let head: Clock = row.try_get("head").map_err(RetrievalError::storage)?;
484        let attestation_bytes: Vec<Vec<u8>> = row.try_get("attestations").map_err(RetrievalError::storage)?;
485        let attestations = attestation_bytes
486            .into_iter()
487            .map(|bytes| bincode::deserialize(&bytes))
488            .collect::<Result<Vec<Attestation>, _>>()
489            .map_err(RetrievalError::storage)?;
490
491        Ok(Attested {
492            payload: EntityState {
493                entity_id: id,
494                collection: self.collection_id.clone(),
495                state: State { state_buffers: StateBuffers(state_buffers), head },
496            },
497            attestations: AttestationSet(attestations),
498        })
499    }
500
501    async fn fetch_states(&self, selection: &ankql::ast::Selection) -> Result<Vec<Attested<EntityState>>, RetrievalError> {
502        debug!("fetch_states: {:?}", selection);
503        let client = self.pool.get().await.map_err(|err| RetrievalError::StorageError(Box::new(err)))?;
504
505        let mut results = Vec::new();
506        let mut builder = SqlBuilder::with_fields(vec!["id", "state_buffer", "head", "attestations"]);
507        builder.table_name(self.state_table());
508        builder.selection(selection)?;
509
510        let (sql, args) = builder.build()?;
511        debug!("PostgresBucket({}).fetch_states: SQL: {} with args: {:?}", self.collection_id, sql, args);
512
513        let stream = match client.query_raw(&sql, args).await {
514            Ok(stream) => stream,
515            Err(err) => {
516                let kind = error_kind(&err);
517                match kind {
518                    ErrorKind::UndefinedTable { table } => {
519                        if table == self.state_table() {
520                            // Table doesn't exist yet, return empty results
521                            return Ok(Vec::new());
522                        }
523                    }
524                    ErrorKind::UndefinedColumn { table, column } => {
525                        // this means we didn't write the column yet, which suggests that all values are null
526                        // So we can recompute the predicate to treat this column as always NULL and retry
527                        debug!("Undefined column: {} in table: {:?}, {}", column, table, self.state_table());
528                        let new_selection = selection.assume_null(&[column]);
529                        return self.fetch_states(&new_selection).await;
530                    }
531                    _ => {}
532                }
533
534                return Err(RetrievalError::StorageError(err.into()));
535            }
536        };
537        pin_mut!(stream);
538
539        while let Some(row) = stream.try_next().await.map_err(RetrievalError::storage)? {
540            let id: EntityId = row.try_get(0).map_err(RetrievalError::storage)?;
541            let state_buffer: Vec<u8> = row.try_get(1).map_err(RetrievalError::storage)?;
542            let state_buffers: BTreeMap<String, Vec<u8>> = bincode::deserialize(&state_buffer).map_err(RetrievalError::storage)?;
543            let head: Clock = row.try_get("head").map_err(RetrievalError::storage)?;
544            let attestation_bytes: Vec<Vec<u8>> = row.try_get("attestations").map_err(RetrievalError::storage)?;
545            let attestations = attestation_bytes
546                .into_iter()
547                .map(|bytes| bincode::deserialize(&bytes))
548                .collect::<Result<Vec<Attestation>, _>>()
549                .map_err(RetrievalError::storage)?;
550
551            results.push(Attested {
552                payload: EntityState {
553                    entity_id: id,
554                    collection: self.collection_id.clone(),
555                    state: State { state_buffers: StateBuffers(state_buffers), head },
556                },
557                attestations: AttestationSet(attestations),
558            });
559        }
560
561        Ok(results)
562    }
563
564    async fn add_event(&self, entity_event: &Attested<Event>) -> Result<bool, MutationError> {
565        let operations = bincode::serialize(&entity_event.payload.operations)?;
566        let attestations = bincode::serialize(&entity_event.attestations)?;
567
568        let query = format!(
569            r#"INSERT INTO "{0}"("id", "entity_id", "operations", "parent", "attestations") VALUES($1, $2, $3, $4, $5)
570               ON CONFLICT ("id") DO NOTHING"#,
571            self.event_table(),
572        );
573
574        let mut client = self.pool.get().await.map_err(|err| MutationError::General(err.into()))?;
575        debug!("PostgresBucket({}).add_event: {}", self.collection_id, query);
576        let affected = match client
577            .execute(
578                &query,
579                &[&entity_event.payload.id(), &entity_event.payload.entity_id, &operations, &entity_event.payload.parent, &attestations],
580            )
581            .await
582        {
583            Ok(affected) => affected,
584            Err(err) => {
585                let kind = error_kind(&err);
586                match kind {
587                    ErrorKind::UndefinedTable { table } => {
588                        if table == self.event_table() {
589                            self.create_event_table(&mut client).await?;
590                            return self.add_event(entity_event).await; // retry
591                        }
592                    }
593                    _ => {
594                        error!("PostgresBucket({}).add_event: Error: {:?}", self.collection_id, err);
595                    }
596                }
597
598                return Err(StateError::DMLError(Box::new(err)).into());
599            }
600        };
601
602        Ok(affected > 0)
603    }
604
605    async fn get_events(&self, event_ids: Vec<EventId>) -> Result<Vec<Attested<Event>>, RetrievalError> {
606        if event_ids.is_empty() {
607            return Ok(Vec::new());
608        }
609
610        let query = format!(
611            r#"SELECT "id", "entity_id", "operations", "parent", "attestations" FROM "{0}" WHERE "id" = ANY($1)"#,
612            self.event_table(),
613        );
614
615        let client = self.pool.get().await.map_err(RetrievalError::storage)?;
616        let rows = match client.query(&query, &[&event_ids]).await {
617            Ok(rows) => rows,
618            Err(err) => {
619                let kind = error_kind(&err);
620                match kind {
621                    ErrorKind::UndefinedTable { table } if table == self.event_table() => return Ok(Vec::new()),
622                    _ => return Err(RetrievalError::storage(err)),
623                }
624            }
625        };
626
627        let mut events = Vec::new();
628        for row in rows {
629            let entity_id: EntityId = row.try_get("entity_id").map_err(RetrievalError::storage)?;
630            let operations: OperationSet = row.try_get("operations").map_err(RetrievalError::storage)?;
631            let parent: Clock = row.try_get("parent").map_err(RetrievalError::storage)?;
632            let attestations_binary: Vec<u8> = row.try_get("attestations").map_err(RetrievalError::storage)?;
633            let attestations: Vec<Attestation> = bincode::deserialize(&attestations_binary).map_err(RetrievalError::storage)?;
634
635            let event = Attested {
636                payload: Event { collection: self.collection_id.clone(), entity_id, operations, parent },
637                attestations: AttestationSet(attestations),
638            };
639            events.push(event);
640        }
641        Ok(events)
642    }
643
644    async fn dump_entity_events(&self, entity_id: EntityId) -> Result<Vec<Attested<Event>>, ankurah_core::error::RetrievalError> {
645        let query =
646            format!(r#"SELECT "id", "operations", "parent", "attestations" FROM "{0}" WHERE "entity_id" = $1"#, self.event_table(),);
647
648        let client = self.pool.get().await.map_err(RetrievalError::storage)?;
649        debug!("PostgresBucket({}).get_events: {}", self.collection_id, query);
650        let rows = match client.query(&query, &[&entity_id]).await {
651            Ok(rows) => rows,
652            Err(err) => {
653                let kind = error_kind(&err);
654                if let ErrorKind::UndefinedTable { table } = kind {
655                    if table == self.event_table() {
656                        return Ok(Vec::new());
657                    }
658                }
659
660                return Err(RetrievalError::storage(err));
661            }
662        };
663
664        let mut events = Vec::new();
665        for row in rows {
666            // let event_id: EventId = row.try_get("id").map_err(|err| RetrievalError::storage(err))?;
667            let operations_binary: Vec<u8> = row.try_get("operations").map_err(RetrievalError::storage)?;
668            let operations = bincode::deserialize(&operations_binary).map_err(RetrievalError::storage)?;
669            let parent: Clock = row.try_get("parent").map_err(RetrievalError::storage)?;
670            let attestations_binary: Vec<u8> = row.try_get("attestations").map_err(RetrievalError::storage)?;
671            let attestations: Vec<Attestation> = bincode::deserialize(&attestations_binary).map_err(RetrievalError::storage)?;
672
673            events.push(Attested {
674                payload: Event { collection: self.collection_id.clone(), entity_id, operations, parent },
675                attestations: AttestationSet(attestations),
676            });
677        }
678
679        Ok(events)
680    }
681}
682
683// Some hacky shit because rust-postgres doesn't let us ask for the error kind
684// TODO: remove this when https://github.com/sfackler/rust-postgres/pull/1185
685//       gets merged
686#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
687pub enum ErrorKind {
688    RowCount,
689    UndefinedTable { table: String },
690    UndefinedColumn { table: Option<String>, column: String },
691    Unknown,
692    PostgresError(String),
693}
694
695pub fn error_kind(err: &tokio_postgres::Error) -> ErrorKind {
696    let string = err.as_db_error().map(|e| e.message()).unwrap_or_default().trim().to_owned();
697    let _db_error = err.as_db_error();
698    let sql_code = err.code().cloned();
699
700    // Check the error's Display string for RowCount errors (client-side, not db error)
701    let err_string = err.to_string();
702    if err_string.contains("query returned an unexpected number of rows") || string == "query returned an unexpected number of rows" {
703        return ErrorKind::RowCount;
704    }
705
706    // Useful for adding new errors
707    // error!("postgres error: {:?}", err);
708    // error!("db_err: {:?}", err.as_db_error());
709    // error!("sql_code: {:?}", err.code());
710    // error!("err: {:?}", err);
711    // error!("err: {:?}", err.to_string());
712    debug!("postgres error: {:?}", err);
713
714    let quote_indices = |s: &str| {
715        let mut quotes = Vec::new();
716        for (index, char) in s.char_indices() {
717            if char == '"' {
718                quotes.push(index)
719            }
720        }
721        quotes
722    };
723
724    match sql_code {
725        Some(SqlState::UNDEFINED_TABLE) => {
726            // relation "album" does not exist
727            let quotes = quote_indices(&string);
728            if quotes.len() >= 2 {
729                let table = &string[quotes[0] + 1..quotes[1]];
730                ErrorKind::UndefinedTable { table: table.to_owned() }
731            } else {
732                ErrorKind::PostgresError(string.clone())
733            }
734        }
735        Some(SqlState::UNDEFINED_COLUMN) => {
736            // Handle both formats:
737            // "column "name" of relation "album" does not exist"
738            // "column "status" does not exist"
739            let quotes = quote_indices(&string);
740            if quotes.len() >= 2 {
741                let column = string[quotes[0] + 1..quotes[1]].to_owned();
742
743                let table = if quotes.len() >= 4 {
744                    // Full format with table name
745                    Some(string[quotes[2] + 1..quotes[3]].to_owned())
746                } else {
747                    // Short format without table name
748                    None
749                };
750
751                ErrorKind::UndefinedColumn { table, column }
752            } else {
753                ErrorKind::PostgresError(string.clone())
754            }
755        }
756        _ => ErrorKind::Unknown,
757    }
758}
759
760#[allow(unused)]
761pub struct MissingMaterialized {
762    pub name: String,
763}
764
765use bytes::BytesMut;
766use tokio_postgres::types::{to_sql_checked, IsNull, Type};
767
768use crate::sql_builder::SqlBuilder;
769
770#[derive(Debug)]
771struct UntypedNull;
772
773impl ToSql for UntypedNull {
774    fn to_sql(&self, _ty: &Type, _out: &mut BytesMut) -> Result<IsNull, Box<dyn std::error::Error + Sync + Send>> { Ok(IsNull::Yes) }
775
776    fn accepts(_ty: &Type) -> bool {
777        true // Accept all types
778    }
779
780    to_sql_checked!();
781}
ankurah_storage_postgres/lib.rs

ankurah_storage_postgres/
lib.rs
