bsql_core/

executor.rs

//! Query dispatch via [`QueryTarget`] — the runtime contract between generated
//! code and the pool/connection/transaction.
//!
//! Code generated by `bsql::query!` converts the user-supplied executor into a
//! [`QueryTarget`] via `Into`, then calls `query_raw` / `query_raw_readonly` /
//! `execute_raw` on it. This enum dispatch replaces the old `Executor` trait
//! and eliminates `Mutex` from `PoolConnection` and `Transaction`.

use bsql_driver_postgres::arena::release_arena;
use bsql_driver_postgres::codec::Encode;
use bsql_driver_postgres::{Arena, QueryResult};

use crate::error::{BsqlError, BsqlResult};
use crate::pool::{Pool, PoolConnection};
use crate::transaction::Transaction;

/// Owned query result that carries its arena alongside the result metadata.
///
/// Generated code calls `.row(i)` to access individual rows. This struct
/// bundles the arena with the result so callsites don't manage arenas manually.
pub struct OwnedResult {
    pub result: QueryResult,
    arena: Arena,
}

impl std::fmt::Debug for OwnedResult {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("OwnedResult")
            .field("rows", &self.result.len())
            .finish()
    }
}

impl OwnedResult {
    /// Create without arena — for queries that use data_buf instead of arena.
    /// Zero allocation: Arena::empty() allocates nothing.
    pub(crate) fn without_arena(result: QueryResult) -> Self {
        Self {
            result,
            arena: Arena::empty(),
        }
    }

    /// Number of rows.
    #[inline]
    pub fn len(&self) -> usize {
        self.result.len()
    }

    /// Whether the result set is empty.
    #[inline]
    pub fn is_empty(&self) -> bool {
        self.result.is_empty()
    }

    /// Get a row by index.
    #[inline]
    pub fn row(&self, idx: usize) -> bsql_driver_postgres::Row<'_> {
        self.result.row(idx, &self.arena)
    }

    /// Iterate over rows.
    #[inline]
    pub fn iter(&self) -> impl Iterator<Item = bsql_driver_postgres::Row<'_>> {
        self.result.rows(&self.arena)
    }
}

impl Drop for OwnedResult {
    fn drop(&mut self) {
        // Return arena to thread-local pool.
        let arena = std::mem::take(&mut self.arena);
        release_arena(arena);
        // Return data buffer to thread-local pool for reuse by next query.
        if let Some(buf) = self.result.take_data_buf() {
            bsql_driver_postgres::release_resp_buf(buf);
        }
        // Return column offsets buffer to thread-local pool.
        let col_offsets = self.result.take_col_offsets();
        if col_offsets.capacity() > 0 {
            bsql_driver_postgres::release_col_offsets(col_offsets);
        }
    }
}

// ---------------------------------------------------------------------------
// QueryTarget — enum dispatch for Pool / PoolConnection / Transaction
// ---------------------------------------------------------------------------

/// Concrete enum for query dispatch. Replaces the `Executor` trait.
///
/// Generated code converts the user-supplied executor (`&Pool`,
/// `&mut PoolConnection`, or `&mut Transaction`) into `QueryTarget` via
/// `Into`, then calls `query_raw` / `query_raw_readonly` / `execute_raw`.
///
/// `Pool` takes a shared reference (`&Pool`) because it acquires from the
/// pool internally. `PoolConnection` and `Transaction` take exclusive
/// references (`&mut`) — no `Mutex` needed.
pub enum QueryTarget<'a> {
    Pool(&'a Pool),
    Conn(&'a mut PoolConnection),
    Tx(&'a mut Transaction),
}

impl<'a> From<&'a Pool> for QueryTarget<'a> {
    #[inline]
    fn from(pool: &'a Pool) -> Self {
        QueryTarget::Pool(pool)
    }
}

impl<'a> From<&'a mut PoolConnection> for QueryTarget<'a> {
    #[inline]
    fn from(conn: &'a mut PoolConnection) -> Self {
        QueryTarget::Conn(conn)
    }
}

impl<'a> From<&'a mut Transaction> for QueryTarget<'a> {
    #[inline]
    fn from(tx: &'a mut Transaction) -> Self {
        QueryTarget::Tx(tx)
    }
}

// --- Async QueryTarget methods ---

#[cfg(feature = "async")]
impl<'a> QueryTarget<'a> {
    /// Execute a query and return all rows.
    #[inline]
    pub async fn query_raw(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        match self {
            QueryTarget::Pool(pool) => {
                let mut guard = pool.inner.acquire_async().await.map_err(BsqlError::from)?;
                let result = guard
                    .query_async(sql, sql_hash, params)
                    .await
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Conn(conn) => {
                let result = conn
                    .inner
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Tx(tx) => tx.query_inner(sql, sql_hash, params),
        }
    }

    /// Execute a read-only query. Routes to replicas for Pool when configured.
    #[inline]
    pub async fn query_raw_readonly(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        match self {
            QueryTarget::Pool(pool) => {
                let driver_pool = pool.read_pool.as_ref().unwrap_or(&pool.inner);
                let mut guard = driver_pool.acquire_async().await.map_err(BsqlError::from)?;
                let result = guard
                    .query_async(sql, sql_hash, params)
                    .await
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            // PoolConnection and Transaction don't have replicas; same as query_raw.
            QueryTarget::Conn(conn) => {
                let result = conn
                    .inner
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Tx(tx) => tx.query_inner(sql, sql_hash, params),
        }
    }

    /// Execute a query and return the number of affected rows.
    #[inline]
    pub async fn execute_raw(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        match self {
            QueryTarget::Pool(pool) => {
                let mut guard = pool.inner.acquire_async().await.map_err(BsqlError::from)?;
                guard
                    .execute_async(sql, sql_hash, params)
                    .await
                    .map_err(BsqlError::from_driver_query)
            }
            QueryTarget::Conn(conn) => conn
                .inner
                .execute(sql, sql_hash, params)
                .map_err(BsqlError::from_driver_query),
            QueryTarget::Tx(tx) => tx.execute_inner(sql, sql_hash, params),
        }
    }
}

// --- Sync QueryTarget methods ---

#[cfg(not(feature = "async"))]
impl<'a> QueryTarget<'a> {
    /// Execute a query and return all rows.
    #[inline]
    pub fn query_raw(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        match self {
            QueryTarget::Pool(pool) => {
                let mut guard = pool.inner.acquire().map_err(BsqlError::from)?;
                let result = guard
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Conn(conn) => {
                let result = conn
                    .inner
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Tx(tx) => tx.query_inner(sql, sql_hash, params),
        }
    }

    /// Execute a read-only query. Routes to replicas for Pool when configured.
    #[inline]
    pub fn query_raw_readonly(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        match self {
            QueryTarget::Pool(pool) => {
                let driver_pool = pool.read_pool.as_ref().unwrap_or(&pool.inner);
                let mut guard = driver_pool.acquire().map_err(BsqlError::from)?;
                let result = guard
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Conn(conn) => {
                let result = conn
                    .inner
                    .query(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)?;
                Ok(OwnedResult::without_arena(result))
            }
            QueryTarget::Tx(tx) => tx.query_inner(sql, sql_hash, params),
        }
    }

    /// Execute a query and return the number of affected rows.
    #[inline]
    pub fn execute_raw(
        self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        match self {
            QueryTarget::Pool(pool) => {
                let mut guard = pool.inner.acquire().map_err(BsqlError::from)?;
                guard
                    .execute(sql, sql_hash, params)
                    .map_err(BsqlError::from_driver_query)
            }
            QueryTarget::Conn(conn) => conn
                .inner
                .execute(sql, sql_hash, params)
                .map_err(BsqlError::from_driver_query),
            QueryTarget::Tx(tx) => tx.execute_inner(sql, sql_hash, params),
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bsql_driver_postgres::arena::{acquire_arena, release_arena};
    use bsql_driver_postgres::{ColumnDesc, QueryResult};
    use std::sync::Arc;

    /// Helper: build an OwnedResult with `n` rows and `num_cols` columns.
    /// Each column offset entry is a dummy (0, 0) pair — sufficient for
    /// testing len/is_empty/row-count without decoding real data.
    fn make_owned_result(num_rows: usize, num_cols: usize) -> OwnedResult {
        let arena = acquire_arena();
        let cols: Arc<[ColumnDesc]> = (0..num_cols)
            .map(|i| ColumnDesc {
                name: format!("c{i}").into(),
                type_oid: 23, // int4
                type_size: 4,
                table_oid: 0,
                column_id: 0,
            })
            .collect::<Vec<_>>()
            .into();

        let col_offsets: Vec<(usize, i32)> = vec![(0, -1); num_rows * num_cols]; // NULL columns
        let result = QueryResult::from_parts(col_offsets, num_cols, cols, 0);
        OwnedResult { result, arena }
    }

    // --- OwnedResult ---

    #[test]
    fn owned_result_new_zero_rows() {
        let owned = make_owned_result(0, 2);
        assert_eq!(owned.len(), 0);
        assert!(owned.is_empty());
    }

    #[test]
    fn owned_result_new_single_row() {
        let owned = make_owned_result(1, 3);
        assert_eq!(owned.len(), 1);
        assert!(!owned.is_empty());
    }

    #[test]
    fn owned_result_new_multiple_rows() {
        let owned = make_owned_result(5, 2);
        assert_eq!(owned.len(), 5);
        assert!(!owned.is_empty());
    }

    // --- OwnedResult::row ---

    #[test]
    fn owned_result_row_access() {
        let owned = make_owned_result(3, 2);
        // Should not panic for valid indices
        let _r0 = owned.row(0);
        let _r1 = owned.row(1);
        let _r2 = owned.row(2);
    }

    #[test]
    #[should_panic]
    fn owned_result_row_out_of_bounds_panics() {
        let owned = make_owned_result(2, 1);
        let _r = owned.row(2); // out of bounds
    }

    // --- OwnedResult::iter ---

    #[test]
    fn owned_result_iter_count() {
        let owned = make_owned_result(4, 2);
        let count = owned.iter().count();
        assert_eq!(count, 4);
    }

    #[test]
    fn owned_result_iter_empty() {
        let owned = make_owned_result(0, 2);
        let count = owned.iter().count();
        assert_eq!(count, 0);
    }

    // --- OwnedResult::Drop releases arena back to pool ---

    #[test]
    fn owned_result_drop_releases_arena() {
        // Acquire an arena, wrap it in OwnedResult, drop it.
        // After drop, acquiring should succeed (arena was returned to pool).
        let owned = make_owned_result(1, 1);
        drop(owned);
        // If the arena was released, we can acquire again without issue.
        let arena = acquire_arena();
        release_arena(arena);
    }

    // --- OwnedResult with zero columns ---

    #[test]
    fn owned_result_zero_columns() {
        // Commands like INSERT without RETURNING have 0 columns
        let arena = acquire_arena();
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 42);
        let owned = OwnedResult { result, arena };
        assert_eq!(owned.len(), 0);
        assert!(owned.is_empty());
        assert_eq!(owned.result.affected_rows(), 42);
    }

    // --- OwnedResult::without_arena ---

    #[test]
    fn owned_result_without_arena_len_zero() {
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 0);
        let owned = OwnedResult::without_arena(result);
        assert_eq!(owned.len(), 0);
    }

    #[test]
    fn owned_result_without_arena_is_empty() {
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 0);
        let owned = OwnedResult::without_arena(result);
        assert!(owned.is_empty());
    }

    #[test]
    fn owned_result_without_arena_with_rows() {
        let cols: Arc<[ColumnDesc]> = vec![ColumnDesc {
            name: "c0".into(),
            type_oid: 23,
            type_size: 4,
            table_oid: 0,
            column_id: 0,
        }]
        .into();
        let col_offsets = vec![(0, -1); 3]; // 3 rows, 1 col each (all NULL)
        let result = QueryResult::from_parts(col_offsets, 1, cols, 0);
        let owned = OwnedResult::without_arena(result);
        assert_eq!(owned.len(), 3);
        assert!(!owned.is_empty());
    }

    // --- OwnedResult Debug ---

    #[test]
    fn owned_result_debug_format() {
        let owned = make_owned_result(5, 2);
        let dbg = format!("{owned:?}");
        assert!(
            dbg.contains("OwnedResult"),
            "Debug should contain struct name: {dbg}"
        );
        assert!(dbg.contains("5"), "Debug should contain row count: {dbg}");
    }

    // --- OwnedResult drop without_arena variant ---

    #[test]
    fn owned_result_without_arena_drop_does_not_panic() {
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 0);
        let owned = OwnedResult::without_arena(result);
        drop(owned); // Must not panic — arena is Arena::empty()
    }

    // --- Pool / PoolConnection / Transaction Send constraints ---
    // PoolConnection and Transaction are Send but NOT Sync (no Mutex).

    #[test]
    fn pool_is_send_and_sync() {
        fn _assert_send<T: Send>() {}
        fn _assert_sync<T: Sync>() {}
        _assert_send::<crate::pool::Pool>();
        _assert_sync::<crate::pool::Pool>();
    }

    #[test]
    fn pool_connection_is_send() {
        fn _assert_send<T: Send>() {}
        _assert_send::<crate::pool::PoolConnection>();
    }

    #[test]
    fn transaction_is_send() {
        fn _assert_send<T: Send>() {}
        _assert_send::<crate::transaction::Transaction>();
    }

    #[test]
    fn owned_result_is_send_and_sync() {
        fn _assert_send<T: Send>() {}
        fn _assert_sync<T: Sync>() {}
        _assert_send::<OwnedResult>();
        _assert_sync::<OwnedResult>();
    }

    // --- QueryTarget From impls ---

    #[test]
    fn query_target_from_pool_compiles() {
        // Compile-time test: From<&Pool> for QueryTarget exists
        fn _check<'a>(_: &'a Pool) -> QueryTarget<'a> {
            unimplemented!()
        }
    }

    #[test]
    fn query_target_from_pool_connection_compiles() {
        fn _check<'a>(_: &'a mut PoolConnection) -> QueryTarget<'a> {
            unimplemented!()
        }
    }

    #[test]
    fn query_target_from_transaction_compiles() {
        fn _check<'a>(_: &'a mut Transaction) -> QueryTarget<'a> {
            unimplemented!()
        }
    }

    // --- OwnedResult affected_rows on without_arena variant ---

    #[test]
    fn owned_result_without_arena_affected_rows() {
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 42);
        let owned = OwnedResult::without_arena(result);
        assert_eq!(owned.result.affected_rows(), 42);
    }

    #[test]
    fn owned_result_without_arena_affected_rows_zero() {
        let cols: Arc<[ColumnDesc]> = Arc::from(Vec::new());
        let result = QueryResult::from_parts(vec![], 0, cols, 0);
        let owned = OwnedResult::without_arena(result);
        assert_eq!(owned.result.affected_rows(), 0);
    }

    // --- OwnedResult iter yields correct rows ---

    #[test]
    fn owned_result_iter_yields_all_rows() {
        let owned = make_owned_result(3, 1);
        let rows: Vec<_> = owned.iter().collect();
        assert_eq!(rows.len(), 3);
    }

    // --- OwnedResult Debug format with 0 rows ---

    #[test]
    fn owned_result_debug_format_zero_rows() {
        let owned = make_owned_result(0, 2);
        let dbg = format!("{owned:?}");
        assert!(dbg.contains("OwnedResult"), "should contain name: {dbg}");
        assert!(dbg.contains("0"), "should contain 0: {dbg}");
    }

    // --- OwnedResult row panics for empty result ---

    #[test]
    #[should_panic]
    fn owned_result_row_panics_on_empty() {
        let owned = make_owned_result(0, 1);
        let _r = owned.row(0);
    }

    // --- QueryTarget variant discrimination ---

    #[test]
    fn query_target_pool_variant_matches() {
        fn _check_pool<'a>(pool: &'a Pool) {
            let qt: QueryTarget<'a> = pool.into();
            assert!(matches!(qt, QueryTarget::Pool(_)));
        }
    }

    #[test]
    fn query_target_conn_variant_matches() {
        fn _check_conn<'a>(conn: &'a mut PoolConnection) {
            let qt: QueryTarget<'a> = conn.into();
            assert!(matches!(qt, QueryTarget::Conn(_)));
        }
    }

    #[test]
    fn query_target_tx_variant_matches() {
        fn _check_tx<'a>(tx: &'a mut Transaction) {
            let qt: QueryTarget<'a> = tx.into();
            assert!(matches!(qt, QueryTarget::Tx(_)));
        }
    }

    // --- OwnedResult large row count ---

    #[test]
    fn owned_result_large_row_count() {
        let owned = make_owned_result(1000, 2);
        assert_eq!(owned.len(), 1000);
        assert!(!owned.is_empty());
        assert_eq!(owned.iter().count(), 1000);
    }

    // --- OwnedResult single column ---

    #[test]
    fn owned_result_single_column_row_access() {
        let owned = make_owned_result(2, 1);
        let _r0 = owned.row(0);
        let _r1 = owned.row(1);
        assert_eq!(owned.len(), 2);
    }

    // --- Multiple drops don't panic ---

    #[test]
    fn owned_result_drop_twice_via_option() {
        let owned = make_owned_result(1, 1);
        let mut opt = Some(owned);
        opt.take(); // first drop
                    // second drop on None — should not panic
        drop(opt);
    }
}