bashkit 0.18.0

//! Thin wrapper around `turso_core` that hides the choice of `IO` backend.
//!
//! The two backends correspond to the two phases described in
//! `knowledge/runtimes/sqlite-builtin.md`:
//!
//! - **Phase 1** ([`Backend::Memory`]): use turso's `MemoryIO` and snapshot the
//!   raw database bytes for caller-driven persistence. This is what the
//!   builtin uses when the caller asks for `:memory:` databases or wants to
//!   load/flush the entire DB file from the VFS at command boundaries.
//!
//! - **Phase 2** ([`Backend::Vfs`]): use [`super::vfs_io::BashkitVfsIO`], which
//!   reads/writes through bashkit's `Arc<dyn FileSystem>` and persists dirty
//!   bytes back on `flush_dirty`.
//!
//! Both expose the same query API. The builtin layer above is agnostic to
//! which backend is active.

use crate::time_compat::Instant;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex};

use turso_core::{
    Connection, Database, IO, MemoryIO, Numeric, OpenFlags, SqliteDialect, StepResult, Value,
};

use super::vfs_io::BashkitVfsIO;

/// Shared wall-clock budget for an engine's lifetime. Each `step()` cycle
/// checks whether the deadline has passed and trips an interrupt if so.
#[derive(Debug, Clone, Copy)]
pub(super) struct Deadline {
    pub deadline: Option<Instant>,
}

impl Deadline {
    pub(super) fn new(max_duration: std::time::Duration) -> Self {
        // Treat zero/near-zero as "no deadline" — useful in tests and for
        // operators that explicitly opt out via `Duration::ZERO`.
        let deadline = if max_duration.is_zero() {
            None
        } else {
            Some(Instant::now() + max_duration)
        };
        Self { deadline }
    }

    /// Returns true once the budget is exhausted.
    pub(super) fn expired(&self) -> bool {
        self.deadline.map(|d| Instant::now() >= d).unwrap_or(false)
    }
}

/// VM instructions between progress-handler callbacks.
///
/// Design: the step loop in [`SqliteEngine::execute`] can only check limits
/// *between* `Statement::step()` calls, and a step returns only when a row is
/// produced, the program halts, or IO is needed. A query whose rows are all
/// consumed inside the VDBE (a recursive CTE feeding an aggregate, a filtered
/// cross join) therefore performs unbounded work inside one step. Turso's
/// SQLite-compatible progress handler is consulted every N VM instructions
/// *inside* the instruction loop, which is the only place work of that shape
/// can be counted and interrupted. `THREAT[TM-SQL-014]`.
///
/// 1024 keeps the callback cost in the noise (one `Instant::now()` plus two
/// atomic loads per 1024 instructions) while bounding overshoot past the
/// deadline to microseconds.
const PROGRESS_HANDLER_OPS: u64 = 1024;

/// Why the VDBE was interrupted from inside a step, recorded by the progress
/// handler so the step loop can report the real limit rather than a generic
/// "interrupted".
type InterruptReason = Arc<Mutex<Option<String>>>;

/// Installs a progress handler on a connection and clears it on drop, so an
/// early `?` return never leaves a stale callback (and its cloned budget)
/// attached to a pooled connection.
struct ProgressHandlerGuard<'a> {
    conn: &'a Connection,
}

impl<'a> ProgressHandlerGuard<'a> {
    /// Install a handler that interrupts the VDBE once the invocation deadline
    /// passes or the request's execution budget is exhausted.
    fn install(
        conn: &'a Connection,
        deadline: Deadline,
        budget: Option<crate::limits::ExecutionBudget>,
    ) -> (Self, InterruptReason) {
        let reason: InterruptReason = Arc::new(Mutex::new(None));
        let callback_reason = Arc::clone(&reason);
        conn.set_progress_handler(
            PROGRESS_HANDLER_OPS,
            Some(Box::new(move || {
                let stop = if deadline.expired() {
                    Some("query timed out".to_string())
                } else if let Some(budget) = &budget {
                    budget
                        .consume_work(1)
                        .err()
                        .map(|failure| failure.to_string())
                } else {
                    None
                };
                match stop {
                    Some(message) => {
                        let mut slot = callback_reason.lock().expect("interrupt reason lock");
                        slot.get_or_insert(message);
                        true
                    }
                    None => false,
                }
            })),
        );
        (Self { conn }, reason)
    }
}

impl Drop for ProgressHandlerGuard<'_> {
    fn drop(&mut self) {
        self.conn.set_progress_handler(0, None);
    }
}

/// Result alias for engine operations. The error string is intended to be
/// shown directly to the user via `ExecResult::err`, so it should not include
/// host paths or other sensitive details.
pub(super) type EngineResult<T> = std::result::Result<T, String>;

/// Each engine creates its own unique in-memory path. Turso bypasses its
/// process-wide `DATABASE_MANAGER` registry for paths starting with `:memory:`,
/// so we prefix with that to keep concurrent engines isolated even when
/// multiple instances live in the same process (e.g. parallel tests).
fn unique_memory_path() -> String {
    static COUNTER: AtomicU64 = AtomicU64::new(0);
    let n = COUNTER.fetch_add(1, Ordering::Relaxed);
    format!(":memory:bashkit-{n}")
}

/// Selects which `IO` impl backs the engine.
pub(super) enum Backend {
    /// Pure in-process `MemoryIO`. The owner of the engine is responsible for
    /// calling [`SqliteEngine::snapshot_bytes`] when it wants to persist.
    Memory(Arc<MemoryIO>),
    /// VFS-backed `BashkitVfsIO`. The owner calls `flush_dirty` on the IO
    /// when it wants the in-memory pages flushed back to the VFS.
    Vfs(Arc<BashkitVfsIO>),
}

/// Query materialisation caps applied before values are cloned into Bashkit-owned memory.
#[derive(Debug, Clone)]
pub(super) struct QueryLimits {
    pub max_rows: usize,
    pub max_value_bytes: usize,
    pub max_result_bytes: usize,
    pub execution_budget: Option<crate::limits::ExecutionBudget>,
}

/// Outcome of executing a single SQL statement.
#[derive(Debug, Default)]
pub(super) struct StatementOutcome {
    /// Column names if the statement produced a result set.
    pub columns: Vec<String>,
    /// Rows materialised from the result set. Empty for non-SELECT statements.
    pub rows: Vec<Vec<Value>>,
    /// Number of rows changed (for INSERT/UPDATE/DELETE; 0 otherwise).
    pub changes: i64,
}

/// Wraps a turso `Database`/`Connection` pair plus the backing `IO`.
pub(super) struct SqliteEngine {
    backend: Backend,
    _db: Arc<Database>,
    conn: Arc<Connection>,
    /// Path used to register the database file inside the IO. For
    /// [`Backend::Memory`] this is a unique `:memory:bashkit-N` string so
    /// concurrent engines never share state through turso's process-wide
    /// `DATABASE_MANAGER` registry.
    memory_path: Option<String>,
}

impl SqliteEngine {
    /// Open a fresh in-memory database. If `initial_bytes` is `Some`, the
    /// bytes are written into a temporary file inside the `MemoryIO` first
    /// so that turso opens an existing database rather than a blank one.
    ///
    /// We always route through a named in-memory file (rather than `:memory:`)
    /// so that [`SqliteEngine::snapshot_bytes`] can read the resulting
    /// database pages back. When there are no initial bytes we still seed an
    /// empty file to ensure the path exists.
    pub(super) fn open_memory(initial_bytes: Option<&[u8]>) -> EngineResult<Self> {
        let io: Arc<MemoryIO> = Arc::new(MemoryIO::new());
        let path = unique_memory_path();
        if let Some(bytes) = initial_bytes
            && !bytes.is_empty()
        {
            seed_memory_io(&io, &path, bytes).map_err(turso_msg)?;
        }
        let io_dyn: Arc<dyn IO> = io.clone();
        let db = Database::open_file(io_dyn, &path, Arc::new(SqliteDialect)).map_err(turso_msg)?;
        let conn = db.connect().map_err(turso_msg)?;
        Ok(Self {
            backend: Backend::Memory(io),
            _db: db,
            conn,
            memory_path: Some(path),
        })
    }

    /// Open a true `:memory:` database (no file backing, no persistence).
    /// Use this when the caller never intends to extract bytes.
    pub(super) fn open_pure_memory() -> EngineResult<Self> {
        let io: Arc<MemoryIO> = Arc::new(MemoryIO::new());
        let io_dyn: Arc<dyn IO> = io.clone();
        let db =
            Database::open_file(io_dyn, ":memory:", Arc::new(SqliteDialect)).map_err(turso_msg)?;
        let conn = db.connect().map_err(turso_msg)?;
        Ok(Self {
            backend: Backend::Memory(io),
            _db: db,
            conn,
            memory_path: None,
        })
    }

    /// Open a database backed by the bashkit VFS via [`BashkitVfsIO`].
    /// `path_in_io` is the path string passed verbatim to turso (and used as
    /// a key in the VFS).
    pub(super) fn open_vfs(io: Arc<BashkitVfsIO>, path_in_io: &str) -> EngineResult<Self> {
        let io_dyn: Arc<dyn IO> = io.clone();
        let db =
            Database::open_file(io_dyn, path_in_io, Arc::new(SqliteDialect)).map_err(turso_msg)?;
        let conn = db.connect().map_err(turso_msg)?;
        Ok(Self {
            backend: Backend::Vfs(io),
            _db: db,
            conn,
            memory_path: Some(path_in_io.to_string()),
        })
    }

    /// Execute a single statement, materialising rows up-front so that the
    /// caller doesn't need to drive the step loop.
    ///
    /// `deadline` carries the wall-clock budget shared across all statements
    /// in this invocation. Once it expires, we issue `stmt.interrupt()` and
    /// return a timeout error rather than continuing the step loop. Work done
    /// inside a single `step()` is bounded by the progress handler installed
    /// here, which charges the execution budget one unit per
    /// [`PROGRESS_HANDLER_OPS`] VM instructions and interrupts the VDBE when
    /// the deadline or the budget is exhausted.
    pub(super) fn execute(
        &self,
        sql: &str,
        deadline: Deadline,
        limits: QueryLimits,
    ) -> EngineResult<StatementOutcome> {
        // Bounds work done *inside* a step; the loop below bounds work
        // *between* steps. Both are needed: see `PROGRESS_HANDLER_OPS`.
        let (_progress_guard, interrupt_reason) =
            ProgressHandlerGuard::install(&self.conn, deadline, limits.execution_budget.clone());
        let mut stmt = self.conn.prepare(sql).map_err(turso_msg)?;
        let mut outcome = StatementOutcome::default();
        for idx in 0..stmt.num_columns() {
            outcome.columns.push(stmt.get_column_name(idx).to_string());
        }
        let mut result_bytes = 0usize;
        loop {
            if let Some(budget) = &limits.execution_budget {
                budget.consume_work(1).map_err(|e| e.to_string())?;
            }
            if deadline.expired() {
                stmt.interrupt();
                return Err("query timed out".to_string());
            }
            match stmt.step().map_err(turso_msg)? {
                StepResult::Row => {
                    let next_row = outcome.rows.len() + 1;
                    if next_row > limits.max_rows {
                        return Err(format!(
                            "result set exceeds row cap ({next_row} > {})",
                            limits.max_rows
                        ));
                    }
                    if let Some(row) = stmt.row() {
                        let mut row_bytes = 0usize;
                        for idx in 0..stmt.num_columns() {
                            let value_bytes = Self::value_size_bytes(row.get_value(idx));
                            if value_bytes > limits.max_value_bytes {
                                return Err(format!(
                                    "result value exceeds byte cap ({value_bytes} > {})",
                                    limits.max_value_bytes
                                ));
                            }
                            row_bytes = row_bytes.saturating_add(value_bytes);
                        }
                        let next_result_bytes = result_bytes.saturating_add(row_bytes);
                        if next_result_bytes > limits.max_result_bytes {
                            return Err(format!(
                                "result set exceeds byte cap ({next_result_bytes} > {})",
                                limits.max_result_bytes
                            ));
                        }
                        result_bytes = next_result_bytes;
                        let values: Vec<Value> = (0..stmt.num_columns())
                            .map(|idx| row.get_value(idx).clone())
                            .collect();
                        outcome.rows.push(values);
                    }
                }
                StepResult::Done => break,
                // `IO` means the VDBE is blocked on outstanding completions;
                // `Yield` (added in turso 0.7) means it voluntarily paused to let
                // the program state machines advance. `Sleep` (added in turso
                // 0.8.0-pre.4) asks us to back off for a duration before
                // re-stepping, e.g. after a busy handler chose to retry. In all
                // three cases the caller drives pending completions and re-steps.
                // Our backends are synchronous, so `io_step` returns immediately
                // when nothing is pending, and the deadline check above bounds
                // the loop.
                //
                // We deliberately do not honour `Sleep`'s duration. Blocking the
                // thread would be wrong here: the engine is single-threaded and
                // driven synchronously, so nothing else can make the progress the
                // sleep is waiting for, and `std::thread::sleep` is unavailable on
                // the wasm targets this builtin also ships to. Upstream documents
                // that callers which don't track time may treat `Sleep` exactly
                // like `IO`. Spinning is bounded by both the wall-clock deadline
                // (TM-SQL-005a) and the per-iteration execution budget above.
                StepResult::IO | StepResult::Yield | StepResult::Sleep { .. } => {
                    self.io_step()?;
                }
                // The progress handler records why it asked for the
                // interrupt, so a deadline or budget breach inside one step
                // reports the same message it would between steps.
                StepResult::Interrupt => {
                    let reason = interrupt_reason
                        .lock()
                        .expect("interrupt reason lock")
                        .take();
                    return Err(reason.unwrap_or_else(|| "query was interrupted".to_string()));
                }
                StepResult::Busy => {
                    return Err("database is busy".to_string());
                }
            }
        }
        outcome.changes = self.conn.changes();
        Ok(outcome)
    }

    fn value_size_bytes(value: &Value) -> usize {
        match value {
            Value::Null => 0,
            Value::Numeric(Numeric::Integer(_)) | Value::Numeric(Numeric::Float(_)) => 8,
            Value::Text(text) => text.as_str().len(),
            Value::Blob(bytes) => bytes.len(),
        }
    }

    fn io_step(&self) -> EngineResult<()> {
        match &self.backend {
            Backend::Memory(io) => io.step().map_err(turso_msg),
            Backend::Vfs(io) => io.step().map_err(turso_msg),
        }
    }

    /// Snapshot the current database file bytes for cache invalidation.
    ///
    /// We force a TRUNCATE-mode checkpoint before reading so that any pages
    /// still in the WAL are folded into the main file. Without this step the
    /// snapshot would be missing the just-written transaction.
    pub(super) fn snapshot_bytes(&self) -> Option<Vec<u8>> {
        let path = self.memory_path.as_deref()?;
        let _ = self.conn.checkpoint(turso_core::CheckpointMode::Truncate {
            upper_bound_inclusive: None,
        });
        match &self.backend {
            Backend::Memory(io) => {
                let file = io.open_file(path, OpenFlags::None, false).ok()?;
                let size = file.size().ok()? as usize;
                if size == 0 {
                    return Some(Vec::new());
                }
                Some(read_all(&file, size))
            }
            Backend::Vfs(io) => io.file_bytes(path),
        }
    }

    /// For the VFS backend, flush any pages dirtied in memory back to the
    /// underlying `FileSystem`. Returns the number of files persisted.
    pub(super) async fn flush_dirty(&self) -> EngineResult<usize> {
        match &self.backend {
            Backend::Memory(_) => Ok(0),
            Backend::Vfs(io) => io.flush_dirty().await,
        }
    }

    /// Close the connection, releasing any cached pages. Best-effort.
    pub(super) fn close(&self) {
        let _ = self.conn.close();
    }
}

impl Drop for SqliteEngine {
    fn drop(&mut self) {
        // turso's Connection has its own Drop, but we want to be explicit
        // about checkpoints to keep the on-disk image consistent.
        self.close();
    }
}

fn read_all(file: &Arc<dyn turso_core::File>, size: usize) -> Vec<u8> {
    use turso_core::{Buffer, Completion};
    let mut out = vec![0u8; size];
    let chunk_size: usize = 4096;
    let mut pos = 0usize;
    while pos < size {
        let remaining = size - pos;
        let take = remaining.min(chunk_size);
        let chunk = Arc::new(Buffer::new(vec![0u8; take]));
        // The completion runs synchronously for MemoryIO; the closure receives
        // the buffer back via the Result tuple. We copy bytes after pread()
        // returns rather than from the closure, since the closure has to be
        // 'static and copying from there is awkward.
        let completion = Completion::new_read(chunk.clone(), |_res| None);
        let _ = file.pread(pos as u64, completion);
        out[pos..pos + take].copy_from_slice(&chunk.as_slice()[..take]);
        pos += take;
    }
    out
}

/// Pre-seed a `MemoryIO`-backed file with bytes by writing them as a single
/// `pwrite` operation. This is how we hand turso an existing database image.
fn seed_memory_io(
    io: &Arc<MemoryIO>,
    path: &str,
    bytes: &[u8],
) -> std::result::Result<(), turso_core::LimboError> {
    use turso_core::{Buffer, Completion, OpenFlags};
    let file = io.open_file(path, OpenFlags::Create, false)?;
    if bytes.is_empty() {
        return Ok(());
    }
    let buf = Arc::new(Buffer::new(bytes.to_vec()));
    let completion = Completion::new_write(|_| {});
    let _completion = file.pwrite(0, buf, completion)?;
    Ok(())
}

/// Map a turso error to a sanitised user-facing string.
fn turso_msg(e: turso_core::LimboError) -> String {
    e.to_string()
}