spg-sqlx 7.18.0

//! v7.16.0 — `sqlx::Connection` for SPG.
//!
//! Wraps [`spg_embedded_tokio::AsyncDatabase`]. The engine is
//! single-writer at the lock level, but read-only statements
//! short-circuit through [`spg_embedded_tokio::AsyncReadHandle`]
//! — each `SpgConnection` lazily attaches a read handle on first
//! readonly statement and refreshes it per-statement so PG
//! read-committed semantics hold (every statement sees the
//! latest committed state). Writes / DDL / TX-control take the
//! writer lock as before.
//!
//! Result: `SpgPoolOptions::new().max_connections(20)` behaves
//! like its `PgPool` analogue — concurrent SELECTs run truly in
//! parallel, transactions serialise. Stock sqlx code drops in
//! unchanged.

use std::sync::Arc;

use futures_core::future::BoxFuture;
use futures_core::stream::BoxStream;
use sqlx_core::HashMap;
use sqlx_core::connection::Connection;
use sqlx_core::error::Error;
use sqlx_core::executor::Executor;
use sqlx_core::transaction::Transaction;

use spg_embedded::QueryResult as EngineQueryResult;
use spg_embedded_tokio::{AsyncDatabase, AsyncReadHandle};

use crate::column::SpgColumn;
use crate::database::Spg;
use crate::error::engine_to_sqlx;
use crate::options::SpgConnectOptions;
use crate::query_result::SpgQueryResult;
use crate::row::SpgRow;
use crate::type_info::SpgTypeInfo;

/// One sqlx connection backed by an in-process SPG.
///
/// Holds two engine handles:
///
/// - `inner: AsyncDatabase` — writer path. Used for DDL / DML /
///   transaction control and for statements running inside a
///   transaction (where snapshot semantics would conflict with
///   the user's TX isolation).
///
/// - `read_handle: Option<AsyncReadHandle>` — readonly fan-out
///   path. Lazily built on the first readonly statement seen,
///   refreshed per statement so each SELECT sees the latest
///   committed state (PG read-committed default).
///
/// Both handles share the same underlying engine (the
/// `AsyncDatabase` and the read handle clone the same `Arc` —
/// snapshots are cheap trie-root clones).
#[derive(Debug)]
pub struct SpgConnection {
    pub(crate) inner: AsyncDatabase,
    /// v7.18 — lazy fan-out reader. `None` until the connection
    /// runs its first readonly statement outside a transaction.
    /// `refresh()` is called per-statement so writes committed
    /// elsewhere become visible without manual intervention.
    pub(crate) read_handle: Option<AsyncReadHandle>,
    pub(crate) tx_depth: usize,
    pub(crate) pending_rollback: bool,
}

impl Clone for SpgConnection {
    fn clone(&self) -> Self {
        // The read_handle holds a snapshot frozen at its own
        // creation time; cloning a connection should NOT inherit
        // that snapshot — the clone gets a fresh None and will
        // lazy-init its own. Same `inner` (the AsyncDatabase is
        // `Arc`-shared by design) so writes are still consistent.
        Self {
            inner: self.inner.clone(),
            read_handle: None,
            tx_depth: self.tx_depth,
            pending_rollback: self.pending_rollback,
        }
    }
}

impl SpgConnection {
    /// Build a connection from a ready `AsyncDatabase`. Called
    /// internally by [`SpgConnectOptions::connect`] and by
    /// [`crate::SpgPool::connect_in_memory`].
    pub fn new(inner: AsyncDatabase) -> Self {
        Self {
            inner,
            read_handle: None,
            tx_depth: 0,
            pending_rollback: false,
        }
    }

    /// Borrow the underlying `AsyncDatabase`. Lets advanced
    /// callers reach for the spg-embedded API directly. The
    /// per-connection [`AsyncReadHandle`] used for readonly
    /// fan-out is internal — sqlx code paths don't need to
    /// know it exists.
    #[must_use]
    pub const fn engine(&self) -> &AsyncDatabase {
        &self.inner
    }

    /// v7.18 — return a `&mut AsyncReadHandle` with a snapshot
    /// refreshed to the engine's latest committed state. Lazily
    /// constructs the handle on first call; subsequent calls
    /// refresh the existing one (cheap — clone_snapshot is an
    /// Arc-clone of the trie roots). Used by the routing
    /// helper below for every readonly statement so the
    /// connection sees the same read-committed view a PG client
    /// would see across statement boundaries.
    pub(crate) async fn fresh_read_handle(&mut self) -> &mut AsyncReadHandle {
        if let Some(rh) = self.read_handle.as_mut() {
            rh.refresh().await;
        } else {
            self.read_handle = Some(self.inner.read_handle().await);
        }
        self.read_handle.as_mut().expect("set above")
    }
}

impl Connection for SpgConnection {
    type Database = Spg;
    type Options = SpgConnectOptions;

    fn close(self) -> BoxFuture<'static, Result<(), Error>> {
        // In-process — dropping the last `AsyncDatabase` clone
        // releases the engine. Nothing to send.
        Box::pin(async move { Ok(()) })
    }

    fn close_hard(self) -> BoxFuture<'static, Result<(), Error>> {
        Box::pin(async move { Ok(()) })
    }

    fn ping(&mut self) -> BoxFuture<'_, Result<(), Error>> {
        // The engine doesn't time-out; a quick no-op query
        // exercises the lock path.
        Box::pin(async move {
            self.inner
                .execute("SELECT 1")
                .await
                .map_err(engine_to_sqlx)?;
            Ok(())
        })
    }

    fn begin(&mut self) -> BoxFuture<'_, Result<Transaction<'_, Self::Database>, Error>>
    where
        Self: Sized,
    {
        Transaction::begin(self, None)
    }

    fn shrink_buffers(&mut self) {
        // No-op — engine doesn't expose per-connection buffers.
    }

    fn flush(&mut self) -> BoxFuture<'_, Result<(), Error>> {
        Box::pin(async move { Ok(()) })
    }

    fn should_flush(&self) -> bool {
        false
    }
}

// v7.16.0 — Executor on &mut SpgConnection. fetch_many returns
// `Either<QueryResult, Row>` per sqlx-core's stream contract.
//
// v7.18 — fetch_many / fetch_optional take `&mut SpgConnection`
// across the future (allowed by sqlx's `'c: 'e` bound) so the
// run_one routing can lazy-init / refresh the per-connection
// read handle without cloning state. Readonly statements
// outside a transaction fan out through the snapshot path;
// writer statements + everything inside BEGIN keep using the
// writer mutex.
impl<'c> Executor<'c> for &'c mut SpgConnection {
    type Database = Spg;

    fn fetch_many<'e, 'q: 'e, E>(
        self,
        mut query: E,
    ) -> BoxStream<
        'e,
        Result<
            either::Either<
                <Self::Database as sqlx_core::database::Database>::QueryResult,
                crate::SpgRow,
            >,
            Error,
        >,
    >
    where
        'c: 'e,
        E: 'q + sqlx_core::executor::Execute<'q, Self::Database>,
    {
        use futures_util::stream::{self, StreamExt};
        let sql = query.sql().to_string();
        let arguments = match query.take_arguments() {
            Ok(args) => args,
            Err(e) => {
                return Box::pin(stream::iter(std::iter::once(Err(Error::Encode(e)))));
            }
        };
        let outcome_fut = async move { run_one(self, &sql, arguments).await };
        Box::pin(stream::once(outcome_fut).flat_map(|outcome| {
            let items: Vec<Result<either::Either<SpgQueryResult, SpgRow>, Error>> = match outcome {
                Ok(Outcome::Affected(qr)) => vec![Ok(either::Either::Left(qr))],
                Ok(Outcome::Rows(rows)) => rows
                    .into_iter()
                    .map(|r| Ok(either::Either::Right(r)))
                    .collect(),
                Err(e) => vec![Err(e)],
            };
            stream::iter(items)
        }))
    }

    fn fetch_optional<'e, 'q: 'e, E>(
        self,
        mut query: E,
    ) -> BoxFuture<'e, Result<Option<crate::SpgRow>, Error>>
    where
        'c: 'e,
        E: 'q + sqlx_core::executor::Execute<'q, Self::Database>,
    {
        let sql = query.sql().to_string();
        let arguments = query.take_arguments();
        Box::pin(async move {
            let args = arguments.map_err(Error::Encode)?;
            match run_one(self, &sql, args).await? {
                Outcome::Rows(mut rows) => Ok(rows.drain(..).next()),
                Outcome::Affected(_) => Ok(None),
            }
        })
    }

    fn prepare_with<'e, 'q: 'e>(
        self,
        sql: &'q str,
        _parameters: &'e [<Self::Database as sqlx_core::database::Database>::TypeInfo],
    ) -> BoxFuture<
        'e,
        Result<<Self::Database as sqlx_core::database::Database>::Statement<'q>, Error>,
    >
    where
        'c: 'e,
    {
        let inner = self.inner.clone();
        let sql_str = sql.to_string();
        Box::pin(async move {
            let stmt = inner.prepare(&sql_str).await.map_err(engine_to_sqlx)?;
            // The AsyncStatement wraps the embedded::Statement
            // in Arc — pull it out for the sqlx-side handle.
            // We expose the underlying handle via a tiny adapter
            // method on AsyncStatement (added on the
            // spg-embedded-tokio side).
            let inner_stmt = spg_embedded_tokio::async_statement_inner(&stmt);
            Ok(crate::SpgStatement {
                sql: std::borrow::Cow::Owned(sql_str),
                inner: Some(inner_stmt),
                columns: std::sync::Arc::new(Vec::new()),
                by_name: std::sync::Arc::new(sqlx_core::HashMap::new()),
            })
        })
    }

    fn describe<'e, 'q: 'e>(
        self,
        sql: &'q str,
    ) -> BoxFuture<'e, Result<sqlx_core::describe::Describe<Self::Database>, Error>>
    where
        'c: 'e,
    {
        // v7.17.0 Phase 3.P0-66 — real Describe wired through to
        // `Engine::describe_prepared`. Surfaces column metadata
        // (name / type / nullable) and a parameter count for
        // `sqlx::query!()` compile-time validation. Statement
        // shapes the describe planner can't resolve (JOIN /
        // subquery / unknown table) return an empty `columns`
        // vec — sqlx tolerates this as "no row description
        // available" and the macros fall back to offline mode
        // for those shapes.
        let inner = self.inner.clone();
        let sql_str = sql.to_string();
        Box::pin(async move {
            let (params, cols) = inner.describe(&sql_str).await.map_err(engine_to_sqlx)?;
            let nullable: Vec<Option<bool>> = cols.iter().map(|c| Some(c.nullable)).collect();
            let columns: Vec<SpgColumn> = cols
                .iter()
                .enumerate()
                .map(|(i, c)| {
                    let ti = SpgTypeInfo::from_data_type(c.ty);
                    SpgColumn::new(i, c.name.clone(), ti)
                })
                .collect();
            let parameters = if params.is_empty() {
                None
            } else {
                Some(either::Either::Right(params.len()))
            };
            Ok(sqlx_core::describe::Describe {
                columns,
                parameters,
                nullable,
            })
        })
    }
}

/// Outcome of a single dispatch — either rows-affected (DML / DDL)
/// or a row stream (SELECT). The fetch helpers below convert this
/// to sqlx's `Either<QueryResult, Row>` stream shape.
enum Outcome {
    /// DML / DDL statement returning a rows-affected counter.
    Affected(SpgQueryResult),
    /// SELECT result — every row already converted to an
    /// [`SpgRow`].
    Rows(Vec<SpgRow>),
}

async fn run_one(
    conn: &mut SpgConnection,
    sql: &str,
    arguments: Option<crate::SpgArguments<'_>>,
) -> Result<Outcome, Error> {
    // v7.18 routing decision. Inside a transaction the writer
    // lock has to stay held end-to-end so the user's isolation
    // contract holds; we never route to the snapshot path
    // there. Outside a transaction, parse the SQL and look at
    // `Statement::is_readonly()`. A parse error falls through
    // to the writer path so the user sees the same parse error
    // they'd get from a simple-query SELECT (no behaviour
    // change vs pre-v7.18 for unsupported syntax).
    let in_tx = conn.tx_depth > 0;
    let use_readonly = !in_tx && spg_embedded::Engine::is_readonly_sql(sql);

    let result: EngineQueryResult = if use_readonly {
        let rh = conn.fresh_read_handle().await;
        if let Some(args) = arguments {
            let stmt = rh.prepare(sql).await.map_err(engine_to_sqlx)?;
            rh.execute_prepared(&stmt, args.into_engine_values())
                .await
                .map_err(engine_to_sqlx)?
        } else {
            rh.query(sql).await.map_err(engine_to_sqlx)?
        }
    } else {
        let db = &conn.inner;
        if let Some(args) = arguments {
            let stmt = db.prepare(sql).await.map_err(engine_to_sqlx)?;
            db.execute_prepared(&stmt, args.into_engine_values())
                .await
                .map_err(engine_to_sqlx)?
        } else {
            db.execute(sql).await.map_err(engine_to_sqlx)?
        }
    };
    match result {
        EngineQueryResult::Rows { columns, rows } => {
            let row_values: Vec<Vec<spg_embedded::Value>> =
                rows.into_iter().map(|r| r.values).collect();
            Ok(Outcome::Rows(build_rows(&columns, row_values)))
        }
        EngineQueryResult::CommandOk { affected, .. } => Ok(Outcome::Affected(
            SpgQueryResult::new(u64::try_from(affected).unwrap_or(0)),
        )),
        _ => Ok(Outcome::Affected(SpgQueryResult::default())),
    }
}

#[allow(dead_code)]
fn affected_from(qr: &EngineQueryResult) -> u64 {
    match qr {
        EngineQueryResult::CommandOk { affected, .. } => u64::try_from(*affected).unwrap_or(0),
        EngineQueryResult::Rows { rows, .. } => u64::try_from(rows.len()).unwrap_or(0),
        _ => 0,
    }
}

fn build_rows(
    cols: &[spg_embedded::ColumnSchema],
    rows: Vec<Vec<spg_embedded::Value>>,
) -> Vec<SpgRow> {
    let columns: Arc<Vec<SpgColumn>> = Arc::new(
        cols.iter()
            .enumerate()
            .map(|(i, c)| SpgColumn::new(i, c.name.clone(), SpgTypeInfo::from_data_type(c.ty)))
            .collect(),
    );
    let mut by_name: HashMap<String, usize> = HashMap::new();
    for (i, c) in cols.iter().enumerate() {
        by_name.insert(c.name.clone(), i);
    }
    let by_name = Arc::new(by_name);
    rows.into_iter()
        .map(|values| SpgRow::new(Arc::clone(&columns), Arc::clone(&by_name), values))
        .collect()
}