rusticx-postgres 1.0.0

use async_trait::async_trait;
use rusticx_core::{
    adapter::DatabaseAdapter,
    error::{Result, RusticxError},
    model::TableSchema,
    query::QueryBuilder,
    value::{Row, Value},
};
use rusticx_sql::{compiler::SqlCompiler, dialect::PostgresDialect};
use sqlx::{postgres::PgPoolOptions, PgPool, Row as SqlxRow};
use tracing::debug;

use crate::convert::pg_row_to_row;

/// Async PostgreSQL adapter backed by an `sqlx` connection pool.
///
/// Obtain one via [`PostgresAdapter::connect_url`] (quick) or
/// [`PostgresAdapter::connect`] (full [`PostgresConfig`] options).
pub struct PostgresAdapter {
    pool: PgPool,
    dialect: PostgresDialect,
}

/// Configuration for the PostgreSQL connection pool.
///
/// ```rust,ignore
/// let adapter = PostgresAdapter::connect(
///     PostgresConfig::new("postgres://user:pass@localhost/mydb")
///         .max_connections(20)
///         .min_connections(2)
/// ).await?;
/// ```
#[derive(Debug, Clone)]
pub struct PostgresConfig {
    pub url: String,
    pub max_connections: u32,
    pub min_connections: u32,
}

impl PostgresConfig {
    pub fn new(url: impl Into<String>) -> Self {
        Self {
            url: url.into(),
            max_connections: 10,
            min_connections: 1,
        }
    }

    pub fn max_connections(mut self, n: u32) -> Self {
        self.max_connections = n;
        self
    }

    pub fn min_connections(mut self, n: u32) -> Self {
        self.min_connections = n;
        self
    }
}

impl PostgresAdapter {
    pub async fn connect(config: PostgresConfig) -> Result<Self> {
        let pool = PgPoolOptions::new()
            .max_connections(config.max_connections)
            .min_connections(config.min_connections)
            .connect(&config.url)
            .await
            .map_err(|e| RusticxError::Connection(e.to_string()))?;

        Ok(Self { pool, dialect: PostgresDialect })
    }

    pub async fn connect_url(url: impl Into<String>) -> Result<Self> {
        Self::connect(PostgresConfig::new(url)).await
    }

    fn compiler(&self) -> SqlCompiler<'_, PostgresDialect> {
        SqlCompiler::new(&self.dialect)
    }

    /// Bind rusticx Values onto a sqlx PgQuery.
    async fn execute_with_bindings(&self, sql: &str, bindings: Vec<Value>) -> Result<u64> {
        let mut q = sqlx::query(sql);
        for val in bindings {
            q = bind_value(q, val);
        }
        let res = q.execute(&self.pool).await.map_err(|e| RusticxError::Query(e.to_string()))?;
        Ok(res.rows_affected())
    }

    async fn fetch_rows_with_bindings(&self, sql: &str, bindings: Vec<Value>) -> Result<Vec<Row>> {
        let mut q = sqlx::query(sql);
        for val in bindings {
            q = bind_value(q, val);
        }
        let rows = q.fetch_all(&self.pool).await.map_err(|e| RusticxError::Query(e.to_string()))?;
        Ok(rows.into_iter().map(pg_row_to_row).collect())
    }
}

fn bind_value<'q>(
    q: sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments>,
    val: Value,
) -> sqlx::query::Query<'q, sqlx::Postgres, sqlx::postgres::PgArguments> {
    match val {
        Value::Null => q.bind(Option::<String>::None),
        Value::Bool(b) => q.bind(b),
        Value::Int(i) => q.bind(i),
        Value::Float(f) => q.bind(f),
        Value::Text(s) => q.bind(s),
        Value::Bytes(b) => q.bind(b),
        Value::Uuid(u) => q.bind(u),
        Value::DateTime(dt) => q.bind(dt),
        Value::Json(j) => q.bind(j),
        Value::Array(arr) => {
            let json = serde_json::to_value(&arr).unwrap_or(serde_json::Value::Null);
            q.bind(json)
        }
        Value::Map(m) => {
            let json = serde_json::to_value(&m).unwrap_or(serde_json::Value::Null);
            q.bind(json)
        }
    }
}

#[async_trait]
impl DatabaseAdapter for PostgresAdapter {
    fn name(&self) -> &'static str {
        "postgres"
    }

    async fn ping(&self) -> Result<()> {
        sqlx::query("SELECT 1")
            .execute(&self.pool)
            .await
            .map_err(|e| RusticxError::Connection(e.to_string()))?;
        Ok(())
    }

    async fn close(&self) -> Result<()> {
        self.pool.close().await;
        Ok(())
    }

    async fn create_table(&self, schema: &TableSchema) -> Result<()> {
        let sql = self.compiler().create_table(schema);
        debug!(sql = %sql, "create_table");
        // May contain multiple statements (table + indexes), execute each
        for stmt in sql.split(';').map(str::trim).filter(|s| !s.is_empty()) {
            sqlx::query(stmt)
                .execute(&self.pool)
                .await
                .map_err(|e| RusticxError::Schema(e.to_string()))?;
        }
        Ok(())
    }

    async fn drop_table(&self, table: &str) -> Result<()> {
        let sql = self.compiler().drop_table(table);
        sqlx::query(&sql)
            .execute(&self.pool)
            .await
            .map_err(|e| RusticxError::Schema(e.to_string()))?;
        Ok(())
    }

    async fn table_exists(&self, table: &str) -> Result<bool> {
        let row = sqlx::query(
            "SELECT COUNT(*) as cnt FROM information_schema.tables WHERE table_name = $1",
        )
        .bind(table)
        .fetch_one(&self.pool)
        .await
        .map_err(|e| RusticxError::Query(e.to_string()))?;
        let cnt: i64 = row.try_get("cnt").unwrap_or(0);
        Ok(cnt > 0)
    }

    async fn insert(&self, table: &str, row: Row) -> Result<Row> {
        let pairs: Vec<(String, Value)> = row.into_iter().collect();
        let (sql, bindings) = self.compiler().insert(table, &pairs);
        debug!(sql = %sql, "insert");

        let mut q = sqlx::query(&sql);
        for val in bindings {
            q = bind_value(q, val);
        }

        let pg_row = q
            .fetch_one(&self.pool)
            .await
            .map_err(|e| RusticxError::Query(e.to_string()))?;

        Ok(pg_row_to_row(pg_row))
    }

    async fn insert_many(&self, table: &str, rows: Vec<Row>) -> Result<u64> {
        let mut tx = self.pool.begin().await.map_err(|e| RusticxError::Transaction(e.to_string()))?;
        let mut count = 0u64;
        for row in rows {
            let pairs: Vec<(String, Value)> = row.into_iter().collect();
            let (sql, bindings) = self.compiler().insert(table, &pairs);
            let mut q = sqlx::query(&sql);
            for val in bindings {
                q = bind_value(q, val);
            }
            q.execute(&mut *tx).await.map_err(|e| RusticxError::Query(e.to_string()))?;
            count += 1;
        }
        tx.commit().await.map_err(|e| RusticxError::Transaction(e.to_string()))?;
        Ok(count)
    }

    async fn find(&self, query: &QueryBuilder) -> Result<Vec<Row>> {
        let (sql, bindings) = self.compiler().select(query);
        debug!(sql = %sql, "find");
        self.fetch_rows_with_bindings(&sql, bindings).await
    }

    async fn find_one(&self, query: &QueryBuilder) -> Result<Option<Row>> {
        let mut qb = query.clone();
        qb.limit = Some(1);
        let (sql, bindings) = self.compiler().select(&qb);
        debug!(sql = %sql, "find_one");
        let rows = self.fetch_rows_with_bindings(&sql, bindings).await?;
        Ok(rows.into_iter().next())
    }

    async fn update(&self, query: &QueryBuilder) -> Result<u64> {
        let (sql, bindings) = self.compiler().update(query);
        debug!(sql = %sql, "update");
        self.execute_with_bindings(&sql, bindings).await
    }

    async fn delete(&self, query: &QueryBuilder) -> Result<u64> {
        let (sql, bindings) = self.compiler().delete(query);
        debug!(sql = %sql, "delete");
        self.execute_with_bindings(&sql, bindings).await
    }

    async fn count(&self, query: &QueryBuilder) -> Result<u64> {
        let (sql, bindings) = self.compiler().count(query);
        debug!(sql = %sql, "count");
        let mut q = sqlx::query(&sql);
        for val in bindings {
            q = bind_value(q, val);
        }
        let row = q.fetch_one(&self.pool).await.map_err(|e| RusticxError::Query(e.to_string()))?;
        let cnt: i64 = row.try_get("count").unwrap_or(0);
        Ok(cnt as u64)
    }

    async fn execute_raw(&self, sql: &str, bindings: Vec<Value>) -> Result<u64> {
        self.execute_with_bindings(sql, bindings).await
    }

    async fn query_raw(&self, sql: &str, bindings: Vec<Value>) -> Result<Vec<Row>> {
        self.fetch_rows_with_bindings(sql, bindings).await
    }
}