yugendb_postgres/
driver.rs

1//! PostgreSQL driver implementation.
2//!
3//! PostgreSQL is used as a SQL-backed implementation detail for the yugendb
4//! storage model. The public API remains namespace, collection, key, and
5//! serialised value based, so application code can stay driver-neutral.
6
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use parking_lot::Mutex;
10use sqlx::{postgres::PgPoolOptions, PgPool, Row};
11use yugendb_core::{
12    Batch, BatchOperation, BatchResult, Capabilities, CollectionName, DeleteOptions, Driver, Key,
13    Namespace, ReadOptions, Result, ScanOptions, StoredValue, ValueBytes, ValueMetadata,
14    WriteOptions, YugenDbError,
15};
16
17use crate::schema::{schema_statements, PostgresDriverStorageSchema};
18
19/// Options used to configure the PostgreSQL driver.
20#[derive(Debug, Clone, PartialEq, Eq)]
21pub struct PostgresDriverOptions {
22    /// Backend connection string.
23    pub connection_string: String,
24    /// Whether the storage schema should be created during initialisation.
25    pub create_schema_on_initialise: bool,
26    /// Maximum number of pooled connections.
27    pub max_connections: u32,
28}
29
30impl PostgresDriverOptions {
31    /// Creates default options for a connection string.
32    #[must_use]
33    pub fn new(connection_string: impl Into<String>) -> Self {
34        Self {
35            connection_string: connection_string.into(),
36            create_schema_on_initialise: true,
37            max_connections: 5,
38        }
39    }
40}
41
42/// PostgreSQL driver for durable external-service storage.
43#[derive(Debug)]
44pub struct PostgresDriver {
45    options: PostgresDriverOptions,
46    pool: Mutex<Option<PgPool>>,
47}
48
49impl PostgresDriver {
50    /// Creates a new PostgreSQL driver using default options.
51    #[must_use]
52    pub fn new(connection_string: impl Into<String>) -> Self {
53        Self::with_options(PostgresDriverOptions::new(connection_string))
54    }
55
56    /// Creates a new PostgreSQL driver with explicit options.
57    #[must_use]
58    pub fn with_options(options: PostgresDriverOptions) -> Self {
59        Self {
60            options,
61            pool: Mutex::new(None),
62        }
63    }
64
65    /// Returns the configured options.
66    #[must_use]
67    pub const fn options(&self) -> &PostgresDriverOptions {
68        &self.options
69    }
70
71    /// Returns the storage schema.
72    #[must_use]
73    pub const fn schema(&self) -> PostgresDriverStorageSchema {
74        PostgresDriverStorageSchema::current()
75    }
76
77    fn map_error(error: sqlx::Error) -> YugenDbError {
78        match error {
79            sqlx::Error::RowNotFound => YugenDbError::not_found("PostgreSQL row was not found"),
80            sqlx::Error::PoolTimedOut | sqlx::Error::Io(_) | sqlx::Error::Tls(_) => {
81                YugenDbError::connection_error(error.to_string())
82            }
83            other => YugenDbError::driver_error("postgres", other.to_string()),
84        }
85    }
86
87    fn pool(&self) -> Result<PgPool> {
88        self.pool.lock().as_ref().cloned().ok_or_else(|| {
89            YugenDbError::connection_error("PostgreSQL driver has not been initialised")
90        })
91    }
92
93    fn parse_timestamp(value: String) -> Result<DateTime<Utc>> {
94        DateTime::parse_from_rfc3339(&value)
95            .map(|value| value.with_timezone(&Utc))
96            .map_err(|error| YugenDbError::driver_error("postgres", error.to_string()))
97    }
98
99    fn parse_optional_timestamp(value: Option<String>) -> Result<Option<DateTime<Utc>>> {
100        value.map(Self::parse_timestamp).transpose()
101    }
102
103    fn stored_value_from_row(row: &sqlx::postgres::PgRow) -> Result<StoredValue> {
104        let bytes: Vec<u8> = row.try_get("value").map_err(Self::map_error)?;
105        let codec: String = row.try_get("codec").map_err(Self::map_error)?;
106        let created_at: String = row.try_get("created_at").map_err(Self::map_error)?;
107        let updated_at: String = row.try_get("updated_at").map_err(Self::map_error)?;
108        let expires_at: Option<String> = row.try_get("expires_at").map_err(Self::map_error)?;
109        Ok(StoredValue::new(
110            ValueBytes::from(bytes),
111            ValueMetadata {
112                codec,
113                created_at: Self::parse_timestamp(created_at)?,
114                updated_at: Self::parse_timestamp(updated_at)?,
115                expires_at: Self::parse_optional_timestamp(expires_at)?,
116            },
117        ))
118    }
119
120    async fn insert_value(
121        executor: impl sqlx::Executor<'_, Database = sqlx::Postgres>,
122        namespace: &Namespace,
123        collection: &CollectionName,
124        key: &Key,
125        value: &StoredValue,
126    ) -> Result<()> {
127        sqlx::query(
128            r#"INSERT INTO yugendb_store
129(namespace, collection, key, value, codec, created_at, updated_at, expires_at)
130VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
131ON CONFLICT(namespace, collection, key) DO UPDATE SET
132    value = excluded.value,
133    codec = excluded.codec,
134    updated_at = excluded.updated_at,
135    expires_at = excluded.expires_at"#,
136        )
137        .bind(namespace.as_str())
138        .bind(collection.as_str())
139        .bind(key.as_str())
140        .bind(value.bytes.as_slice())
141        .bind(value.metadata.codec.as_str())
142        .bind(value.metadata.created_at.to_rfc3339())
143        .bind(value.metadata.updated_at.to_rfc3339())
144        .bind(value.metadata.expires_at.as_ref().map(DateTime::to_rfc3339))
145        .execute(executor)
146        .await
147        .map_err(Self::map_error)?;
148        Ok(())
149    }
150}
151
152impl Clone for PostgresDriver {
153    fn clone(&self) -> Self {
154        Self::with_options(self.options.clone())
155    }
156}
157
158/// Capabilities for the implemented PostgreSQL driver.
159#[must_use]
160pub const fn postgres_capabilities() -> Capabilities {
161    Capabilities {
162        transactions: true,
163        ttl: true,
164        prefix_scan: true,
165        atomic_increment: false,
166        batch_write: true,
167        raw_sql: false,
168        document_query: false,
169        json_query: false,
170        migrations: true,
171        connection_pooling: true,
172        watch: false,
173        backup: false,
174    }
175}
176
177
178
179/// Convenience constructor for the PostgreSQL driver.
180#[must_use]
181pub fn postgres(connection_string: impl Into<String>) -> PostgresDriver {
182    PostgresDriver::new(connection_string)
183}
184
185#[async_trait]
186impl Driver for PostgresDriver {
187    fn name(&self) -> &'static str {
188        "postgres"
189    }
190
191    fn capabilities(&self) -> Capabilities {
192        postgres_capabilities()
193    }
194
195    async fn initialise(&self) -> Result<()> {
196        let pool = PgPoolOptions::new()
197            .max_connections(self.options.max_connections)
198            .connect(&self.options.connection_string)
199            .await
200            .map_err(Self::map_error)?;
201        if self.options.create_schema_on_initialise {
202            for statement in schema_statements() {
203                sqlx::query(statement)
204                    .execute(&pool)
205                    .await
206                    .map_err(Self::map_error)?;
207            }
208        }
209        *self.pool.lock() = Some(pool);
210        Ok(())
211    }
212
213    async fn finalise(&self) -> Result<()> {
214        let pool = {
215            let mut guard = self.pool.lock();
216            guard.take()
217        };
218
219        if let Some(pool) = pool {
220            pool.close().await;
221        }
222
223        Ok(())
224    }
225
226    async fn get(
227        &self,
228        namespace: &Namespace,
229        collection: &CollectionName,
230        key: &Key,
231        options: &ReadOptions,
232    ) -> Result<Option<StoredValue>> {
233        let pool = self.pool()?;
234        let row = sqlx::query("SELECT value, codec, created_at, updated_at, expires_at FROM yugendb_store WHERE namespace = $1 AND collection = $2 AND key = $3")
235            .bind(namespace.as_str()).bind(collection.as_str()).bind(key.as_str())
236            .fetch_optional(&pool).await.map_err(Self::map_error)?;
237        let Some(row) = row else {
238            return Ok(None);
239        };
240        let value = Self::stored_value_from_row(&row)?;
241        Ok((options.allow_expired || !value.is_expired()).then_some(value))
242    }
243
244    async fn set(
245        &self,
246        namespace: &Namespace,
247        collection: &CollectionName,
248        key: &Key,
249        value: StoredValue,
250        options: &WriteOptions,
251    ) -> Result<()> {
252        if !options.overwrite
253            && self
254                .get(namespace, collection, key, &ReadOptions::default())
255                .await?
256                .is_some()
257        {
258            return Err(YugenDbError::conflict(
259                "PostgreSQL driver refused to overwrite an existing value",
260            ));
261        }
262        let pool = self.pool()?;
263        Self::insert_value(&pool, namespace, collection, key, &value).await
264    }
265
266    async fn delete(
267        &self,
268        namespace: &Namespace,
269        collection: &CollectionName,
270        key: &Key,
271        options: &DeleteOptions,
272    ) -> Result<bool> {
273        let pool = self.pool()?;
274        let removed = sqlx::query(
275            "DELETE FROM yugendb_store WHERE namespace = $1 AND collection = $2 AND key = $3",
276        )
277        .bind(namespace.as_str())
278        .bind(collection.as_str())
279        .bind(key.as_str())
280        .execute(&pool)
281        .await
282        .map_err(Self::map_error)?
283        .rows_affected()
284            > 0;
285        if options.must_exist && !removed {
286            return Err(YugenDbError::not_found(
287                "PostgreSQL driver could not delete a missing value",
288            ));
289        }
290        Ok(removed)
291    }
292
293    async fn exists(
294        &self,
295        namespace: &Namespace,
296        collection: &CollectionName,
297        key: &Key,
298    ) -> Result<bool> {
299        Ok(self
300            .get(namespace, collection, key, &ReadOptions::default())
301            .await?
302            .is_some())
303    }
304
305    async fn scan_prefix(
306        &self,
307        namespace: &Namespace,
308        collection: &CollectionName,
309        options: &ScanOptions,
310    ) -> Result<Vec<(Key, StoredValue)>> {
311        let pool = self.pool()?;
312        let prefix = options.prefix.as_deref().unwrap_or("");
313        let rows = sqlx::query("SELECT key, value, codec, created_at, updated_at, expires_at FROM yugendb_store WHERE namespace = $1 AND collection = $2 AND key LIKE $3 ORDER BY key ASC")
314            .bind(namespace.as_str()).bind(collection.as_str()).bind(format!("{prefix}%"))
315            .fetch_all(&pool).await.map_err(Self::map_error)?;
316        let mut values = Vec::new();
317        for row in rows {
318            let key_text: String = row.try_get("key").map_err(Self::map_error)?;
319            let value = Self::stored_value_from_row(&row)?;
320            if value.is_expired() && !options.include_expired {
321                continue;
322            }
323            values.push((Key::try_from(key_text)?, value));
324            if let Some(limit) = options.limit {
325                if values.len() >= limit {
326                    break;
327                }
328            }
329        }
330        Ok(values)
331    }
332
333    async fn batch(&self, batch: Batch) -> Result<BatchResult> {
334        let pool = self.pool()?;
335        let mut transaction = pool.begin().await.map_err(Self::map_error)?;
336        let mut set_count = 0;
337        let mut delete_count = 0;
338        for operation in batch.into_operations() {
339            match operation {
340                BatchOperation::Set {
341                    namespace,
342                    collection,
343                    key,
344                    value,
345                } => {
346                    Self::insert_value(&mut *transaction, &namespace, &collection, &key, &value)
347                        .await?;
348                    set_count += 1;
349                }
350                BatchOperation::Delete {
351                    namespace,
352                    collection,
353                    key,
354                } => {
355                    let removed = sqlx::query("DELETE FROM yugendb_store WHERE namespace = $1 AND collection = $2 AND key = $3")
356                        .bind(namespace.as_str()).bind(collection.as_str()).bind(key.as_str())
357                        .execute(&mut *transaction).await.map_err(Self::map_error)?.rows_affected();
358                    if removed > 0 {
359                        delete_count += 1;
360                    }
361                }
362            }
363        }
364        transaction.commit().await.map_err(Self::map_error)?;
365        Ok(BatchResult::new(set_count, delete_count))
366    }
367}
yugendb_postgres/driver.rs

yugendb_postgres/
driver.rs
