akar-main 0.2.0

Akar - pure Rust embedded graph database for AI agent memory
//! Connection — used to execute queries against a Database.
//!
//! Manages the full query lifecycle: parse → bind → plan → optimize → execute.
//! DDL statements (CREATE/DROP TABLE) update the catalog directly and return
//! a message result. DML statements (MATCH/RETURN) produce DataChunk results.
//!
//! Supports prepared statements via `prepare()` and `execute()` for
//! parameterized queries.
//!
//! # Concurrent Multi-Writer Support
//!
//! Write transactions use `TransactionManager::begin_write()` / `commit()` /
//! `rollback()` with per-transaction `LocalStorage`, `LocalWAL`, and
//! `ShadowFile` resources held in `txn_resources`. The full commit pipeline
//! flushes: LocalStorage → tables, LocalWAL → global WAL, ShadowFile → BM.

pub mod copy;
pub mod ddl;
pub mod dml;
pub mod fts_estimate;
pub mod plan_cache;
pub mod query;
pub mod standalone_call;
pub mod substitute;
pub mod transaction;
pub mod utils;

use crate::database::Database;
use crate::prepared_statement::PreparedStatement;
use plan_cache::{CachedPlan, PlanCache};
use std::collections::HashMap;
use std::sync::{Arc, Mutex};

/// Default capacity for the query plan LRU cache.
const PLAN_CACHE_CAPACITY: usize = 100;

/// Per-transaction resources held during an active write transaction.
pub(crate) struct TxnResources {
    pub local_storage: akar_storage::LocalStorage,
    pub local_wal: akar_storage::LocalWAL,
    pub shadow_file: akar_storage::ShadowFile,
}

/// A connection to the database for executing Cypher queries.
///
/// Created via [`Connection::new`] with a reference to a [`Database`].
/// Supports both ad-hoc queries and prepared statements.
///
/// # Examples
///
/// ```no_run
/// use akar_main::database::{Database, SystemConfig};
/// use akar_main::connection::Connection;
///
/// let db = std::sync::Arc::new(Database::new("./my_db", SystemConfig::default())?);
/// let conn = Connection::new(&db);
///
/// // DDL
/// conn.query("CREATE NODE TABLE Person(name STRING, PRIMARY KEY(name))")?;
///
/// // DML
/// conn.query("CREATE (:Person {name: 'Alice'})")?;
///
/// // Query
/// let result = conn.query("MATCH (p:Person) RETURN p.name")?;
/// assert!(result.success);
/// # Ok::<(), String>(())
/// ```
pub struct Connection {
    pub(crate) database: Arc<Database>,
    /// Cache of prepared statements (query → PreparedStatement).
    pub(crate) statement_cache: Mutex<HashMap<String, PreparedStatement>>,
    /// LRU cache of optimized query plans keyed by normalized query string.
    /// Entries are validated against the catalog version on lookup, so DDL
    /// (via any connection) invalidates them implicitly.
    pub(crate) plan_cache: Mutex<PlanCache<CachedPlan>>,
    /// Per-transaction resources keyed by transaction ID.
    /// Set up when `begin_write()` is called, cleaned up on commit/rollback.
    pub(crate) txn_resources: Mutex<HashMap<u64, TxnResources>>,
    /// Set while an explicit `BEGIN TRANSACTION` is open on this connection.
    /// DDL statements bypass the txn's LocalStorage/ShadowFile and mutate the
    /// catalog directly, so they cannot be rolled back — they are rejected
    /// while this flag is set instead of falsely reporting success (P52.28).
    pub(crate) explicit_txn_active: std::sync::atomic::AtomicBool,
    /// Lazily-built processor handler callbacks (sequence, schema DDL, query,
    /// subquery, standalone-call registry), shared across every query on this
    /// connection. Cached here rather than on `Database` so the handlers'
    /// strong `Arc<Database>` captures do not create a reference cycle that
    /// would keep the database (and its file lock) alive forever (P51.47).
    processor_handlers: std::sync::OnceLock<Arc<crate::connection::query::ProcessorHandlers>>,
}

impl Connection {
    /// Create a new connection to the given database.
    ///
    /// A connection owns a statement cache and transaction context.
    /// Multiple connections can coexist on the same `Database` — they
    /// share the buffer pool and catalog but hold independent transactions.
    pub fn new(database: &Arc<Database>) -> Self {
        Self {
            database: database.clone(),
            statement_cache: Mutex::new(HashMap::new()),
            plan_cache: Mutex::new(PlanCache::new(PLAN_CACHE_CAPACITY)),
            txn_resources: Mutex::new(HashMap::new()),
            explicit_txn_active: std::sync::atomic::AtomicBool::new(false),
            processor_handlers: std::sync::OnceLock::new(),
        }
    }

    /// Clear the prepared statement and plan caches.
    pub fn clear_cache(&self) {
        if let Ok(mut cache) = self.statement_cache.lock() {
            cache.clear();
        }
        if let Ok(mut cache) = self.plan_cache.lock() {
            cache.clear();
        }
    }

    /// Number of cached prepared statements.
    pub fn cache_size(&self) -> usize {
        self.statement_cache.lock().map(|c| c.len()).unwrap_or(0)
    }

    /// Number of cached query plans.
    pub fn plan_cache_size(&self) -> usize {
        self.plan_cache.lock().map(|c| c.len()).unwrap_or(0)
    }
}

impl Drop for Connection {
    fn drop(&mut self) {
        // A dropped connection must roll back any abandoned write transactions.
        // Without this, a client that disconnects after `BEGIN` (without
        // COMMIT/ROLLBACK) leaves `active_write_count` raised and table locks
        // held forever, wedging single-writer mode (P52.15).
        let txn_ids: Vec<u64> = self
            .txn_resources
            .lock()
            .map(|map| map.keys().copied().collect())
            .unwrap_or_default();
        if txn_ids.is_empty() {
            return;
        }

        let tm = &self.database.transaction_manager;
        let txns: Vec<akar_transaction::Transaction> = tm
            .active_snapshot()
            .ok()
            .map(|mut active| txn_ids.iter().filter_map(|id| active.remove(id)).collect())
            .unwrap_or_default();

        for mut txn in txns {
            let id = txn.transaction_id;
            match self.rollback_write_txn(&mut txn) {
                Ok(_) => tracing::info!("Rolled back abandoned txn#{id} on connection drop"),
                Err(e) => tracing::warn!("Failed to roll back abandoned txn#{id} on connection drop: {e}"),
            }
        }
    }
}