bsql_core/

executor.rs

//! The `Executor` trait — the runtime contract between generated code and the pool.
//!
//! Code generated by `bsql::query!` calls methods on this trait. `Pool`,
//! `PoolConnection`, and `Transaction` all implement it.
//!
//! The `query_raw` / `query_raw_readonly` methods use the bsql-driver's arena-based
//! row storage. Generated code decodes columns from `Row` via typed getters.

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.
    pub fn len(&self) -> usize {
        self.result.len()
    }

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

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

    /// Iterate over rows.
    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);
        }
    }
}

/// Execute a prepared query and return rows.
///
/// The generated code calls `query_raw`, `query_raw_readonly`, and
/// `execute_raw` on `&Pool`, `&PoolConnection`, or `&Transaction`.
///
/// When the `async` feature is enabled and the pool connects via TCP,
/// `acquire_async()` returns true async connections that use tokio I/O
/// instead of blocking the worker thread. UDS connections remain sync
/// (sub-millisecond, acceptable for tokio).
pub trait Executor {
    /// Execute a query and return all rows.
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult>;

    /// Execute a read-only query. May route to replicas in the future.
    fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult>;

    /// Execute a query and return the number of affected rows.
    fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64>;
}

/// When async feature is enabled, use `acquire_async()` which auto-detects
/// UDS vs TCP: UDS gets sync Connection (fast, sub-ms), TCP gets AsyncConnection
/// (true async I/O via tokio, doesn't block the worker thread).
///
/// The `query_async` / `execute_async` methods on PoolGuard dispatch to the
/// correct backend: sync I/O for UDS connections, async I/O for TCP.
/// Since we need `.await` inside a sync trait method, we use
/// `tokio::task::block_in_place` which allows blocking the current worker
/// while letting other tokio tasks make progress.
#[cfg(feature = "async")]
impl Executor for Pool {
    #[inline]
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let mut guard = tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(self.inner.acquire_async())
        })
        .map_err(BsqlError::from)?;
        let result = tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(guard.query_async(sql, sql_hash, params))
        })
        .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::without_arena(result))
    }

    #[inline]
    fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let pool = self.read_pool.as_ref().unwrap_or(&self.inner);
        let mut guard = tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(pool.acquire_async())
        })
        .map_err(BsqlError::from)?;
        let result = tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(guard.query_async(sql, sql_hash, params))
        })
        .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::without_arena(result))
    }

    #[inline]
    fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        let mut guard = tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(self.inner.acquire_async())
        })
        .map_err(BsqlError::from)?;
        tokio::task::block_in_place(|| {
            tokio::runtime::Handle::current().block_on(guard.execute_async(sql, sql_hash, params))
        })
        .map_err(BsqlError::from_driver_query)
    }
}

/// When async feature is NOT enabled, use plain sync `acquire()` + `query()`.
#[cfg(not(feature = "async"))]
impl Executor for Pool {
    #[inline]
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let mut guard = self.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))
    }

    #[inline]
    fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let pool = self.read_pool.as_ref().unwrap_or(&self.inner);
        let mut guard = 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))
    }

    #[inline]
    fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        let mut guard = self.inner.acquire().map_err(BsqlError::from)?;
        guard
            .execute(sql, sql_hash, params)
            .map_err(BsqlError::from_driver_query)
    }
}

impl Executor for PoolConnection {
    #[inline]
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let mut guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        let result = guard
            .query(sql, sql_hash, params)
            .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::without_arena(result))
    }

    #[inline]
    fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        self.query_raw(sql, sql_hash, params)
    }

    #[inline]
    fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        let mut guard = self.inner.lock().unwrap_or_else(|e| e.into_inner());
        guard
            .execute(sql, sql_hash, params)
            .map_err(BsqlError::from_driver_query)
    }
}

impl Executor for Transaction {
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        self.query_inner(sql, sql_hash, params)
    }

    #[inline]
    fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        self.query_raw(sql, sql_hash, params)
    }

    #[inline]
    fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        self.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()
    }
}