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::{acquire_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 OwnedResult {
    /// Create from a result and its arena.
    pub(crate) fn new(result: QueryResult, arena: Arena) -> Self {
        Self { result, arena }
    }

    /// 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) {
        // Swap out the arena to return it to the thread-local pool.
        // Arena implements Default, so take() avoids the explicit Arena::new().
        let arena = std::mem::take(&mut self.arena);
        release_arena(arena);
    }
}

/// Execute a prepared query and return rows.
///
/// This trait is sealed — it cannot be implemented outside of bsql-core.
/// The generated code calls `query_raw`, `query_raw_readonly`, and
/// `execute_raw` on `&Pool`, `&PoolConnection`, or `&Transaction`.
pub trait Executor: sealed::Sealed {
    /// Execute a query and return all rows.
    fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> impl std::future::Future<Output = BsqlResult<OwnedResult>> + Send;

    /// 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)],
    ) -> impl std::future::Future<Output = BsqlResult<OwnedResult>> + Send;

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

mod sealed {
    pub trait Sealed {}
    impl Sealed for super::Pool {}
    impl Sealed for super::PoolConnection {}
    impl Sealed for super::Transaction {}
}

impl Executor for Pool {
    async fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let mut guard = self.inner.acquire().await.map_err(BsqlError::from)?;
        let mut arena = acquire_arena();
        let result = guard
            .query(sql, sql_hash, params, &mut arena)
            .await
            .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::new(result, arena))
    }

    async fn query_raw_readonly(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        // Route to replica pool when configured; fall back to primary otherwise.
        let pool = self.read_pool.as_ref().unwrap_or(&self.inner);
        let mut guard = pool.acquire().await.map_err(BsqlError::from)?;
        let mut arena = acquire_arena();
        let result = guard
            .query(sql, sql_hash, params, &mut arena)
            .await
            .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::new(result, arena))
    }

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

impl Executor for PoolConnection {
    async fn query_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<OwnedResult> {
        let mut guard = self.inner.lock().await;
        let mut arena = acquire_arena();
        let result = guard
            .query(sql, sql_hash, params, &mut arena)
            .await
            .map_err(BsqlError::from_driver_query)?;
        Ok(OwnedResult::new(result, arena))
    }

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

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

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

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

    async fn execute_raw(
        &self,
        sql: &str,
        sql_hash: u64,
        params: &[&(dyn Encode + Sync)],
    ) -> BsqlResult<u64> {
        self.execute_inner(sql, sql_hash, params).await
    }
}

#[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::new(result, arena)
    }

    // --- OwnedResult::new ---

    #[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::new(result, arena);
        assert_eq!(owned.len(), 0);
        assert!(owned.is_empty());
        assert_eq!(owned.result.affected_rows(), 42);
    }
}