lora_database/

transaction.rs

use std::collections::BTreeMap;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex, MutexGuard};
use std::time::{Duration, Instant};

use anyhow::Result;
use thiserror::Error;

use lora_analyzer::Analyzer;
use lora_compiler::{CompiledQuery, Compiler};
use lora_executor::{
    classify_stream, compiled_result_columns, project_rows, ExecuteOptions, ExecutionContext,
    Executor, LoraValue, MutableExecutionContext, MutableExecutor, MutablePullExecutor,
    PullExecutor, QueryResult, Row, RowSource,
};
use lora_parser::parse_query;
use lora_store::{InMemoryGraph, MutationEvent, MutationRecorder};
use lora_wal::WalRecorder;

use crate::error::LoraError;
use crate::explain::{OperatorMetrics, ProfileMetrics, QueryPlan, QueryProfile};
use crate::live_store::LiveStore;
use crate::snapshot::ManagedSnapshotStore;
use crate::stream::QueryStream;
use crate::wal::write_scope::ensure_wal_not_poisoned;
use lora_compiler::plan_tree_from_compiled;
use lora_executor::{plan_result_columns, CollectorGuard, MetricsCollector};

/// Transaction-lifecycle invariant violations.
///
/// All variants used to be raised as `anyhow!("...")` strings. Surfacing
/// them as a typed enum lets [`crate::LoraError`] route them onto stable
/// `LoraErrorCode`s without phrase-matching the `Display` text.
#[derive(Debug, Clone, PartialEq, Eq, Error)]
pub enum TransactionError {
    #[error("transaction is already closed")]
    AlreadyClosed,

    #[error("transaction has no live graph guard")]
    NoGraphGuard,

    #[error("transaction has no staged graph")]
    NoStagedGraph,

    #[error("cannot commit transaction while a streaming cursor is still active")]
    CursorActiveCommit,

    #[error("cannot start a new statement while a streaming cursor is still active")]
    CursorActiveStatement,

    #[error("cannot execute mutating query in read-only transaction")]
    ReadOnlyMutation,

    #[error("streaming write cursor requires a ReadWrite transaction")]
    StreamingRequiresReadWrite,

    #[error("read-only transaction cannot publish staged graph")]
    ReadOnlyCommit,

    #[error("transaction state lock is poisoned")]
    Poisoned,
}

/// Transaction execution mode.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransactionMode {
    /// Use the read-only executor. Write operators return read-only errors.
    ReadOnly,
    /// Execute reads and writes against a staged graph, then publish on commit.
    ReadWrite,
}

/// What a transaction holds onto for the duration of its statements.
///
/// `Read` simply pins an `Arc<InMemoryGraph>` snapshot — readers don't
/// take any lock, so nothing observable changes when a writer commits
/// a new version mid-transaction. `Write` holds the writer Mutex
/// (serializing commit ordering) and a snapshot of the live graph at
/// the point the transaction began; the working copy is built lazily
/// in `TxInner::staged` on first mutation, mirroring the previous
/// "clone on first mutation" behavior.
pub(crate) enum LiveStoreGuard<'db> {
    Read(Arc<InMemoryGraph>),
    Write(WriteLease<'db>),
}

/// Writer lease for a `ReadWrite` transaction. Holds the per-database
/// writer Mutex plus a read snapshot of the graph at lease open time.
/// The mutating working copy lives in `TxInner::staged` and is
/// cloned from `snapshot` lazily.
pub(crate) struct WriteLease<'db> {
    /// Held for the tx lifetime so concurrent ReadWrite txns serialize.
    pub(crate) _writer_lock: MutexGuard<'db, ()>,
    /// Pointer back to the live store so commit can publish.
    pub(crate) store: Arc<LiveStore<InMemoryGraph>>,
    /// Read-only view of the graph at lease open time. The first
    /// mutating statement clones from this into `TxInner::staged`.
    pub(crate) snapshot: Arc<InMemoryGraph>,
}

impl LiveStoreGuard<'_> {
    fn as_graph(&self) -> &InMemoryGraph {
        match self {
            Self::Read(arc) => arc,
            Self::Write(lease) => &lease.snapshot,
        }
    }
}

/// Captures the staged graph and tx-local mutation buffer at the point
/// a statement is opened, so a failed/dropped statement can be rolled
/// back to that point without affecting earlier work in the same
/// transaction.
pub(crate) struct Savepoint {
    staged: Option<InMemoryGraph>,
    buffer_len: usize,
}

/// How a transaction-bound stream finished. Exhaustion commits that
/// statement's staged changes into the transaction; interruption means
/// drop or runtime error before all rows were observed.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) enum TxStreamOutcome {
    Exhausted,
    Interrupted,
}

impl TxStreamOutcome {
    fn should_restore_savepoint(self, rollback_on_drop: bool) -> bool {
        matches!(self, Self::Interrupted) && rollback_on_drop
    }
}

/// Owns the active-cursor token for a transaction-bound stream.
///
/// The lease is created only after a cursor has opened successfully. If stream
/// code forgets to finalize it explicitly, `Drop` treats the cursor as
/// interrupted so a mutating statement cannot accidentally commit partial work
/// into the transaction.
pub(crate) struct TxCursorLease {
    handle: Arc<Mutex<TxInner>>,
    rollback_on_drop: bool,
    finalized: bool,
}

impl TxCursorLease {
    pub(crate) fn new(handle: Arc<Mutex<TxInner>>, rollback_on_drop: bool) -> Self {
        Self {
            handle,
            rollback_on_drop,
            finalized: false,
        }
    }

    pub(crate) fn finalize(&mut self, outcome: TxStreamOutcome) {
        if self.finalized {
            return;
        }
        finalize_tx_stream(&self.handle, outcome, self.rollback_on_drop);
        self.finalized = true;
    }
}

impl Drop for TxCursorLease {
    fn drop(&mut self) {
        self.finalize(TxStreamOutcome::Interrupted);
    }
}

/// Buffers `MutationEvent`s emitted by the staged graph while a
/// transaction is in progress. The buffer replaces direct WAL writes
/// during the transaction body; on commit the host replays the
/// buffer into the real `WalRecorder` as a single durable
/// transaction. Statement rollback truncates the buffer back to its
/// pre-statement length; transaction rollback drops it entirely.
pub(crate) struct BufferingRecorder {
    buffer: Arc<Mutex<Vec<MutationEvent>>>,
    failed: Arc<AtomicBool>,
}

impl BufferingRecorder {
    pub(crate) fn new(buffer: Arc<Mutex<Vec<MutationEvent>>>, failed: Arc<AtomicBool>) -> Self {
        Self { buffer, failed }
    }
}

impl MutationRecorder for BufferingRecorder {
    fn record(&self, event: MutationEvent) {
        match self.buffer.lock() {
            Ok(mut buf) => buf.push(event),
            Err(_) => self.failed.store(true, Ordering::Release),
        }
    }
}

/// Shared transaction state. Wrapped in `Arc<Mutex<>>` so a
/// `QueryStream` opened against the transaction can release its
/// cursor token and signal savepoint-rollback intent on drop without
/// borrowing the [`Transaction`] handle.
pub(crate) struct TxInner {
    /// The cloned staging graph. Mutated by write statements through
    /// the [`MutableExecutor`]; read by read-only statements through
    /// [`PullExecutor`]. `None` once the transaction has been closed.
    pub(crate) staged: Option<InMemoryGraph>,
    /// Tx-local mutation log, populated by the [`BufferingRecorder`]
    /// installed on `staged`. Replayed into the real WAL exactly once
    /// at commit time.
    pub(crate) buffer: Arc<Mutex<Vec<MutationEvent>>>,
    /// Latched if a mutation event could not be appended to `buffer`.
    /// Commit must fail rather than publishing a graph whose WAL event stream
    /// is incomplete.
    pub(crate) buffer_failed: Arc<AtomicBool>,
    /// Per-statement savepoint snapshot. Set when a statement opens,
    /// cleared on successful completion, restored on
    /// failure/premature drop.
    pub(crate) pending_savepoint: Option<Savepoint>,
    /// True while a `QueryStream` opened against this transaction is
    /// alive. Blocks new statements and prevents commit until the
    /// cursor is released.
    pub(crate) cursor_active: bool,
    /// True after `commit` or `rollback` has run, regardless of
    /// outcome. Subsequent operations fail loudly instead of silently
    /// running on stale state.
    pub(crate) closed: bool,
    /// Transaction execution mode chosen at `begin_transaction` time.
    pub(crate) mode: TransactionMode,
    /// Whether this transaction needs a mutation buffer for durable WAL
    /// replay. Databases without a WAL can skip recorder installation and
    /// avoid cloning mutation payloads into an unused buffer.
    pub(crate) buffer_mutations: bool,
}

/// Explicit transaction over the in-memory graph.
///
/// The implementation is conservative: read-only transactions hold a
/// database read lock, and read-write transactions hold the database
/// write lock. Read-write transactions lazily create a cloned staging
/// graph on the first mutating statement, then either swap that graph
/// into place on commit or drop it on rollback. Explicit mutating
/// statements capture a graph +
/// WAL-buffer savepoint so a failed or dropped streaming statement
/// only rolls back its own effects, not the transaction as a whole.
///
/// When a WAL is attached, mutation events fire into a tx-local
/// buffer rather than the durable log. The buffer is replayed into
/// the WAL exactly once at commit, so recovery never observes
/// partial / aborted / dropped statements.
pub struct Transaction<'db> {
    pub(crate) live: Option<LiveStoreGuard<'db>>,
    pub(crate) inner: Arc<Mutex<TxInner>>,
    pub(crate) wal: Option<Arc<WalRecorder>>,
    pub(crate) snapshots: Option<Arc<ManagedSnapshotStore>>,
    mode: TransactionMode,
}

impl<'db> Transaction<'db> {
    /// Build a fresh transaction. Used by `Database::begin_transaction`.
    pub(crate) fn new(
        live: LiveStoreGuard<'db>,
        wal: Option<Arc<WalRecorder>>,
        snapshots: Option<Arc<ManagedSnapshotStore>>,
        mode: TransactionMode,
    ) -> Self {
        let buffer_mutations = wal.is_some();
        let inner = TxInner {
            staged: None,
            buffer: Arc::new(Mutex::new(Vec::new())),
            buffer_failed: Arc::new(AtomicBool::new(false)),
            pending_savepoint: None,
            cursor_active: false,
            closed: false,
            mode,
            buffer_mutations,
        };
        Self {
            live: Some(live),
            inner: Arc::new(Mutex::new(inner)),
            wal,
            snapshots,
            mode,
        }
    }

    /// Transaction mode chosen at begin time.
    pub fn mode(&self) -> TransactionMode {
        self.mode
    }

    /// Execute a query inside the transaction and return a materialized
    /// `QueryResult`.
    pub fn execute(
        &mut self,
        query: &str,
        options: Option<ExecuteOptions>,
    ) -> Result<QueryResult, LoraError> {
        self.execute_with_params(query, options, BTreeMap::new())
    }

    /// Execute a query inside the transaction with a cooperative deadline.
    pub fn execute_with_timeout(
        &mut self,
        query: &str,
        options: Option<ExecuteOptions>,
        timeout: Duration,
    ) -> Result<QueryResult, LoraError> {
        let rows = self.execute_rows_with_params_deadline(
            query,
            BTreeMap::new(),
            Some(deadline_after(timeout)),
        )?;
        Ok(project_rows(rows, options.unwrap_or_default()))
    }

    /// Execute a parameterised query inside the transaction.
    pub fn execute_with_params(
        &mut self,
        query: &str,
        options: Option<ExecuteOptions>,
        params: BTreeMap<String, LoraValue>,
    ) -> Result<QueryResult, LoraError> {
        let rows = self.execute_rows_with_params_deadline(query, params, None)?;
        Ok(project_rows(rows, options.unwrap_or_default()))
    }

    /// Execute a parameterised query inside the transaction with a cooperative
    /// deadline.
    pub fn execute_with_params_timeout(
        &mut self,
        query: &str,
        options: Option<ExecuteOptions>,
        params: BTreeMap<String, LoraValue>,
        timeout: Duration,
    ) -> Result<QueryResult, LoraError> {
        let rows =
            self.execute_rows_with_params_deadline(query, params, Some(deadline_after(timeout)))?;
        Ok(project_rows(rows, options.unwrap_or_default()))
    }

    /// Execute a query inside the transaction and return hydrated rows before
    /// final result-format projection.
    pub fn execute_rows(&mut self, query: &str) -> Result<Vec<Row>, LoraError> {
        self.execute_rows_with_params(query, BTreeMap::new())
    }

    /// Execute a parameterised query inside the transaction and return hydrated
    /// rows before final result-format projection.
    pub fn execute_rows_with_params(
        &mut self,
        query: &str,
        params: BTreeMap<String, LoraValue>,
    ) -> Result<Vec<Row>, LoraError> {
        Ok(self.execute_rows_with_params_deadline(query, params, None)?)
    }

    fn execute_rows_with_params_deadline(
        &mut self,
        query: &str,
        params: BTreeMap<String, LoraValue>,
        deadline: Option<Instant>,
    ) -> Result<Vec<Row>> {
        let compiled = self.compile_in_tx(query)?;
        self.execute_rows_compiled_deadline(&compiled, params, deadline)
    }

    fn execute_rows_compiled_deadline(
        &mut self,
        compiled: &CompiledQuery,
        params: BTreeMap<String, LoraValue>,
        deadline: Option<Instant>,
    ) -> Result<Vec<Row>> {
        // ReadOnly tx: never clones, runs straight against live.
        if self.is_read_only_unchecked() {
            self.precheck_open_no_savepoint()?;
            return self.execute_live_compiled(compiled, params, deadline);
        }

        // ReadWrite tx, lazy-clone aware.
        let mut inner = self.begin_statement()?;
        let is_mutating = classify_stream(compiled).is_mutating();

        if !is_mutating {
            // Read-only statement in a ReadWrite tx. Run against
            // staged if it has been materialized (so the read
            // sees prior in-tx writes), otherwise straight off
            // the live graph — which equals staged-as-it-would-be
            // because no writes have happened yet.
            return self.execute_read_statement(inner, compiled, params, deadline);
        }

        // Mutating statement: lazy-clone the live graph if this
        // is the first write in the tx, then capture a savepoint
        // and run the mutable executor.
        let savepoint = self.prepare_mutating_statement(&mut inner)?;

        let exec_result: ExecResultRows = {
            let staged = inner.staged_mut()?;
            execute_mutable_compiled(staged, compiled, params, deadline)
        };

        match exec_result {
            Ok(rows) => Ok(rows),
            Err(err) => {
                restore_savepoint(&mut inner, savepoint);
                Err(err)
            }
        }
    }

    /// Open a streaming write cursor over the staged graph for a
    /// pre-compiled mutating plan, used by the hidden auto-commit
    /// stream path in `Database::stream_with_params`.
    ///
    /// The returned `Box<dyn RowSource + 'static>` may be either a
    /// real per-row [`StreamingWriteCursor`][lora_executor::StreamingWriteCursor],
    /// a mutable UNION cursor, or a [`BufferedRowSource`][lora_executor::BufferedRowSource]
    /// for the remaining materialized leaves. Either way:
    ///
    /// * The cursor mutates the *staged* graph, never the live store.
    /// * Mutations fire the [`BufferingRecorder`] installed on staged
    ///   by [`Self::ensure_staged_locked`], which accumulates into
    ///   `inner.buffer` and is replayed into the WAL on commit.
    /// * `cursor_active` is set to `true` here. The caller MUST clear
    ///   it before invoking [`Self::commit`] or [`Self::rollback`] —
    ///   the cursor itself does not.
    ///
    /// # Safety
    ///
    /// The cursor is `'static` because it owns its compiled query (via
    /// the supplied `Arc`) and aliases the staged graph through a raw
    /// pointer. Soundness depends on the invariant that
    /// `inner.staged` remains `Some(_)` at a stable address for the
    /// cursor's lifetime. That invariant holds while
    /// `cursor_active = true` blocks every other path that could
    /// move or drop staged: explicit statements (`begin_statement`
    /// rejects), `commit` and `rollback` (rejected until the caller
    /// clears `cursor_active`).
    pub(crate) fn open_streaming_compiled_autocommit(
        &mut self,
        compiled: Arc<CompiledQuery>,
        params: BTreeMap<String, LoraValue>,
    ) -> Result<Box<dyn RowSource + 'static>> {
        if self.is_read_only_unchecked() {
            return Err(TransactionError::StreamingRequiresReadWrite.into());
        }

        let mut inner = self.begin_statement()?;
        self.ensure_staged_locked(&mut inner)?;
        inner.activate_cursor();

        // SAFETY: `inner.staged` is `Some` after `ensure_staged_locked`,
        // and stays at the same address while `cursor_active = true`
        // (see method-level safety note).
        let staged_ptr: *mut InMemoryGraph = inner
            .staged
            .as_mut()
            .ok_or(TransactionError::NoStagedGraph)?
            as *mut _;
        drop(inner);

        // SAFETY: `compiled` (Arc held by the caller / AutoCommit guard)
        // keeps the plan alive; `staged_ptr` is valid for the cursor's
        // lifetime per the invariant above. We extend both lifetimes
        // to `'static` so the resulting cursor can sit inside the
        // `'static`-shaped AutoCommit variant of `QueryStream`.
        let storage_static: &'static mut InMemoryGraph = unsafe { &mut *staged_ptr };
        let compiled_static: &'static CompiledQuery =
            unsafe { std::mem::transmute::<&CompiledQuery, _>(compiled.as_ref()) };

        // `MutablePullExecutor::open_compiled` picks the narrowest
        // cursor shape it can: per-row write cursor, branch-wise
        // mutable UNION cursor, or a buffered materialized leaf.
        let cursor = MutablePullExecutor::new(storage_static, params)
            .open_compiled(compiled_static)
            .map_err(|e| {
                // Roll back: the cursor build never happened, so the
                // tx is in a clean-but-poisoned state. Discard
                // everything and let the caller bubble the error.
                if let Ok(mut inner) = self.inner.lock() {
                    discard_transaction_state(&mut inner);
                }
                self.live.take();
                anyhow::Error::from(e)
            })?;

        // The Arc<CompiledQuery> is the safety anchor for the
        // `'static` plan reference. Keep it alive for the cursor's
        // lifetime by leaking a clone into the cursor's owned data.
        // We can't store it on the cursor itself (it's a Box<dyn>),
        // so we wrap the cursor in a guard that owns the Arc.
        Ok(Box::new(StreamingCursorWithArc {
            cursor,
            _compiled: compiled,
        }))
    }

    /// Compile a query in this transaction's view of the world:
    /// against `staged` if it has been materialized, otherwise
    /// straight against `live`. The two are equivalent before the
    /// first mutating statement, so the resulting plan is valid
    /// either way.
    fn compile_in_tx(&self, query: &str) -> Result<CompiledQuery> {
        let document = parse_query(query)?;
        let (resolved, stats) = {
            let inner = self.lock_inner()?;
            if let Some(staged) = &inner.staged {
                let mut analyzer = Analyzer::new(staged);
                let resolved = analyzer.analyze(&document)?;
                let stats = staged.graph_stats();
                (resolved, stats)
            } else {
                drop(inner);
                let live = self.live.as_ref().ok_or(TransactionError::NoGraphGuard)?;
                let graph = live.as_graph();
                let mut analyzer = Analyzer::new(graph);
                let resolved = analyzer.analyze(&document)?;
                let stats = graph.graph_stats();
                (resolved, stats)
            }
        };
        Ok(Compiler::compile(&resolved, &stats))
    }

    /// Materialize `inner.staged` if it doesn't exist yet —
    /// ReadWrite transactions defer this clone until the first
    /// mutating statement.
    fn ensure_staged_locked(&self, inner: &mut MutexGuard<'_, TxInner>) -> Result<()> {
        if inner.staged.is_some() {
            return Ok(());
        }
        let live = self.live.as_ref().ok_or(TransactionError::NoGraphGuard)?;
        let mut staged: InMemoryGraph = live.as_graph().clone();
        if matches!(inner.mode, TransactionMode::ReadWrite) && inner.buffer_mutations {
            staged.set_mutation_recorder(Some(Arc::new(BufferingRecorder::new(
                inner.buffer.clone(),
                inner.buffer_failed.clone(),
            )) as Arc<dyn MutationRecorder>));
        }
        inner.staged = Some(staged);
        Ok(())
    }

    fn execute_live_compiled(
        &self,
        compiled: &CompiledQuery,
        params: BTreeMap<String, LoraValue>,
        deadline: Option<Instant>,
    ) -> Result<Vec<Row>> {
        let live = self.live.as_ref().ok_or(TransactionError::NoGraphGuard)?;
        execute_read_compiled(live.as_graph(), compiled, params, deadline)
    }

    fn execute_read_statement(
        &self,
        inner: MutexGuard<'_, TxInner>,
        compiled: &CompiledQuery,
        params: BTreeMap<String, LoraValue>,
        deadline: Option<Instant>,
    ) -> Result<Vec<Row>> {
        match inner.staged.as_ref() {
            Some(staged) => execute_read_compiled(staged, compiled, params, deadline),
            None => {
                drop(inner);
                self.execute_live_compiled(compiled, params, deadline)
            }
        }
    }

    fn prepare_mutating_statement(
        &self,
        inner: &mut MutexGuard<'_, TxInner>,
    ) -> Result<Option<Savepoint>> {
        let clone_savepoint_graph = inner.staged.is_some();
        self.ensure_staged_locked(inner)?;
        Ok(Some(take_savepoint(inner, clone_savepoint_graph)))
    }

    /// Compile `query` against the transaction's view of the graph and
    /// return the plan that *would* run. Never executes the query, so
    /// running `explain` on a mutating statement leaves the transaction's
    /// staged graph untouched.
    pub fn explain(
        &self,
        query: &str,
        _params: Option<BTreeMap<String, LoraValue>>,
    ) -> Result<QueryPlan, LoraError> {
        let compiled = self.compile_in_tx(query).map_err(LoraError::from_anyhow)?;
        Ok(query_plan_for(query, &compiled))
    }

    /// Execute `query` inside the transaction and return the plan plus
    /// runtime metrics.
    ///
    /// **PROFILE executes the query for real.** Mutating statements
    /// affect the transaction's staged graph and are persisted at the
    /// usual `commit` point. Use [`Transaction::explain`] to inspect a
    /// plan without running it.
    pub fn profile(
        &mut self,
        query: &str,
        params: Option<BTreeMap<String, LoraValue>>,
    ) -> Result<QueryProfile, LoraError> {
        let params = params.unwrap_or_default();
        let compiled = self.compile_in_tx(query).map_err(LoraError::from_anyhow)?;
        let plan = query_plan_for(query, &compiled);
        let shape = plan.shape;

        let collector = Arc::new(MetricsCollector::new());
        let _guard = CollectorGuard::install(collector.clone());

        let started = Instant::now();
        let rows = self
            .execute_rows_compiled_deadline(&compiled, params, None)
            .map_err(LoraError::from_anyhow)?;
        let total_elapsed_ns = started.elapsed().as_nanos() as u64;

        drop(_guard);
        let per_operator = collector
            .snapshot()
            .into_iter()
            .map(|(id, op)| {
                (
                    id,
                    OperatorMetrics {
                        rows: op.rows,
                        elapsed_ns: op.elapsed_ns,
                        next_calls: op.next_calls,
                        db_hits: 0,
                    },
                )
            })
            .collect();

        let metrics = ProfileMetrics {
            total_elapsed_ns,
            total_rows: rows.len() as u64,
            mutated: shape.is_mutating(),
            per_operator,
        };

        Ok(QueryProfile { plan, metrics })
    }

    /// Execute a query inside the transaction and return an owning row stream.
    pub fn stream(&mut self, query: &str) -> Result<QueryStream<'static>, LoraError> {
        self.stream_with_params(query, BTreeMap::new())
    }

    /// Execute a parameterised query inside the transaction and return an
    /// owning row stream.
    pub fn stream_with_params(
        &mut self,
        query: &str,
        params: BTreeMap<String, LoraValue>,
    ) -> Result<QueryStream<'static>, LoraError> {
        let compiled = Arc::new(self.compile_in_tx(query)?);
        let columns = compiled_result_columns(&compiled);
        Ok(self.stream_compiled(compiled, columns, params)?)
    }

    /// Open a tx-bound stream for an already-compiled plan. Lets
    /// `Database::stream_with_params` reuse the plan it built for
    /// classification.
    pub(crate) fn stream_compiled(
        &mut self,
        compiled: Arc<CompiledQuery>,
        columns: Vec<String>,
        params: BTreeMap<String, LoraValue>,
    ) -> Result<QueryStream<'static>> {
        let mut inner = self.begin_statement()?;
        let is_mutating = classify_stream(&compiled).is_mutating();
        ensure_stream_allowed(&inner, is_mutating)?;

        let rollback_on_drop = stream_rolls_back_on_drop(is_mutating);
        let staged_ptr = self.prepare_stream_staging(&mut inner, rollback_on_drop)?;
        drop(inner);

        let cursor = open_tx_stream_cursor(staged_ptr, compiled, params, is_mutating);

        match cursor {
            Ok(cursor) => Ok(QueryStream::for_tx_cursor(
                cursor,
                columns,
                TxCursorLease::new(self.inner.clone(), rollback_on_drop),
            )),
            Err(err) => {
                finalize_tx_stream(&self.inner, TxStreamOutcome::Interrupted, rollback_on_drop);
                Err(err)
            }
        }
    }

    fn prepare_stream_staging(
        &self,
        inner: &mut MutexGuard<'_, TxInner>,
        rollback_on_drop: bool,
    ) -> Result<*mut InMemoryGraph> {
        // Transaction streams borrow from the staged graph. Even read-only
        // streams materialize staging when needed so the cursor can outlive
        // the `&mut Transaction` borrow without borrowing from the
        // transaction-owned live write guard.
        let clone_savepoint_graph = inner.staged.is_some();
        self.ensure_staged_locked(inner)?;
        inner.activate_cursor();

        if rollback_on_drop {
            inner.pending_savepoint = Some(take_savepoint(inner, clone_savepoint_graph));
        } else {
            inner.pending_savepoint = None;
        }

        Ok(inner
            .staged
            .as_mut()
            .ok_or(TransactionError::NoStagedGraph)? as *mut _)
    }

    /// Commit the transaction and publish staged changes.
    ///
    /// When WAL is attached the buffered tx-local mutation log is
    /// replayed into the durable WAL as a single committed
    /// transaction; recovery therefore observes either every write
    /// in this transaction or none.
    pub fn commit(mut self) -> Result<(), LoraError> {
        let CommitState {
            staged,
            buffer_events,
            mode,
        } = self.take_commit_state()?;

        let wrote_wal_commit = self.replay_commit_wal(mode, buffer_events)?;
        self.publish_staged_graph(mode, staged, wrote_wal_commit)?;

        self.live.take();
        Ok(())
    }

    fn take_commit_state(&self) -> Result<CommitState> {
        let mut inner = self.lock_inner()?;
        if inner.cursor_active {
            return Err(TransactionError::CursorActiveCommit.into());
        }
        if inner.closed {
            return Err(TransactionError::AlreadyClosed.into());
        }

        let mode = inner.mode;
        if inner.buffer_failed.load(Ordering::Acquire) {
            inner.closed = true;
            return Err(TransactionError::Poisoned.into());
        }
        let buffer_events = {
            let mut buffer = inner
                .buffer
                .lock()
                .map_err(|_| TransactionError::Poisoned)?;
            std::mem::take(&mut *buffer)
        };
        // Both modes can have `staged = None`: ReadOnly never clones,
        // and ReadWrite transactions that performed no writes leave
        // staging unmaterialized too.
        let staged = inner.staged.take();
        inner.closed = true;

        Ok(CommitState {
            staged,
            buffer_events,
            mode,
        })
    }

    fn replay_commit_wal(
        &self,
        mode: TransactionMode,
        buffer_events: Vec<MutationEvent>,
    ) -> Result<bool> {
        let Some(rec) = &self.wal else {
            return Ok(false);
        };

        if !matches!(mode, TransactionMode::ReadWrite) {
            ensure_wal_not_poisoned(rec)?;
            return Ok(false);
        }

        Ok(rec.commit_events(buffer_events)?.wrote())
    }

    fn publish_staged_graph(
        &mut self,
        mode: TransactionMode,
        staged: Option<InMemoryGraph>,
        wrote_wal_commit: bool,
    ) -> Result<()> {
        if !matches!(mode, TransactionMode::ReadWrite) {
            return Ok(());
        }

        let Some(mut staged) = staged else {
            return Ok(());
        };

        // Strip the buffering recorder from the staged graph before
        // publishing it as the live store; the live store either has
        // the durable WAL recorder reinstalled below or no recorder at
        // all (for non-WAL databases).
        staged.set_mutation_recorder(None);
        let wal = self.wal.clone();
        if let Some(rec) = &wal {
            staged.set_mutation_recorder(Some(rec.clone() as Arc<dyn MutationRecorder>));
        }

        let live = self.live.as_mut().ok_or(TransactionError::NoGraphGuard)?;
        let lease = match live {
            LiveStoreGuard::Write(lease) => lease,
            LiveStoreGuard::Read(_) => {
                return Err(TransactionError::ReadOnlyCommit.into());
            }
        };

        if wrote_wal_commit {
            if let (Some(snapshots), Some(rec)) = (&self.snapshots, wal.as_ref()) {
                snapshots.observe_commit(&staged, rec)?;
            }
        }

        // Atomic publish — concurrent readers will see the new state on
        // their next `load_full()`, while in-flight readers keep their
        // existing `Arc<InMemoryGraph>` snapshot until they drop it.
        lease.store.store(Arc::new(staged));

        Ok(())
    }

    /// Roll back the transaction. Staged graph changes and buffered
    /// mutations are discarded; the WAL is never armed.
    pub fn rollback(mut self) -> Result<(), LoraError> {
        let mut inner = self.lock_inner()?;
        if inner.closed {
            return Err(TransactionError::AlreadyClosed.into());
        }
        discard_transaction_state(&mut inner);
        drop(inner);
        self.live.take();
        Ok(())
    }

    /// Acquire the inner state for a new statement. Validates that
    /// the transaction is still open and no cursor is active. The
    /// staged graph is *not* required: ReadWrite transactions
    /// defer the staging clone until the first mutating statement
    /// (see [`Transaction::ensure_staged_locked`]).
    fn begin_statement(&self) -> Result<MutexGuard<'_, TxInner>> {
        let inner = self.lock_inner()?;
        if inner.closed {
            return Err(TransactionError::AlreadyClosed.into());
        }
        if inner.cursor_active {
            return Err(TransactionError::CursorActiveStatement.into());
        }
        Ok(inner)
    }

    /// Cheap state check for the ReadOnly fast path: closed +
    /// cursor_active. No staged-graph check — ReadOnly tx has no
    /// staged graph by construction.
    fn precheck_open_no_savepoint(&self) -> Result<()> {
        let inner = self.lock_inner()?;
        if inner.closed {
            return Err(TransactionError::AlreadyClosed.into());
        }
        if inner.cursor_active {
            return Err(TransactionError::CursorActiveStatement.into());
        }
        Ok(())
    }

    /// True if the transaction was begun in `ReadOnly` mode. Cheap
    /// — `mode` doesn't change after `begin_transaction`, so we
    /// pay one small state-lock acquisition.
    fn is_read_only_unchecked(&self) -> bool {
        matches!(self.mode, TransactionMode::ReadOnly)
    }

    fn lock_inner(&self) -> Result<MutexGuard<'_, TxInner>> {
        self.inner
            .lock()
            .map_err(|_| TransactionError::Poisoned.into())
    }

    pub(crate) fn release_streaming_cursor(&self) {
        if let Ok(mut inner) = self.inner.lock() {
            inner.release_cursor();
        }
    }
}

type ExecResultRows = Result<Vec<Row>>;

fn deadline_after(timeout: Duration) -> Instant {
    Instant::now()
        .checked_add(timeout)
        .unwrap_or_else(Instant::now)
}

fn execute_read_compiled(
    storage: &InMemoryGraph,
    compiled: &CompiledQuery,
    params: BTreeMap<String, LoraValue>,
    deadline: Option<Instant>,
) -> Result<Vec<Row>> {
    let executor = Executor::with_deadline(ExecutionContext { storage, params }, deadline);
    executor
        .execute_compiled_rows(compiled)
        .map_err(anyhow::Error::from)
}

fn execute_mutable_compiled(
    storage: &mut InMemoryGraph,
    compiled: &CompiledQuery,
    params: BTreeMap<String, LoraValue>,
    deadline: Option<Instant>,
) -> Result<Vec<Row>> {
    let mut executor =
        MutableExecutor::with_deadline(MutableExecutionContext { storage, params }, deadline);
    executor
        .execute_compiled_rows(compiled)
        .map_err(anyhow::Error::from)
}

fn query_plan_for(query: &str, compiled: &CompiledQuery) -> QueryPlan {
    QueryPlan {
        query: query.to_string(),
        tree: plan_tree_from_compiled(compiled),
        shape: classify_stream(compiled).into(),
        result_columns: plan_result_columns(&compiled.physical),
    }
}

fn ensure_stream_allowed(inner: &TxInner, is_mutating: bool) -> Result<()> {
    if matches!(inner.mode, TransactionMode::ReadOnly) && is_mutating {
        Err(TransactionError::ReadOnlyMutation.into())
    } else {
        Ok(())
    }
}

fn stream_rolls_back_on_drop(is_mutating: bool) -> bool {
    is_mutating
}

fn open_tx_stream_cursor(
    staged_ptr: *mut InMemoryGraph,
    compiled: Arc<CompiledQuery>,
    params: BTreeMap<String, LoraValue>,
    is_mutating: bool,
) -> Result<Box<dyn RowSource + 'static>> {
    // SAFETY: `stream_compiled` calls this only after `prepare_stream_staging`
    // has materialized `inner.staged` and activated the cursor guard. While the
    // guard is active, statements/commit/rollback cannot move or drop staged.
    let compiled_static: &'static CompiledQuery =
        unsafe { std::mem::transmute::<&CompiledQuery, _>(compiled.as_ref()) };

    if is_mutating {
        // SAFETY: see function-level note above. Mutating streams need exclusive
        // access, and the transaction cursor token prevents any second borrow.
        let storage_static: &'static mut InMemoryGraph = unsafe { &mut *staged_ptr };
        MutablePullExecutor::new(storage_static, params)
            .open_compiled(compiled_static)
            .map(|cursor| boxed_streaming_cursor(cursor, compiled))
            .map_err(anyhow::Error::from)
    } else {
        // SAFETY: see function-level note above. Read streams borrow staged
        // immutably until the cursor lease is finalized.
        let storage_static: &'static InMemoryGraph = unsafe { &*staged_ptr };
        PullExecutor::new(storage_static, params)
            .open_compiled(compiled_static)
            .map(|cursor| boxed_streaming_cursor(cursor, compiled))
            .map_err(anyhow::Error::from)
    }
}

fn boxed_streaming_cursor(
    cursor: Box<dyn RowSource + 'static>,
    compiled: Arc<CompiledQuery>,
) -> Box<dyn RowSource + 'static> {
    Box::new(StreamingCursorWithArc {
        cursor,
        _compiled: compiled,
    })
}

struct CommitState {
    staged: Option<InMemoryGraph>,
    buffer_events: Vec<MutationEvent>,
    mode: TransactionMode,
}

impl TxInner {
    fn staged_mut(&mut self) -> Result<&mut InMemoryGraph> {
        self.staged
            .as_mut()
            .ok_or(TransactionError::NoStagedGraph.into())
    }

    fn activate_cursor(&mut self) {
        self.cursor_active = true;
    }

    fn release_cursor(&mut self) {
        self.cursor_active = false;
    }

    fn clear_pending_savepoint(&mut self) {
        self.pending_savepoint = None;
    }

    fn restore_pending_savepoint(&mut self) {
        if let Some(sp) = self.pending_savepoint.take() {
            apply_savepoint(self, sp);
        }
    }

    fn finalize_stream(&mut self, outcome: TxStreamOutcome, rollback_on_drop: bool) {
        self.release_cursor();

        if self.closed {
            discard_transaction_state(self);
            return;
        }

        if outcome.should_restore_savepoint(rollback_on_drop) {
            self.restore_pending_savepoint();
        } else {
            self.clear_pending_savepoint();
        }
    }
}

/// `RowSource` adapter that owns an `Arc<CompiledQuery>` so the
/// inner cursor's `'static` borrows into the plan stay valid for the
/// life of the wrapper. The inner cursor is stored first so it drops
/// before the Arc, releasing any borrows back into the plan.
struct StreamingCursorWithArc {
    cursor: Box<dyn RowSource + 'static>,
    _compiled: Arc<CompiledQuery>,
}

impl RowSource for StreamingCursorWithArc {
    fn next_row(&mut self) -> lora_executor::ExecResult<Option<Row>> {
        self.cursor.next_row()
    }
}

fn finalize_tx_stream(
    handle: &Arc<Mutex<TxInner>>,
    outcome: TxStreamOutcome,
    rollback_on_drop: bool,
) {
    if let Ok(mut inner) = handle.lock() {
        inner.finalize_stream(outcome, rollback_on_drop);
    }
}

fn discard_transaction_state(inner: &mut TxInner) {
    // A full transaction rollback supersedes any pending cursor savepoint.
    inner.clear_pending_savepoint();
    inner.release_cursor();
    inner.staged = None;
    if let Ok(mut buf) = inner.buffer.lock() {
        buf.clear();
    } else {
        inner.buffer_failed.store(true, Ordering::Release);
    }
    inner.closed = true;
}

fn take_savepoint(inner: &TxInner, clone_staged: bool) -> Savepoint {
    let buffer_len = match inner.buffer.lock() {
        Ok(buffer) => buffer.len(),
        Err(_) => {
            inner.buffer_failed.store(true, Ordering::Release);
            0
        }
    };
    Savepoint {
        staged: if clone_staged {
            inner.staged.as_ref().cloned()
        } else {
            None
        },
        buffer_len,
    }
}

fn restore_savepoint(inner: &mut TxInner, savepoint: Option<Savepoint>) {
    if let Some(sp) = savepoint {
        apply_savepoint(inner, sp);
    }
}

fn apply_savepoint(inner: &mut TxInner, sp: Savepoint) {
    if let Ok(mut buf) = inner.buffer.lock() {
        buf.truncate(sp.buffer_len);
    } else {
        inner.buffer_failed.store(true, Ordering::Release);
    }

    let Some(mut graph) = sp.staged else {
        inner.staged = None;
        return;
    };

    // Rebuild the staged graph from the snapshot and re-install the
    // buffering recorder. `InMemoryGraph::clone` deliberately drops
    // recorders, so the snapshot has none until we put it back.
    if matches!(inner.mode, TransactionMode::ReadWrite) && inner.buffer_mutations {
        graph.set_mutation_recorder(Some(Arc::new(BufferingRecorder::new(
            inner.buffer.clone(),
            inner.buffer_failed.clone(),
        )) as Arc<dyn MutationRecorder>));
    }
    inner.staged = Some(graph);
}

impl Drop for Transaction<'_> {
    fn drop(&mut self) {
        // If the user never called commit/rollback, treat it as a
        // rollback: drop staged changes and the buffered mutations.
        // The live RwLock guard is released as part of dropping `self.live`.
        if let Ok(mut inner) = self.inner.lock() {
            if !inner.closed {
                if inner.cursor_active {
                    // A tx-bound stream may still be borrowing the
                    // staged graph through `inner`. Leave that graph
                    // in place until the stream drops, but mark the
                    // transaction closed so finalization discards it
                    // instead of making it commit-eligible.
                    inner.closed = true;
                } else {
                    discard_transaction_state(&mut inner);
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use super::*;

    #[test]
    fn buffering_recorder_latches_poisoned_buffer() {
        let buffer = Arc::new(Mutex::new(Vec::new()));
        let failed = Arc::new(AtomicBool::new(false));

        let poisoned_buffer = buffer.clone();
        let _ = thread::spawn(move || {
            let _guard = poisoned_buffer.lock().unwrap();
            panic!("poison mutation buffer");
        })
        .join();

        let recorder = BufferingRecorder::new(buffer, failed.clone());
        recorder.record(MutationEvent::Clear);

        assert!(failed.load(Ordering::Acquire));
    }
}