use super::Connection;
use super::plan_cache::{CachedPlan, normalize_query};
use crate::prepared_statement::PreparedStatement;
use crate::query_result::QueryResult;
use akar_binder::Binder;
use akar_binder::bound_statement::BoundStatement;
use akar_common::error::ProcessorError;
use akar_common::types::Value;
use akar_optimizer::Optimizer;
use akar_parser::parse;
use akar_planner::QueryPlanner;
use akar_planner::logical_operator::LogicalOperator;
use akar_processor::QueryProcessor;
use std::collections::HashMap;
use std::sync::Arc;
impl Connection {
pub fn query(&self, query_str: &str) -> Result<QueryResult, String> {
let trimmed = query_str.trim();
if trimmed.is_empty() {
return Ok(QueryResult::new(Vec::new()));
}
if let Some(value) = trimmed
.strip_prefix("SET")
.and_then(|s| s.trim().strip_prefix("spill_threshold"))
.and_then(|s| s.trim().strip_prefix("="))
.map(|s| s.trim())
{
let bytes: u64 = value.parse().map_err(|_| {
format!("Invalid spill_threshold value '{value}'. Expected a positive integer (bytes).")
})?;
self.database.set_spill_threshold(bytes);
return Ok(QueryResult::success_message(format!(
"spill_threshold set to {bytes} bytes"
)));
}
if let Some(value) = trimmed
.strip_prefix("SET")
.and_then(|s| s.trim().strip_prefix("concurrent_writes"))
.and_then(|s| s.trim().strip_prefix("="))
.map(|s| s.trim())
{
let enabled = match value.to_lowercase().as_str() {
"true" | "1" | "yes" => true,
"false" | "0" | "no" => false,
_ => return Err("Invalid value for concurrent_writes. Use true or false.".into()),
};
self.database.transaction_manager.set_concurrent_writes(enabled);
return Ok(QueryResult::success_message(format!(
"concurrent_writes set to {enabled}"
)));
}
let normalized = normalize_query(trimmed);
let catalog_version = self
.database
.catalog
.lock()
.map_err(|e| format!("Catalog lock error: {e}"))?
.version();
{
let mut cache = self.plan_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
if let Some(cached) = cache.get(&normalized).filter(|c| c.catalog_version == catalog_version) {
let bound = cached.bound.clone();
let plan = cached.plan.clone();
drop(cache);
return self.execute_with_plan(&bound, Some(&plan));
}
}
let statement = parse(trimmed).map_err(|e| format!("Parse error: {e}"))?;
let binder = Binder::new(self.database.catalog.clone());
let bound = binder.bind(statement).map_err(|e| format!("Bind error: {e}"))?;
let plan_opt: Option<Vec<LogicalOperator>> = if is_plan_cachable(&bound) {
let plan = self.build_optimized_plan(&bound)?;
let mut cache = self.plan_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
cache.insert(
normalized,
CachedPlan {
bound: bound.clone(),
plan: plan.clone(),
catalog_version,
},
);
Some(plan)
} else {
None
};
self.execute_with_plan(&bound, plan_opt.as_ref())
}
fn execute_with_plan(
&self,
bound: &BoundStatement,
plan: Option<&Vec<LogicalOperator>>,
) -> Result<QueryResult, String> {
if self.database.config.read_only && Connection::is_write_statement(bound) {
return Err("Database is in read-only mode; write statements are not allowed".into());
}
let concurrent_mode = self.database.transaction_manager.allow_concurrent_writes();
let is_write = concurrent_mode && Connection::is_write_statement(bound);
let mut txn_opt: Option<akar_transaction::Transaction> =
if is_write { Some(self.begin_write_txn()?) } else { None };
let query_result = self.execute_query_inner(bound, txn_opt.as_mut(), plan);
match (is_write, &query_result) {
(true, Ok(_)) => {
if let Some(ref mut txn) = txn_opt {
self.commit_write_txn(txn)?;
}
}
(true, Err(e)) => {
if let Some(ref mut txn) = txn_opt {
match self.rollback_write_txn(txn) {
Ok(_records) => {
tracing::warn!("Transaction rolled back due to error: {e}");
}
Err(rollback_err) => {
tracing::error!("Transaction rollback ALSO failed: {rollback_err} (original error: {e})");
}
}
}
}
_ => {}
}
query_result
}
fn build_optimized_plan(&self, bound: &BoundStatement) -> Result<Vec<LogicalOperator>, String> {
let planner = QueryPlanner::new();
let logical_plan = planner.plan(bound.clone()).map_err(|e| format!("Plan error: {e}"))?;
let optimizer = Optimizer::with_stats(self.database.stats_store.clone());
Ok(optimizer.optimize(logical_plan))
}
pub(crate) fn execute_query_inner(
&self,
bound: &BoundStatement,
txn_opt: Option<&mut akar_transaction::Transaction>,
cached_plan: Option<&Vec<LogicalOperator>>,
) -> Result<QueryResult, String> {
if let Some(result) = self.handle_ddl(bound)? {
self.database.persist_catalog()?;
self.maybe_auto_checkpoint()?;
return Ok(result);
}
if let Some(ref txn) = txn_opt {
if !self.database.transaction_manager.allow_concurrent_writes() {
let write_tables = Connection::extract_write_tables(bound);
for tid in write_tables {
self.database.transaction_manager.lock_table(txn.transaction_id, tid)?;
}
}
}
let optimized_plan: Vec<LogicalOperator> = match cached_plan {
Some(plan) => plan.clone(),
None => self.build_optimized_plan(bound)?,
};
if optimized_plan.is_empty() {
return Ok(QueryResult::success_message("Query executed (no result)".into()));
}
let (snapshot_ts, commit_history) = if let Some(ref txn) = txn_opt {
(
txn.snapshot_ts,
self.database.transaction_manager.commit_history_snapshot(),
)
} else {
let ts = self.database.transaction_manager.current_commit_ts();
let history = self.database.transaction_manager.commit_history_snapshot();
(Some(ts), history)
};
let processor = self.create_processor().with_snapshot(snapshot_ts, commit_history);
let chunks = processor
.execute(&optimized_plan)
.map_err(|e| format!("Execute error: {e}"))?;
if let Some(ref txn) = txn_opt {
let written_rows = processor.take_written_rows();
let tm = &self.database.transaction_manager;
for (table_id, row_id) in written_rows {
tm.record_write(txn.transaction_id, table_id, row_id);
}
}
if txn_opt.is_none() && Connection::is_write_statement(bound) {
self.database
.storage_manager
.persist_all_tables()
.map_err(|e| format!("Failed to persist tables: {e}"))?;
}
self.maybe_auto_checkpoint()?;
Ok(QueryResult::new(chunks))
}
pub fn prepare(&self, query_str: &str) -> Result<PreparedStatement, String> {
let trimmed = query_str.trim();
{
let cache = self.statement_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
if let Some(cached) = cache.get(trimmed) {
return Ok(cached.clone());
}
}
let statement = parse(trimmed).map_err(|e| format!("Parse error: {e}"))?;
let binder = Binder::new(self.database.catalog.clone());
let bound = binder.bind(statement).map_err(|e| format!("Bind error: {e}"))?;
let prepared = PreparedStatement::new(trimmed.to_string(), bound);
{
let mut cache = self.statement_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
cache.insert(trimmed.to_string(), prepared.clone());
}
Ok(prepared)
}
pub fn execute(&self, prepared: &PreparedStatement, params: Vec<(&str, Value)>) -> Result<QueryResult, String> {
let mut param_map = HashMap::new();
let num_expected = prepared.parameters.len();
for (name, value) in ¶ms {
param_map.insert(name.to_string(), value.clone());
}
for p in &prepared.parameters {
if !param_map.contains_key(p) {
return Err(format!("Missing parameter: ${}", p));
}
}
if params.len() > num_expected {
return Err(format!("Expected {} parameter(s), got {}", num_expected, params.len()));
}
if self.database.config.read_only && Connection::is_write_statement(&prepared.bound_statement) {
return Err("Database is in read-only mode; write statements are not allowed".into());
}
if let Some(result) = self.handle_ddl(&prepared.bound_statement)? {
self.database.persist_catalog()?;
self.maybe_auto_checkpoint()?;
return Ok(result);
}
let substituted =
crate::connection::substitute::substitute_params_in_statement(&prepared.bound_statement, ¶m_map)?;
let planner = QueryPlanner::new();
let logical_plan = planner.plan(substituted).map_err(|e| format!("Plan error: {e}"))?;
if logical_plan.is_empty() {
return Ok(QueryResult::success_message("Query executed (no result)".into()));
}
let optimizer = Optimizer::with_stats(self.database.stats_store.clone());
let optimized_plan = optimizer.optimize(logical_plan);
let ts = self.database.transaction_manager.current_commit_ts();
let history = self.database.transaction_manager.commit_history_snapshot();
let processor = self.create_processor().with_snapshot(Some(ts), history);
let chunks = processor
.execute(&optimized_plan)
.map_err(|e| format!("Execute error: {e}"))?;
if Connection::is_write_statement(&prepared.bound_statement) {
self.database
.storage_manager
.persist_all_tables()
.map_err(|e| format!("Failed to persist tables: {e}"))?;
}
self.maybe_auto_checkpoint()?;
Ok(QueryResult::new(chunks))
}
pub(crate) fn create_processor(&self) -> QueryProcessor {
let seq_fn = super::utils::make_sequence_callback(self.database.catalog.clone());
let db = self.database.clone();
let db_sddl = db.clone();
let schema_ddl_fn: akar_processor::processor::SchemaDdlFn = Arc::new(
move |op: akar_processor::processor::SchemaDdlOp| -> Result<String, ProcessorError> {
match op {
akar_processor::processor::SchemaDdlOp::CreateSequence {
name,
if_not_exists,
start_value,
increment,
min_value,
max_value,
cycle,
} => {
let mut catalog = db_sddl.catalog.lock().map_err(|e| format!("Catalog lock: {e}"))?;
match catalog.create_sequence(name.clone(), start_value, increment, min_value, max_value, cycle)
{
akar_catalog::CatalogResult::Created { .. } => Ok(format!("Sequence '{}' created", name)),
akar_catalog::CatalogResult::AlreadyExists => {
if if_not_exists {
Ok(format!("Sequence '{}' already exists", name))
} else {
Err(ProcessorError::Execution(format!("Sequence '{}' already exists", name)))
}
}
other => Err(ProcessorError::Execution(format!(
"Failed to create sequence: {:?}",
other
))),
}
}
akar_processor::processor::SchemaDdlOp::DropSequence { name, if_exists } => {
let mut catalog = db_sddl.catalog.lock().map_err(|e| format!("Catalog lock: {e}"))?;
match catalog.drop_sequence(&name) {
akar_catalog::CatalogResult::Dropped { .. } => Ok(format!("Sequence '{}' dropped", name)),
akar_catalog::CatalogResult::NotFound => {
if if_exists {
Ok(format!("Sequence '{}' not found", name))
} else {
Err(ProcessorError::Execution(format!("Sequence '{}' not found", name)))
}
}
other => Err(ProcessorError::Execution(format!(
"Failed to drop sequence: {:?}",
other
))),
}
}
akar_processor::processor::SchemaDdlOp::ExportDatabase {
file_path,
file_type,
schema_only,
} => {
use std::fs;
use std::path::Path;
let dir = Path::new(&file_path);
fs::create_dir_all(dir)
.map_err(|err| format!("Cannot create export directory '{}': {err}", file_path))?;
let catalog = db_sddl.catalog.lock().map_err(|e| format!("Catalog lock: {e}"))?;
let mut schema = String::new();
for entry in catalog.all_entries() {
match entry {
akar_catalog::CatalogEntry::NodeTable(t) => {
let cols: Vec<String> = t
.columns
.iter()
.map(|c| format!(" {} {:?}", c.name, c.logical_type))
.collect();
schema.push_str(&format!(
"CREATE NODE TABLE {} (\n{}\n);\n\n",
t.name,
cols.join(",\n")
));
}
akar_catalog::CatalogEntry::RelTable(t) => {
let cols: Vec<String> = t
.columns
.iter()
.map(|c| format!(" {} {:?}", c.name, c.logical_type))
.collect();
schema.push_str(&format!(
"CREATE REL TABLE {} (\n{}\n);\n\n",
t.name,
cols.join(",\n")
));
}
_ => {}
}
}
fs::write(dir.join("schema.cypher"), &schema)
.map_err(|err| format!("Cannot write schema.cypher: {err}"))?;
if !schema_only {
let mut copy = String::new();
for entry in catalog.all_entries() {
let name = match entry {
akar_catalog::CatalogEntry::NodeTable(t) => Some(t.name.as_str()),
akar_catalog::CatalogEntry::RelTable(t) => Some(t.name.as_str()),
_ => None,
};
if let Some(table_name) = name {
let ext = if file_type == "parquet" { "parquet" } else { "csv" };
let file_name = format!("{}.{}", table_name, ext);
copy.push_str(&format!("COPY {} FROM '{}';\n", table_name, file_name));
}
}
fs::write(dir.join("copy.cypher"), ©)
.map_err(|err| format!("Cannot write copy.cypher: {err}"))?;
}
Ok(format!("Database exported to '{}'", file_path))
}
akar_processor::processor::SchemaDdlOp::ImportDatabase {
file_path,
query,
index_query,
} => {
let stmts: Vec<&str> = query
.lines()
.chain(index_query.lines())
.filter(|l| !l.trim().is_empty() && !l.trim().starts_with("//"))
.collect();
let count = stmts.len();
Ok(format!("Imported {} statements from '{}'", count, file_path))
}
}
},
);
let db_qf = db.clone();
let query_fn: crate::connection::standalone_call::QueryFn = Arc::new({
let schema_ddl_qf = schema_ddl_fn.clone();
move |query_str: &str| -> Result<crate::query_result::QueryResult, String> {
let stmt = akar_parser::parse(query_str).map_err(|e| format!("Parse error: {e}"))?;
let binder = Binder::new(db_qf.catalog.clone());
let bound = binder.bind(stmt).map_err(|e| format!("Bind error: {e}"))?;
let planner = QueryPlanner::new();
let logical_plan = planner.plan(bound).map_err(|e| format!("Plan error: {e}"))?;
let optimizer = Optimizer::with_stats(db_qf.stats_store.clone());
let optimized_plan = optimizer.optimize(logical_plan);
let processor = QueryProcessor::with_catalog(
db_qf.function_registry.clone(),
db_qf.table_catalog(),
db_qf.vfs.clone(),
)
.with_schema_ddl_fn(schema_ddl_qf.clone())
.with_standalone_call_handler(Arc::new(
crate::connection::standalone_call::DbStandaloneCallHandler::new(db_qf.clone()),
))
.with_snapshot(
Some(db_qf.transaction_manager.current_commit_ts()),
db_qf.transaction_manager.commit_history_snapshot(),
);
let chunks = processor
.execute(&optimized_plan)
.map_err(|e| format!("Execute error: {e}"))?;
let num_rows: usize = chunks.iter().map(|c| c.size).sum();
let num_columns = chunks.first().map(|c| c.num_fields()).unwrap_or(0);
Ok(crate::query_result::QueryResult {
chunks,
num_rows,
num_columns,
success: true,
error_message: None,
message: None,
summary: None,
})
}
});
let subquery_fn: Arc<
dyn Fn(&akar_parser::ast::Query) -> Result<Vec<akar_common::vector::DataChunk>, ProcessorError>
+ Send
+ Sync,
> = Arc::new({
let schema_ddl_sq = schema_ddl_fn.clone();
move |query: &akar_parser::ast::Query| -> Result<Vec<akar_common::vector::DataChunk>, ProcessorError> {
let stmt = akar_parser::ast::Statement::Query(query.clone());
let binder = Binder::new(db.catalog.clone());
let bound = binder.bind(stmt).map_err(|e| format!("Bind error: {e}"))?;
let planner = QueryPlanner::new();
let logical_plan = planner.plan(bound).map_err(|e| format!("Plan error: {e}"))?;
let optimizer = Optimizer::with_stats(db.stats_store.clone());
let optimized_plan = optimizer.optimize(logical_plan);
let catalog_inner = db.catalog.clone();
let seq_fn_inner = super::utils::make_sequence_callback(catalog_inner);
let processor =
QueryProcessor::with_catalog(db.function_registry.clone(), db.table_catalog(), db.vfs.clone())
.with_sequence_fn(seq_fn_inner)
.with_schema_ddl_fn(schema_ddl_sq.clone())
.with_standalone_call_handler(Arc::new(
crate::connection::standalone_call::DbStandaloneCallHandler::new(db.clone()),
))
.with_snapshot(
Some(db.transaction_manager.current_commit_ts()),
db.transaction_manager.commit_history_snapshot(),
);
processor
.execute(&optimized_plan)
.map_err(|e| ProcessorError::Execution(format!("Execute error: {e}")))
}
});
QueryProcessor::with_catalog(
self.database.function_registry.clone(),
self.database.table_catalog(),
self.database.vfs.clone(),
)
.with_sequence_fn(seq_fn)
.with_subquery_fn(subquery_fn)
.with_schema_ddl_fn(schema_ddl_fn)
.with_standalone_call_handler(Arc::new(
crate::connection::standalone_call::DbStandaloneCallHandler::with_query_executor(
self.database.clone(),
Some(query_fn),
),
))
}
pub(crate) fn maybe_auto_checkpoint(&self) -> Result<(), String> {
let threshold = self.database.config.checkpoint_threshold;
if threshold == 0 {
return Ok(()); }
let should_checkpoint = if threshold < 0 {
true
} else {
self.database.storage_manager.wal_size() > threshold as usize
};
if should_checkpoint {
self.database.transaction_manager.schedule_auto_checkpoint();
tracing::debug!("Auto-checkpoint signaled to background worker");
}
Ok(())
}
pub(crate) fn do_sync_checkpoint(&self) -> Result<(), String> {
let tm = &self.database.transaction_manager;
let drain_fn = |timeout: std::time::Duration| -> bool { tm.stop_new_txns_and_wait_until_all_leave(timeout) };
self.database
.storage_manager
.checkpoint_with_drain(Some(&drain_fn))
.map_err(|e| format!("Checkpoint failed: {e}"))?;
tracing::debug!("Sync checkpoint completed");
Ok(())
}
}
fn is_plan_cachable(bound: &BoundStatement) -> bool {
matches!(
bound,
BoundStatement::BoundQuery(_)
| BoundStatement::BoundUnion(_)
| BoundStatement::BoundMerge(_)
| BoundStatement::BoundCreateDml(_)
)
}