floz_orm/

db.rs

//! Database connection, transaction management, and the Executor trait.
//!
//! ```ignore
//! use floz::prelude::*;
//!
//! let db = Db::connect("postgres://user:pass@localhost/mydb").await?;
//! let result = db.execute_raw("SELECT 1", vec![]).await?;
//!
//! // Transaction
//! let mut tx = db.begin().await?;
//! tx.execute_raw("INSERT INTO ...", vec![...]).await?;
//! tx.commit().await?;
//! ```

use sqlx::postgres::{PgPool, PgPoolOptions, PgRow};
use tokio::sync::Mutex;

use crate::error::FlozError;
use crate::value::Value;

// ═══════════════════════════════════════════════════════════════
// Db — Connection Pool Wrapper
// ═══════════════════════════════════════════════════════════════

/// A database connection pool wrapping `sqlx::PgPool`.
///
/// `Db` is cheap to clone (internally reference-counted) and safe to share
/// across tokio tasks.
#[derive(Debug, Clone)]
pub struct Db {
    pool: PgPool,
}

impl Db {
    /// Connect to a PostgreSQL database.
    ///
    /// ```ignore
    /// let db = Db::connect("postgres://user:pass@localhost/mydb").await?;
    /// ```
    pub async fn connect(url: &str) -> Result<Self, FlozError> {
        let pool = PgPoolOptions::new()
            .max_connections(10)
            .connect(url)
            .await?;
        Ok(Self { pool })
    }

    /// Connect using an environment variable, falling back to a default URL.
    ///
    /// ```ignore
    /// let db = Db::connect_env("DATABASE_URL", "postgres://localhost:5432/mydb").await?;
    /// ```
    pub async fn connect_env(env_var: &str, default_url: &str) -> Result<Self, FlozError> {
        let url = std::env::var(env_var).unwrap_or_else(|_| default_url.to_string());
        Self::connect(&url).await
    }

    /// Wrap an existing sqlx pool.
    ///
    /// Useful for integrating floz into an existing application.
    pub fn from_pool(pool: PgPool) -> Self {
        Self { pool }
    }

    /// Access the underlying sqlx pool (escape hatch).
    pub fn pool(&self) -> &PgPool {
        &self.pool
    }

    /// Begin a new transaction.
    pub async fn begin(&self) -> Result<Tx, FlozError> {
        let mut conn = self.pool.acquire().await?;
        sqlx::query("BEGIN").execute(&mut *conn).await?;
        Ok(Tx {
            conn: Mutex::new(conn),
            committed: false,
        })
    }
}

// ═══════════════════════════════════════════════════════════════
// Tx — Transaction Wrapper
// ═══════════════════════════════════════════════════════════════

/// A database transaction.
///
/// Wraps a pooled connection with an active transaction.
/// If `commit()` is not called before the `Tx` is dropped,
/// the transaction is automatically rolled back when the connection
/// returns to the pool.
pub struct Tx {
    conn: Mutex<sqlx::pool::PoolConnection<sqlx::Postgres>>,
    committed: bool,
}

impl Tx {
    /// Commit the transaction.
    pub async fn commit(mut self) -> Result<(), FlozError> {
        let mut conn = self.conn.lock().await;
        sqlx::query("COMMIT").execute(&mut **conn).await?;
        self.committed = true;
        Ok(())
    }

    /// Explicitly rollback the transaction.
    pub async fn rollback(self) -> Result<(), FlozError> {
        let mut conn = self.conn.lock().await;
        sqlx::query("ROLLBACK").execute(&mut **conn).await?;
        Ok(())
    }
}

// Debug impl without exposing connection internals
impl std::fmt::Debug for Tx {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("Tx")
            .field("committed", &self.committed)
            .finish()
    }
}

// ═══════════════════════════════════════════════════════════════
// Executor Trait
// ═══════════════════════════════════════════════════════════════

/// The core database execution trait.
///
/// All floz DAO and DSL methods accept `&impl Executor`, allowing
/// seamless switching between connection pools, transactions, and mocks.
///
/// ```ignore
/// // Both work:
/// let users = User::all(&db).await?;
/// let users = User::all(&tx).await?;
/// ```
pub trait Executor {
    /// Execute a raw SQL statement, returning the number of affected rows.
    fn execute_raw(
        &self,
        sql: &str,
        params: Vec<Value>,
    ) -> impl std::future::Future<Output = Result<u64, FlozError>> + Send;

    /// Fetch all matching rows.
    fn fetch_all<T>(
        &self,
        sql: &str,
        params: Vec<Value>,
    ) -> impl std::future::Future<Output = Result<Vec<T>, FlozError>> + Send
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;

    /// Fetch exactly one row. Returns `FlozError::NotFound` if no rows match.
    fn fetch_one<T>(
        &self,
        sql: &str,
        params: Vec<Value>,
    ) -> impl std::future::Future<Output = Result<T, FlozError>> + Send
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;

    /// Fetch zero or one row.
    fn fetch_optional<T>(
        &self,
        sql: &str,
        params: Vec<Value>,
    ) -> impl std::future::Future<Output = Result<Option<T>, FlozError>> + Send
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>;
}

// ═══════════════════════════════════════════════════════════════
// Value → sqlx bind
// ═══════════════════════════════════════════════════════════════

/// Bind a `Vec<Value>` to a `sqlx::query()` builder.
///
/// Each `Value` variant is pattern-matched to call `.bind()` with the
/// concrete type, so sqlx knows the PostgreSQL type OID.
macro_rules! bind_params {
    ($query:expr, $params:expr) => {{
        let mut q = $query;
        for param in $params {
            q = match param {
                // Non-nullable
                Value::Short(v) => q.bind(v),
                Value::Int(v) => q.bind(v),
                Value::BigInt(v) => q.bind(v),
                Value::Real(v) => q.bind(v),
                Value::Double(v) => q.bind(v),
                Value::Bool(v) => q.bind(v),
                Value::String(v) => q.bind(v),
                Value::Bytes(v) => q.bind(v),
                Value::Uuid(v) => q.bind(v),
                Value::DateTime(v) => q.bind(v),
                Value::NaiveDateTime(v) => q.bind(v),
                Value::NaiveDate(v) => q.bind(v),
                Value::NaiveTime(v) => q.bind(v),
                Value::Json(v) => q.bind(v),
                Value::Jsonb(v) => q.bind(v),
                // Nullable (preserves PG type OID for NULLs)
                Value::OptionShort(v) => q.bind(v),
                Value::OptionInt(v) => q.bind(v),
                Value::OptionBigInt(v) => q.bind(v),
                Value::OptionReal(v) => q.bind(v),
                Value::OptionDouble(v) => q.bind(v),
                Value::OptionBool(v) => q.bind(v),
                Value::OptionString(v) => q.bind(v),
                Value::OptionBytes(v) => q.bind(v),
                Value::OptionUuid(v) => q.bind(v),
                Value::OptionDateTime(v) => q.bind(v),
                Value::OptionNaiveDateTime(v) => q.bind(v),
                Value::OptionNaiveDate(v) => q.bind(v),
                Value::OptionNaiveTime(v) => q.bind(v),
                Value::OptionJson(v) => q.bind(v),
                Value::OptionJsonb(v) => q.bind(v),
            };
        }
        q
    }};
}

// ═══════════════════════════════════════════════════════════════
// impl Executor for Db
// ═══════════════════════════════════════════════════════════════

impl Executor for Db {
    async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, FlozError> {
        let query = bind_params!(sqlx::query(sql), params);
        let result = query.execute(&self.pool).await?;
        Ok(result.rows_affected())
    }

    async fn fetch_all<T>(&self, sql: &str, params: Vec<Value>) -> Result<Vec<T>, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let rows = query.fetch_all(&self.pool).await?;
        Ok(rows)
    }

    async fn fetch_one<T>(&self, sql: &str, params: Vec<Value>) -> Result<T, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let row = query
            .fetch_optional(&self.pool)
            .await?
            .ok_or(FlozError::NotFound)?;
        Ok(row)
    }

    async fn fetch_optional<T>(&self, sql: &str, params: Vec<Value>) -> Result<Option<T>, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let row = query.fetch_optional(&self.pool).await?;
        Ok(row)
    }
}

// ═══════════════════════════════════════════════════════════════
// impl Executor for Tx
// ═══════════════════════════════════════════════════════════════

impl Executor for Tx {
    async fn execute_raw(&self, sql: &str, params: Vec<Value>) -> Result<u64, FlozError> {
        let mut conn = self.conn.lock().await;
        let query = bind_params!(sqlx::query(sql), params);
        let result = query.execute(&mut **conn).await?;
        Ok(result.rows_affected())
    }

    async fn fetch_all<T>(&self, sql: &str, params: Vec<Value>) -> Result<Vec<T>, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let mut conn = self.conn.lock().await;
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let rows = query.fetch_all(&mut **conn).await?;
        Ok(rows)
    }

    async fn fetch_one<T>(&self, sql: &str, params: Vec<Value>) -> Result<T, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let mut conn = self.conn.lock().await;
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let row = query
            .fetch_optional(&mut **conn)
            .await?
            .ok_or(FlozError::NotFound)?;
        Ok(row)
    }

    async fn fetch_optional<T>(&self, sql: &str, params: Vec<Value>) -> Result<Option<T>, FlozError>
    where
        T: Send + Unpin + for<'r> sqlx::FromRow<'r, PgRow>,
    {
        let mut conn = self.conn.lock().await;
        let query = bind_params!(sqlx::query_as::<_, T>(sql), params);
        let row = query.fetch_optional(&mut **conn).await?;
        Ok(row)
    }
}

// ═══════════════════════════════════════════════════════════════
// Tests
// ═══════════════════════════════════════════════════════════════

#[cfg(test)]
mod tests {
    use super::*;

    // ── Unit tests (no database needed) ──

    #[test]
    fn db_from_pool_type_check() {
        // Just verifying the type signatures compile
        fn _assert_send<T: Send>() {}
        fn _assert_sync<T: Sync>() {}
        _assert_send::<Db>();
        _assert_sync::<Db>();
    }

    #[test]
    fn tx_is_debug() {
        // Tx should implement Debug
        fn _assert_debug<T: std::fmt::Debug>() {}
        _assert_debug::<Tx>();
    }

    // ── Integration tests (require DATABASE_URL) ──

    /// Get test database URL, or skip the test.
    fn test_url() -> Option<String> {
        std::env::var("DATABASE_URL").ok()
    }

    #[tokio::test]
    async fn connect_to_postgres() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();
        assert!(db.pool().size() > 0 || true); // pool may not have connections yet
    }

    #[tokio::test]
    async fn execute_raw_select() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        // SELECT 1 — should execute without error
        let affected = db.execute_raw("SELECT 1", vec![]).await.unwrap();
        // SELECT doesn't affect rows, but shouldn't error
        assert_eq!(affected, 0);
    }

    #[tokio::test]
    async fn execute_with_params() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        // Create temp table, insert, verify
        db.execute_raw("CREATE TEMP TABLE xql_test_params (id INT, name TEXT)", vec![])
            .await
            .unwrap();

        let affected = db
            .execute_raw(
                "INSERT INTO xql_test_params (id, name) VALUES ($1, $2)",
                vec![Value::Int(1), Value::String("Alice".into())],
            )
            .await
            .unwrap();
        assert_eq!(affected, 1);
    }

    #[tokio::test]
    async fn fetch_all_rows() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_fetch (id INT, name TEXT)", vec![])
            .await
            .unwrap();
        db.execute_raw(
            "INSERT INTO xql_test_fetch VALUES (1, 'Alice'), (2, 'Bob')",
            vec![],
        )
        .await
        .unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
            name: String,
        }

        let rows: Vec<Row> = db
            .fetch_all("SELECT id, name FROM xql_test_fetch ORDER BY id", vec![])
            .await
            .unwrap();
        assert_eq!(rows.len(), 2);
        assert_eq!(rows[0].name, "Alice");
        assert_eq!(rows[1].name, "Bob");
    }

    #[tokio::test]
    async fn fetch_one_found() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_one (id INT, name TEXT)", vec![])
            .await
            .unwrap();
        db.execute_raw("INSERT INTO xql_test_one VALUES (1, 'Alice')", vec![])
            .await
            .unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
            name: String,
        }

        let row: Row = db
            .fetch_one(
                "SELECT id, name FROM xql_test_one WHERE id = $1",
                vec![Value::Int(1)],
            )
            .await
            .unwrap();
        assert_eq!(row.name, "Alice");
    }

    #[tokio::test]
    async fn fetch_one_not_found() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_nf (id INT)", vec![])
            .await
            .unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
        }

        let result: Result<Row, _> = db
            .fetch_one(
                "SELECT id FROM xql_test_nf WHERE id = $1",
                vec![Value::Int(999)],
            )
            .await;
        assert!(matches!(result, Err(FlozError::NotFound)));
    }

    #[tokio::test]
    async fn fetch_optional_some() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_opt (id INT)", vec![])
            .await
            .unwrap();
        db.execute_raw("INSERT INTO xql_test_opt VALUES (42)", vec![])
            .await
            .unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
        }

        let row: Option<Row> = db
            .fetch_optional("SELECT id FROM xql_test_opt WHERE id = $1", vec![Value::Int(42)])
            .await
            .unwrap();
        assert!(row.is_some());
        assert_eq!(row.unwrap().id, 42);
    }

    #[tokio::test]
    async fn fetch_optional_none() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_opt_n (id INT)", vec![])
            .await
            .unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
        }

        let row: Option<Row> = db
            .fetch_optional(
                "SELECT id FROM xql_test_opt_n WHERE id = $1",
                vec![Value::Int(999)],
            )
            .await
            .unwrap();
        assert!(row.is_none());
    }

    // ── Transaction tests ──

    #[tokio::test]
    async fn transaction_commit() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_tx (id INT)", vec![])
            .await
            .unwrap();

        let tx = db.begin().await.unwrap();
        tx.execute_raw(
            "INSERT INTO xql_test_tx VALUES ($1)",
            vec![Value::Int(1)],
        )
        .await
        .unwrap();
        tx.commit().await.unwrap();

        // Row should exist after commit
        #[derive(Debug, sqlx::FromRow)]
        struct Row { id: i32 }

        let rows: Vec<Row> = db
            .fetch_all("SELECT id FROM xql_test_tx", vec![])
            .await
            .unwrap();
        assert_eq!(rows.len(), 1);
    }

    #[tokio::test]
    async fn transaction_rollback_on_drop() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TABLE IF NOT EXISTS xql_test_txd (id INT)", vec![])
            .await
            .unwrap();
        // Clean up any previous test data
        db.execute_raw("DELETE FROM xql_test_txd", vec![])
            .await
            .unwrap();

        {
            let tx = db.begin().await.unwrap();
            tx.execute_raw(
                "INSERT INTO xql_test_txd VALUES ($1)",
                vec![Value::Int(99)],
            )
            .await
            .unwrap();
            // tx dropped here without commit — should rollback
        }

        // Row should NOT exist after rollback
        #[derive(Debug, sqlx::FromRow)]
        struct Row { id: i32 }

        let rows: Vec<Row> = db
            .fetch_all("SELECT id FROM xql_test_txd WHERE id = 99", vec![])
            .await
            .unwrap();
        assert_eq!(rows.len(), 0, "Row should not exist after rollback-on-drop");

        // Clean up
        db.execute_raw("DROP TABLE IF EXISTS xql_test_txd", vec![])
            .await
            .unwrap();
    }

    #[tokio::test]
    async fn transaction_explicit_rollback() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_rb (id INT)", vec![])
            .await
            .unwrap();

        let tx = db.begin().await.unwrap();
        tx.execute_raw(
            "INSERT INTO xql_test_rb VALUES ($1)",
            vec![Value::Int(1)],
        )
        .await
        .unwrap();
        tx.rollback().await.unwrap();

        #[derive(Debug, sqlx::FromRow)]
        struct Row { id: i32 }

        let rows: Vec<Row> = db
            .fetch_all("SELECT id FROM xql_test_rb", vec![])
            .await
            .unwrap();
        assert_eq!(rows.len(), 0);
    }

    #[tokio::test]
    async fn transaction_fetch() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw("CREATE TEMP TABLE xql_test_txf (id INT, name TEXT)", vec![])
            .await
            .unwrap();

        let tx = db.begin().await.unwrap();
        tx.execute_raw(
            "INSERT INTO xql_test_txf VALUES (1, 'Inside Tx')",
            vec![],
        )
        .await
        .unwrap();

        // Fetch from within the transaction — should see uncommitted data
        #[derive(Debug, sqlx::FromRow)]
        struct Row { id: i32, name: String }

        let rows: Vec<Row> = tx
            .fetch_all("SELECT id, name FROM xql_test_txf", vec![])
            .await
            .unwrap();
        assert_eq!(rows.len(), 1);
        assert_eq!(rows[0].name, "Inside Tx");

        tx.rollback().await.unwrap();
    }

    // ── Typed null binding test ──

    #[tokio::test]
    async fn bind_typed_null() {
        let Some(url) = test_url() else { return };
        let db = Db::connect(&url).await.unwrap();

        db.execute_raw(
            "CREATE TEMP TABLE xql_test_null (id INT, name TEXT)",
            vec![],
        )
        .await
        .unwrap();

        // Insert with typed NULL — this is the trap from the spec!
        // Value::OptionString(None) sends TEXT OID, not generic null.
        let affected = db
            .execute_raw(
                "INSERT INTO xql_test_null (id, name) VALUES ($1, $2)",
                vec![Value::Int(1), Value::OptionString(None)],
            )
            .await
            .unwrap();
        assert_eq!(affected, 1);

        #[derive(Debug, sqlx::FromRow)]
        struct Row {
            id: i32,
            name: Option<String>,
        }

        let row: Row = db
            .fetch_one("SELECT id, name FROM xql_test_null WHERE id = 1", vec![])
            .await
            .unwrap();
        assert_eq!(row.id, 1);
        assert!(row.name.is_none());
    }
}