1use super::Connection;
2use super::plan_cache::{CachedPlan, normalize_query};
3use crate::database::Database;
4use crate::prepared_statement::PreparedStatement;
5use crate::query_result::QueryResult;
6use akar_binder::Binder;
7use akar_binder::bound_statement::BoundStatement;
8use akar_common::error::ProcessorError;
9use akar_common::types::Value;
10use akar_optimizer::Optimizer;
11use akar_parser::parse;
12use akar_planner::QueryPlanner;
13use akar_planner::logical_operator::LogicalOperator;
14use akar_processor::QueryProcessor;
15use akar_processor::processor::{SchemaDdlFn, SchemaDdlOp, SequenceFn, StandaloneCallHandler, SubqueryFn};
16use std::collections::HashMap;
17use std::sync::Arc;
18
19impl Connection {
20 pub fn query(&self, query_str: &str) -> Result<QueryResult, String> {
22 let trimmed = query_str.trim();
23
24 if trimmed.is_empty() {
26 return Ok(QueryResult::new(Vec::new()));
27 }
28
29 if let Some(value) = trimmed
31 .strip_prefix("SET")
32 .and_then(|s| s.trim().strip_prefix("spill_threshold"))
33 .and_then(|s| s.trim().strip_prefix("="))
34 .map(|s| s.trim())
35 {
36 let bytes: u64 = value.parse().map_err(|_| {
37 format!("Invalid spill_threshold value '{value}'. Expected a non-negative integer (bytes).")
38 })?;
39 self.database.set_spill_threshold(bytes);
40 return Ok(QueryResult::success_message(format!(
41 "spill_threshold set to {bytes} bytes"
42 )));
43 }
44
45 if let Some(value) = trimmed
46 .strip_prefix("SET")
47 .and_then(|s| s.trim().strip_prefix("concurrent_writes"))
48 .and_then(|s| s.trim().strip_prefix("="))
49 .map(|s| s.trim())
50 {
51 let enabled = match value.to_lowercase().as_str() {
52 "true" | "1" | "yes" => true,
53 "false" | "0" | "no" => false,
54 _ => return Err("Invalid value for concurrent_writes. Use true or false.".into()),
55 };
56 self.database.transaction_manager.set_concurrent_writes(enabled);
57 return Ok(QueryResult::success_message(format!(
58 "concurrent_writes set to {enabled}"
59 )));
60 }
61
62 let normalized = normalize_query(trimmed);
64
65 let catalog_version = self
68 .database
69 .catalog
70 .lock()
71 .map_err(|e| format!("Catalog lock error: {e}"))?
72 .version();
73
74 {
75 let mut cache = self.plan_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
76 if let Some(cached) = cache.get(&normalized).filter(|c| c.catalog_version == catalog_version) {
77 let bound = cached.bound.clone();
80 let plan = cached.plan.clone();
81 drop(cache);
82 return self.execute_with_plan(&bound, Some(&plan));
83 }
84 }
85
86 let statement = parse(trimmed).map_err(|e| format!("Parse error: {e}"))?;
89
90 let binder = Binder::new(self.database.catalog.clone());
92 let bound = binder.bind(statement).map_err(|e| format!("Bind error: {e}"))?;
93
94 let plan_opt: Option<Arc<Vec<LogicalOperator>>> = if is_plan_cachable(&bound) {
98 let plan = Arc::new(self.build_optimized_plan(&bound)?);
99 let mut cache = self.plan_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
100 cache.insert(
101 normalized,
102 CachedPlan {
103 bound: Arc::new(bound.clone()),
104 plan: Arc::clone(&plan),
105 catalog_version,
106 },
107 );
108 Some(plan)
109 } else {
110 None
111 };
112
113 self.execute_with_plan(&bound, plan_opt.as_ref())
114 }
115
116 fn execute_with_plan(
120 &self,
121 bound: &BoundStatement,
122 plan: Option<&Arc<Vec<LogicalOperator>>>,
123 ) -> Result<QueryResult, String> {
124 if self.database.config.read_only && Connection::is_write_statement(bound) {
126 return Err("Database is in read-only mode; write statements are not allowed".into());
127 }
128
129 let is_write = Connection::is_write_statement(bound);
134 let mut txn_opt: Option<akar_transaction::Transaction> =
135 if is_write { Some(self.begin_write_txn()?) } else { None };
136
137 let query_result = self.execute_query_inner(bound, txn_opt.as_mut(), plan);
139
140 match (is_write, &query_result) {
142 (true, Ok(_)) => {
143 if let Some(ref mut txn) = txn_opt {
144 self.commit_write_txn(txn)?;
145 }
146 }
147 (true, Err(e)) => {
148 if let Some(ref mut txn) = txn_opt {
149 match self.rollback_write_txn(txn) {
150 Ok(_records) => {
151 tracing::warn!("Transaction rolled back due to error: {e}");
152 }
153 Err(rollback_err) => {
154 tracing::error!("Transaction rollback ALSO failed: {rollback_err} (original error: {e})");
155 }
156 }
157 }
158 }
159 _ => {}
160 }
161
162 if query_result.is_ok() && Connection::is_write_statement(bound) {
165 let written = Connection::extract_write_tables(bound);
166 if !written.is_empty() {
167 self.database.refresh_vector_indexes(&written);
168 }
169 }
170
171 query_result
172 }
173
174 fn build_optimized_plan(&self, bound: &BoundStatement) -> Result<Vec<LogicalOperator>, String> {
176 let planner = QueryPlanner::new();
177 let logical_plan = planner.plan(bound.clone()).map_err(|e| format!("Plan error: {e}"))?;
178 let optimizer = Optimizer::with_stats_and_fts(
179 self.database.stats_store.clone(),
180 super::fts_estimate::build(self.database.table_catalog()),
181 );
182 Ok(optimizer.optimize(logical_plan))
183 }
184
185 pub(crate) fn execute_query_inner(
191 &self,
192 bound: &BoundStatement,
193 mut txn_opt: Option<&mut akar_transaction::Transaction>,
194 cached_plan: Option<&Arc<Vec<LogicalOperator>>>,
195 ) -> Result<QueryResult, String> {
196 if let Some(result) = self.handle_ddl(bound, txn_opt.as_deref_mut())? {
198 self.database.persist_catalog()?;
201 self.maybe_auto_checkpoint()?;
202 return Ok(result);
203 }
204
205 if let Some(ref txn) = txn_opt {
209 if !self.database.transaction_manager.allow_concurrent_writes() {
210 let write_tables = Connection::extract_write_tables(bound);
211 for tid in write_tables {
212 self.database.transaction_manager.lock_table(txn.transaction_id, tid)?;
213 }
214 }
215 }
216
217 let optimized_plan: Arc<Vec<LogicalOperator>> = match cached_plan {
221 Some(plan) => plan.clone(),
222 None => Arc::new(self.build_optimized_plan(bound)?),
223 };
224
225 if optimized_plan.is_empty() {
226 return Ok(QueryResult::success_message("Query executed (no result)".into()));
227 }
228
229 let (snapshot_ts, commit_history) = if let Some(ref txn) = txn_opt {
233 (
234 txn.snapshot_ts,
235 self.database.transaction_manager.commit_history_snapshot(),
236 )
237 } else {
238 let ts = self.database.transaction_manager.current_commit_ts();
240 let history = self.database.transaction_manager.commit_history_snapshot();
241 (Some(ts), history)
242 };
243
244 let processor = self
246 .create_processor()
247 .with_snapshot(snapshot_ts, commit_history)
248 .with_txn_id(txn_opt.as_ref().map(|t| t.transaction_id));
249 let chunks = processor
250 .execute(&optimized_plan)
251 .map_err(|e| format!("Execute error: {e}"))?;
252
253 if let Some(ref txn) = txn_opt {
255 let written_rows = processor.take_written_rows();
256 let tm = &self.database.transaction_manager;
257 for (table_id, row_id) in written_rows {
258 tm.record_write(txn.transaction_id, table_id, row_id);
259 }
260 }
261
262 if let Some(ref mut txn) = txn_opt {
265 let undo = processor.take_undo_records();
266 txn.undo_records.extend(undo);
267 let wal_records = processor.take_wal_records();
269 self.append_local_wal(txn.transaction_id, wal_records);
270 }
271
272 self.maybe_auto_checkpoint()?;
274
275 Ok(QueryResult::new(chunks))
276 }
277
278 pub fn prepare(&self, query_str: &str) -> Result<PreparedStatement, String> {
284 let trimmed = query_str.trim();
285
286 {
288 let cache = self.statement_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
289 if let Some(cached) = cache.get(trimmed) {
290 return Ok(cached.clone());
291 }
292 }
293
294 let statement = parse(trimmed).map_err(|e| format!("Parse error: {e}"))?;
296
297 let binder = Binder::new(self.database.catalog.clone());
299 let bound = binder.bind(statement).map_err(|e| format!("Bind error: {e}"))?;
300
301 let prepared = PreparedStatement::new(trimmed.to_string(), bound);
302
303 {
305 let mut cache = self.statement_cache.lock().map_err(|e| format!("Lock poisoned: {e}"))?;
306 cache.insert(trimmed.to_string(), prepared.clone());
307 }
308
309 Ok(prepared)
310 }
311
312 pub fn execute(&self, prepared: &PreparedStatement, params: Vec<(&str, Value)>) -> Result<QueryResult, String> {
318 let mut param_map = HashMap::new();
320 let num_expected = prepared.parameters.len();
321 for (name, value) in ¶ms {
322 param_map.insert(name.to_string(), value.clone());
323 }
324
325 for p in &prepared.parameters {
327 if !param_map.contains_key(p) {
328 return Err(format!("Missing parameter: ${}", p));
329 }
330 }
331
332 if params.len() > num_expected {
334 return Err(format!("Expected {} parameter(s), got {}", num_expected, params.len()));
335 }
336
337 if self.database.config.read_only && Connection::is_write_statement(&prepared.bound_statement) {
339 return Err("Database is in read-only mode; write statements are not allowed".into());
340 }
341
342 let substituted =
346 crate::connection::substitute::substitute_params_in_statement(&prepared.bound_statement, ¶m_map)?;
347
348 let is_write = Connection::is_write_statement(&prepared.bound_statement);
356 let mut txn_opt: Option<akar_transaction::Transaction> =
357 if is_write { Some(self.begin_write_txn()?) } else { None };
358
359 match self.handle_ddl(&substituted, txn_opt.as_mut()) {
362 Ok(Some(result)) => {
363 self.database.persist_catalog()?;
364 self.maybe_auto_checkpoint()?;
365 if is_write {
366 if let Some(ref mut txn) = txn_opt {
367 self.commit_write_txn(txn)?;
368 }
369 }
370 if is_write {
371 let written = Connection::extract_write_tables(&substituted);
372 if !written.is_empty() {
373 self.database.refresh_vector_indexes(&written);
374 }
375 }
376 return Ok(result);
377 }
378 Ok(None) => {}
379 Err(e) => {
380 if is_write {
381 if let Some(ref mut txn) = txn_opt {
382 match self.rollback_write_txn(txn) {
383 Ok(_) => tracing::warn!("Prepared DDL/DML rolled back due to error: {e}"),
384 Err(rollback_err) => {
385 tracing::error!("Prepared rollback ALSO failed: {rollback_err} (original: {e})");
386 }
387 }
388 }
389 }
390 return Err(e);
391 }
392 }
393
394 let planner = QueryPlanner::new();
396 let logical_plan = match planner.plan(substituted) {
397 Ok(p) => p,
398 Err(e) => {
399 if is_write {
400 if let Some(ref mut txn) = txn_opt {
401 let _ = self.rollback_write_txn(txn);
402 }
403 }
404 return Err(format!("Plan error: {e}"));
405 }
406 };
407
408 if logical_plan.is_empty() {
409 if is_write {
410 if let Some(ref mut txn) = txn_opt {
411 self.commit_write_txn(txn)?;
412 }
413 }
414 return Ok(QueryResult::success_message("Query executed (no result)".into()));
415 }
416
417 let optimizer = Optimizer::with_stats_and_fts(
419 self.database.stats_store.clone(),
420 super::fts_estimate::build(self.database.table_catalog()),
421 );
422 let optimized_plan = optimizer.optimize(logical_plan);
423
424 let (snapshot_ts, history) = if let Some(ref txn) = txn_opt {
426 (
427 txn.snapshot_ts,
428 self.database.transaction_manager.commit_history_snapshot(),
429 )
430 } else {
431 let ts = self.database.transaction_manager.current_commit_ts();
432 (Some(ts), self.database.transaction_manager.commit_history_snapshot())
433 };
434
435 let processor = self
437 .create_processor()
438 .with_snapshot(snapshot_ts, history)
439 .with_txn_id(txn_opt.as_ref().map(|t| t.transaction_id));
440 let chunks = match processor.execute(&optimized_plan) {
441 Ok(c) => c,
442 Err(e) => {
443 if is_write {
444 if let Some(ref mut txn) = txn_opt {
445 match self.rollback_write_txn(txn) {
446 Ok(_) => tracing::warn!("Prepared write rolled back due to error: {e}"),
447 Err(rollback_err) => {
448 tracing::error!("Prepared rollback ALSO failed: {rollback_err} (original: {e})");
449 }
450 }
451 }
452 }
453 return Err(format!("Execute error: {e}"));
454 }
455 };
456
457 if let Some(ref txn) = txn_opt {
459 let written_rows = processor.take_written_rows();
460 let tm = &self.database.transaction_manager;
461 for (table_id, row_id) in written_rows {
462 tm.record_write(txn.transaction_id, table_id, row_id);
463 }
464 }
465
466 if let Some(ref mut txn) = txn_opt {
468 let undo = processor.take_undo_records();
469 txn.undo_records.extend(undo);
470 let wal_records = processor.take_wal_records();
472 self.append_local_wal(txn.transaction_id, wal_records);
473 }
474
475 if is_write {
479 if let Some(ref mut txn) = txn_opt {
480 self.commit_write_txn(txn)?;
481 }
482 }
483
484 if is_write {
488 let written = Connection::extract_write_tables(&prepared.bound_statement);
489 if !written.is_empty() {
490 self.database.refresh_vector_indexes(&written);
491 }
492 }
493
494 self.maybe_auto_checkpoint()?;
496
497 Ok(QueryResult::new(chunks))
498 }
499
500 pub(crate) fn maybe_auto_checkpoint(&self) -> Result<(), String> {
505 if !self.database.config.auto_checkpoint {
506 return Ok(()); }
508 let threshold = self.database.config.checkpoint_threshold;
509 if threshold == 0 {
510 return Ok(()); }
512
513 let should_checkpoint = if threshold < 0 {
514 true
515 } else {
516 self.database.storage_manager.wal_size() > threshold as usize
517 };
518
519 if should_checkpoint {
520 self.database.transaction_manager.schedule_auto_checkpoint();
522 tracing::debug!("Auto-checkpoint signaled to background worker");
523 }
524
525 Ok(())
526 }
527
528 pub(crate) fn do_sync_checkpoint(&self) -> Result<(), String> {
530 let tm = &self.database.transaction_manager;
531 let drain_fn = |timeout: std::time::Duration| -> bool { tm.stop_new_txns_and_wait_until_all_leave(timeout) };
532 self.database
533 .storage_manager
534 .checkpoint_with_drain(Some(&drain_fn))
535 .map_err(|e| format!("Checkpoint failed: {e}"))?;
536 tracing::debug!("Sync checkpoint completed");
537 Ok(())
538 }
539
540 pub(crate) fn create_processor(&self) -> QueryProcessor {
547 let handlers = self
548 .processor_handlers
549 .get_or_init(|| Arc::new(build_processor_handlers(&self.database)));
550
551 QueryProcessor::with_catalog(
552 self.database.function_registry.clone(),
553 self.database.table_catalog(),
554 self.database.vfs.clone(),
555 )
556 .with_sequence_fn(handlers.sequence_fn.clone())
557 .with_subquery_fn(handlers.subquery_fn.clone())
558 .with_schema_ddl_fn(handlers.schema_ddl_fn.clone())
559 .with_standalone_call_handler(handlers.standalone_call_handler.clone())
560 }
561}
562
563pub(crate) struct ProcessorHandlers {
568 pub sequence_fn: SequenceFn,
569 pub schema_ddl_fn: SchemaDdlFn,
570 pub subquery_fn: SubqueryFn,
571 pub standalone_call_handler: Arc<dyn StandaloneCallHandler>,
572}
573
574fn build_processor_handlers(db: &Arc<Database>) -> ProcessorHandlers {
576 let seq_fn = super::utils::make_sequence_callback(db.catalog.clone());
577
578 let db_sddl = db.clone();
580 let schema_ddl_fn: SchemaDdlFn = Arc::new(move |op: SchemaDdlOp| -> Result<String, ProcessorError> {
581 match op {
582 SchemaDdlOp::CreateSequence {
583 name,
584 if_not_exists,
585 start_value,
586 increment,
587 min_value,
588 max_value,
589 cycle,
590 } => {
591 let mut catalog = db_sddl.catalog.lock().map_err(|e| format!("Catalog lock: {e}"))?;
592 match catalog.create_sequence(name.clone(), start_value, increment, min_value, max_value, cycle) {
593 akar_catalog::CatalogResult::Created { .. } => Ok(format!("Sequence '{}' created", name)),
594 akar_catalog::CatalogResult::AlreadyExists => {
595 if if_not_exists {
596 Ok(format!("Sequence '{}' already exists", name))
597 } else {
598 Err(ProcessorError::Execution(format!("Sequence '{}' already exists", name)))
599 }
600 }
601 other => Err(ProcessorError::Execution(format!(
602 "Failed to create sequence: {:?}",
603 other
604 ))),
605 }
606 }
607 SchemaDdlOp::DropSequence { name, if_exists } => {
608 let mut catalog = db_sddl.catalog.lock().map_err(|e| format!("Catalog lock: {e}"))?;
609 match catalog.drop_sequence(&name) {
610 akar_catalog::CatalogResult::Dropped { .. } => Ok(format!("Sequence '{}' dropped", name)),
611 akar_catalog::CatalogResult::NotFound => {
612 if if_exists {
613 Ok(format!("Sequence '{}' not found", name))
614 } else {
615 Err(ProcessorError::Execution(format!("Sequence '{}' not found", name)))
616 }
617 }
618 other => Err(ProcessorError::Execution(format!(
619 "Failed to drop sequence: {:?}",
620 other
621 ))),
622 }
623 }
624 SchemaDdlOp::ExportDatabase {
625 file_path,
626 file_type,
627 schema_only,
628 } => {
629 let conn = super::Connection::new(&db_sddl);
634 let bound = akar_binder::bound_statement::BoundExportDatabase {
635 file_path,
636 file_type,
637 schema_only,
638 options: Default::default(),
639 };
640 let result = conn
641 .execute_export_database(&bound)
642 .map_err(ProcessorError::Execution)?;
643 let msg = result
644 .and_then(|r| r.message)
645 .unwrap_or_else(|| format!("Database exported to '{}'", bound.file_path));
646 Ok(msg)
647 }
648 SchemaDdlOp::ImportDatabase {
649 file_path,
650 query,
651 index_query,
652 } => {
653 let conn = super::Connection::new(&db_sddl);
657 let mut executed = 0usize;
658 let mut skipped = 0usize;
659 for stmt in super::copy::split_cypher_statements(&query)
660 .into_iter()
661 .chain(super::copy::split_cypher_statements(&index_query))
662 {
663 match conn.query(&stmt) {
664 Ok(_) => executed += 1,
665 Err(e) => {
666 tracing::warn!("Import statement skipped (may be duplicate): {e}");
667 skipped += 1;
668 }
669 }
670 }
671 Ok(format!(
672 "Imported {executed} statement(s) from '{file_path}' ({skipped} skipped)"
673 ))
674 }
675 }
676 });
677
678 let db_qf = db.clone();
680 let query_fn: crate::connection::standalone_call::QueryFn = Arc::new({
681 let schema_ddl_qf = schema_ddl_fn.clone();
682 move |query_str: &str| -> Result<crate::query_result::QueryResult, String> {
683 let stmt = akar_parser::parse(query_str).map_err(|e| format!("Parse error: {e}"))?;
684 let binder = Binder::new(db_qf.catalog.clone());
685 let bound = binder.bind(stmt).map_err(|e| format!("Bind error: {e}"))?;
686 let planner = QueryPlanner::new();
687 let logical_plan = planner.plan(bound).map_err(|e| format!("Plan error: {e}"))?;
688 let optimizer = Optimizer::with_stats_and_fts(
689 db_qf.stats_store.clone(),
690 super::fts_estimate::build(db_qf.table_catalog()),
691 );
692 let optimized_plan = optimizer.optimize(logical_plan);
693
694 let processor = QueryProcessor::with_catalog(
695 db_qf.function_registry.clone(),
696 db_qf.table_catalog(),
697 db_qf.vfs.clone(),
698 )
699 .with_schema_ddl_fn(schema_ddl_qf.clone())
700 .with_standalone_call_handler(Arc::new(
701 crate::connection::standalone_call::DbStandaloneCallHandler::new(db_qf.clone()),
702 ))
703 .with_snapshot(
704 Some(db_qf.transaction_manager.current_commit_ts()),
705 db_qf.transaction_manager.commit_history_snapshot(),
706 );
707
708 let chunks = processor
709 .execute(&optimized_plan)
710 .map_err(|e| format!("Execute error: {e}"))?;
711
712 let num_rows: usize = chunks.iter().map(|c| c.size).sum();
713 let num_columns = chunks.first().map(|c| c.num_fields()).unwrap_or(0);
714 Ok(crate::query_result::QueryResult {
715 chunks,
716 num_rows,
717 num_columns,
718 success: true,
719 error_message: None,
720 message: None,
721 summary: None,
722 })
723 }
724 });
725
726 let db_sq = db.clone();
727 let subquery_fn: SubqueryFn = Arc::new({
728 let schema_ddl_sq = schema_ddl_fn.clone();
729 move |query: &akar_parser::ast::Query| -> Result<Vec<akar_common::vector::DataChunk>, ProcessorError> {
730 let stmt = akar_parser::ast::Statement::Query(query.clone());
731 let binder = Binder::new(db_sq.catalog.clone());
732 let bound = binder.bind(stmt).map_err(|e| format!("Bind error: {e}"))?;
733 let planner = QueryPlanner::new();
734 let logical_plan = planner.plan(bound).map_err(|e| format!("Plan error: {e}"))?;
735 let optimizer = Optimizer::with_stats_and_fts(
736 db_sq.stats_store.clone(),
737 super::fts_estimate::build(db_sq.table_catalog()),
738 );
739 let optimized_plan = optimizer.optimize(logical_plan);
740
741 let catalog_inner = db_sq.catalog.clone();
742 let seq_fn_inner = super::utils::make_sequence_callback(catalog_inner);
743
744 let processor = QueryProcessor::with_catalog(
745 db_sq.function_registry.clone(),
746 db_sq.table_catalog(),
747 db_sq.vfs.clone(),
748 )
749 .with_sequence_fn(seq_fn_inner)
750 .with_schema_ddl_fn(schema_ddl_sq.clone())
751 .with_standalone_call_handler(Arc::new(
752 crate::connection::standalone_call::DbStandaloneCallHandler::new(db_sq.clone()),
753 ))
754 .with_snapshot(
755 Some(db_sq.transaction_manager.current_commit_ts()),
756 db_sq.transaction_manager.commit_history_snapshot(),
757 );
758
759 processor
760 .execute(&optimized_plan)
761 .map_err(|e| ProcessorError::Execution(format!("Execute error: {e}")))
762 }
763 });
764
765 let standalone_call_handler: Arc<dyn StandaloneCallHandler> = Arc::new(
766 crate::connection::standalone_call::DbStandaloneCallHandler::with_query_executor(
767 db.clone(),
768 Some(query_fn.clone()),
769 ),
770 );
771
772 ProcessorHandlers {
773 sequence_fn: seq_fn,
774 schema_ddl_fn,
775 subquery_fn,
776 standalone_call_handler,
777 }
778}
779
780fn is_plan_cachable(bound: &BoundStatement) -> bool {
788 match bound {
789 BoundStatement::BoundQuery(q) => {
790 !(q.clauses.len() == 1
791 && matches!(
792 q.clauses.first(),
793 Some(akar_binder::bound_statement::BoundClause::BoundForeach(_))
794 ))
795 }
796 _ => false,
797 }
798}