use std::collections::BTreeMap;
use std::sync::Arc;
use anyhow::Result;
use lora_executor::{classify_stream, compiled_result_columns, LoraValue, StreamShape};
use lora_store::InMemoryGraph;
use crate::database::Database;
use crate::error::LoraError;
use crate::stream::{AutoCommitGuard, LiveCursor, QueryStream};
use crate::transaction::{LiveStoreGuard, Transaction, TransactionMode, WriteLease};
impl Database<InMemoryGraph> {
pub fn begin_transaction(&self, mode: TransactionMode) -> Result<Transaction<'_>, LoraError> {
let live = match mode {
TransactionMode::ReadOnly => LiveStoreGuard::Read(self.store.load_full()),
TransactionMode::ReadWrite => {
let lock = self
.writer
.lock()
.unwrap_or_else(|poisoned| poisoned.into_inner());
let snapshot = self.store.load_full();
LiveStoreGuard::Write(WriteLease {
_writer_lock: lock,
store: self.store.clone(),
snapshot,
})
}
};
Ok(Transaction::new(
live,
self.wal.clone(),
self.snapshots.clone(),
mode,
))
}
pub fn stream(&self, query: &str) -> Result<QueryStream<'_>, LoraError> {
self.stream_with_params(query, BTreeMap::new())
}
pub fn stream_with_params(
&self,
query: &str,
params: BTreeMap<String, LoraValue>,
) -> Result<QueryStream<'_>, LoraError> {
let store_guard = self.read_store();
let compiled_arc = self.compile_query_cached(query, &*store_guard)?;
let columns = compiled_result_columns(&compiled_arc);
let shape = classify_stream(&compiled_arc);
drop(store_guard);
match shape {
StreamShape::ReadOnly => {
let compiled = (*compiled_arc).clone();
let live = LiveCursor::open(self.store.clone(), compiled, params)?;
Ok(QueryStream::live(live, columns))
}
StreamShape::Mutating => {
let mut tx = self.begin_transaction(TransactionMode::ReadWrite)?;
let cursor =
match tx.open_streaming_compiled_autocommit(compiled_arc.clone(), params) {
Ok(c) => c,
Err(err) => {
return Err(err.into());
}
};
let guard = AutoCommitGuard {
tx: Some(tx),
finalized: false,
};
Ok(QueryStream::auto_commit(cursor, columns, guard))
}
}
}
pub unsafe fn stream_with_params_owned(
self: &Arc<Self>,
query: &str,
params: BTreeMap<String, LoraValue>,
) -> Result<QueryStream<'static>, LoraError> {
let stream = self.stream_with_params(query, params)?;
Ok(std::mem::transmute::<QueryStream<'_>, QueryStream<'static>>(stream))
}
}