nstreams-postgres 0.2.2

use async_trait::async_trait;
use nstreams_core::event::{EventVersion, StreamEvent};
use nstreams_core::namespace::{Namespace, NamespaceResources};
use nstreams_core::store::EventStore;
use sqlx::postgres::PgPoolOptions;
use sqlx::{PgPool, Row};
use thiserror::Error;

#[derive(Debug, Error)]
pub enum PostgresError {
    #[error(transparent)]
    Sqlx(#[from] sqlx::Error),
}

impl From<PostgresError> for nstreams_core::Error {
    fn from(value: PostgresError) -> Self {
        Self::Store(value.to_string())
    }
}

fn store_err(error: sqlx::Error) -> nstreams_core::Error {
    PostgresError::from(error).into()
}

fn row_to_event(
    namespace: &str,
    version: i64,
    payload: serde_json::Value,
    filter_value: Option<String>,
) -> nstreams_core::Result<StreamEvent> {
    Ok(StreamEvent {
        namespace: namespace.to_string(),
        version: version as EventVersion,
        payload: serde_json::to_vec(&payload)
            .map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?,
        filter_value,
    })
}

/// PostgreSQL-backed event store with per-namespace tables.
#[derive(Clone)]
pub struct PostgresEventStore {
    pool: PgPool,
}

impl PostgresEventStore {
    pub async fn connect(database_url: &str) -> Result<Self, PostgresError> {
        let pool = PgPoolOptions::new()
            .max_connections(10)
            .connect(database_url)
            .await?;
        Ok(Self { pool })
    }

    pub fn from_pool(pool: PgPool) -> Self {
        Self { pool }
    }

    pub fn pool(&self) -> &PgPool {
        &self.pool
    }
}

#[async_trait]
impl EventStore for PostgresEventStore {
    async fn ensure_namespace<N: Namespace>(&self, namespace: &N) -> nstreams_core::Result<()> {
        let resources = NamespaceResources::new(namespace);
        let table = resources.postgres_table;

        let ddl = format!(
            r#"
            CREATE TABLE IF NOT EXISTS {table} (
                version BIGINT PRIMARY KEY,
                payload JSONB NOT NULL,
                filter_value TEXT,
                created_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
            );
            CREATE INDEX IF NOT EXISTS idx_{table}_version_desc ON {table} (version DESC);
            ALTER TABLE {table} ADD COLUMN IF NOT EXISTS filter_value TEXT;
            "#
        );

        sqlx::raw_sql(&ddl).execute(&self.pool).await.map_err(store_err)?;
        Ok(())
    }

    async fn persist_next<N: Namespace>(
        &self,
        namespace: &N,
        payload: &[u8],
        filter_value: Option<&str>,
    ) -> nstreams_core::Result<StreamEvent> {
        let resources = NamespaceResources::new(namespace);
        let table = resources.postgres_table;
        let namespace = resources.namespace;

        let json: serde_json::Value = serde_json::from_slice(payload)
            .map_err(|e| nstreams_core::Error::Serialization(e.to_string()))?;

        let mut tx = self.pool.begin().await.map_err(store_err)?;

        let lock_sql = format!("SELECT pg_advisory_xact_lock(hashtext($1))");
        sqlx::query(&lock_sql)
            .bind(&namespace)
            .execute(&mut *tx)
            .await
            .map_err(store_err)?;

        let next_version_sql = format!(
            "SELECT COALESCE(MAX(version), 0) + 1 AS version FROM {table}"
        );
        let row = sqlx::query(&next_version_sql)
            .fetch_one(&mut *tx)
            .await
            .map_err(store_err)?;
        let version: i64 = row.get("version");

        let insert_sql = format!(
            "INSERT INTO {table} (version, payload, filter_value) VALUES ($1, $2, $3)"
        );
        sqlx::query(&insert_sql)
            .bind(version)
            .bind(&json)
            .bind(filter_value)
            .execute(&mut *tx)
            .await
            .map_err(store_err)?;

        tx.commit().await.map_err(store_err)?;

        Ok(StreamEvent {
            namespace,
            version: version as EventVersion,
            payload: payload.to_vec(),
            filter_value: filter_value.map(str::to_string),
        })
    }

    async fn load_history<N: Namespace>(
        &self,
        namespace: &N,
        limit: u64,
    ) -> nstreams_core::Result<Vec<StreamEvent>> {
        let resources = NamespaceResources::new(namespace);
        let table = resources.postgres_table;
        let namespace = resources.namespace;

        let sql = format!(
            r#"
            SELECT version, payload, filter_value
            FROM {table}
            ORDER BY version DESC
            LIMIT $1
            "#
        );

        let rows = sqlx::query(&sql)
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await
            .map_err(store_err)?;

        let mut events = rows
            .into_iter()
            .map(|row| {
                let version: i64 = row.get("version");
                let payload: serde_json::Value = row.get("payload");
                let filter_value: Option<String> = row.get("filter_value");
                row_to_event(&namespace, version, payload, filter_value)
            })
            .collect::<nstreams_core::Result<Vec<_>>>()?;

        events.sort_by_key(|event| event.version);
        Ok(events)
    }

    async fn load_after_version<N: Namespace>(
        &self,
        namespace: &N,
        after_version: EventVersion,
        limit: u64,
    ) -> nstreams_core::Result<Vec<StreamEvent>> {
        let resources = NamespaceResources::new(namespace);
        let table = resources.postgres_table;
        let namespace = resources.namespace;

        let sql = format!(
            r#"
            SELECT version, payload, filter_value
            FROM {table}
            WHERE version > $1
            ORDER BY version ASC
            LIMIT $2
            "#
        );

        let rows = sqlx::query(&sql)
            .bind(after_version as i64)
            .bind(limit as i64)
            .fetch_all(&self.pool)
            .await
            .map_err(store_err)?;

        rows.into_iter()
            .map(|row| {
                let version: i64 = row.get("version");
                let payload: serde_json::Value = row.get("payload");
                let filter_value: Option<String> = row.get("filter_value");
                row_to_event(&namespace, version, payload, filter_value)
            })
            .collect()
    }

    async fn latest_version<N: Namespace>(
        &self,
        namespace: &N,
    ) -> nstreams_core::Result<Option<EventVersion>> {
        let resources = NamespaceResources::new(namespace);
        let table = resources.postgres_table;

        let sql = format!("SELECT MAX(version) AS version FROM {table}");
        let row = sqlx::query(&sql)
            .fetch_optional(&self.pool)
            .await
            .map_err(store_err)?;
        Ok(row.and_then(|row| {
            let version: Option<i64> = row.get("version");
            version.map(|v| v as EventVersion)
        }))
    }
}