llkv_runtime/
runtime_transaction_context.rs

1use std::sync::{Arc, RwLock};
2
3use arrow::record_batch::RecordBatch;
4use llkv_executor::ExecutorRowBatch;
5use llkv_expr::expr::Expr as LlkvExpr;
6use llkv_result::{Error, Result as LlkvResult};
7use llkv_storage::pager::Pager;
8use llkv_table::{CatalogDdl, SingleColumnIndexDescriptor, TableId};
9use llkv_transaction::{TransactionContext, TransactionResult, TransactionSnapshot, TxnId};
10use simd_r_drive_entry_handle::EntryHandle;
11
12use crate::{
13    AlterTablePlan, CreateIndexPlan, CreateTablePlan, DeletePlan, DropIndexPlan, DropTablePlan,
14    InsertPlan, PlanColumnSpec, RenameTablePlan, RuntimeContext, RuntimeStatementResult,
15    SelectExecution, SelectPlan, UpdatePlan,
16};
17use llkv_plan::TruncatePlan;
18
19/// Transaction-scoped façade over [`RuntimeContext`].
20///
21/// This type implements [`TransactionContext`] so the runtime can participate in
22/// the `llkv-transaction` engine without exposing transactional methods on the
23/// shared [`RuntimeContext`] itself. Each instance keeps track of the active
24/// snapshot and delegates operations to the underlying context, applying MVCC
25/// filtering and conversion of statement results into transaction results.
26pub struct RuntimeTransactionContext<P>
27where
28    P: Pager<Blob = EntryHandle> + Send + Sync,
29{
30    ctx: Arc<RuntimeContext<P>>,
31    snapshot: RwLock<TransactionSnapshot>,
32}
33
34impl<P> RuntimeTransactionContext<P>
35where
36    P: Pager<Blob = EntryHandle> + Send + Sync,
37{
38    pub(crate) fn new(ctx: Arc<RuntimeContext<P>>) -> Self {
39        let snapshot = ctx.default_snapshot();
40        Self {
41            ctx,
42            snapshot: RwLock::new(snapshot),
43        }
44    }
45
46    fn update_snapshot(&self, snapshot: TransactionSnapshot) {
47        let mut guard = self.snapshot.write().expect("snapshot lock poisoned");
48        *guard = snapshot;
49    }
50
51    fn current_snapshot(&self) -> TransactionSnapshot {
52        *self.snapshot.read().expect("snapshot lock poisoned")
53    }
54
55    pub(crate) fn context(&self) -> &Arc<RuntimeContext<P>> {
56        &self.ctx
57    }
58
59    pub(crate) fn ctx(&self) -> &RuntimeContext<P> {
60        &self.ctx
61    }
62}
63
64impl<P> CatalogDdl for RuntimeTransactionContext<P>
65where
66    P: Pager<Blob = EntryHandle> + Send + Sync + 'static,
67{
68    type CreateTableOutput = TransactionResult<P>;
69    type DropTableOutput = ();
70    type RenameTableOutput = ();
71    type AlterTableOutput = TransactionResult<P>;
72    type CreateIndexOutput = TransactionResult<P>;
73    type DropIndexOutput = Option<SingleColumnIndexDescriptor>;
74
75    fn create_table(&self, plan: CreateTablePlan) -> LlkvResult<Self::CreateTableOutput> {
76        let ctx = self.context();
77        let result = CatalogDdl::create_table(ctx.as_ref(), plan)?;
78        Ok(convert_statement_result(result))
79    }
80
81    fn drop_table(&self, plan: DropTablePlan) -> LlkvResult<Self::DropTableOutput> {
82        CatalogDdl::drop_table(self.ctx.as_ref(), plan)
83    }
84
85    fn rename_table(&self, plan: RenameTablePlan) -> LlkvResult<Self::RenameTableOutput> {
86        CatalogDdl::rename_table(self.ctx.as_ref(), plan)
87    }
88
89    fn alter_table(&self, plan: AlterTablePlan) -> LlkvResult<Self::AlterTableOutput> {
90        let ctx = self.context();
91        let result = CatalogDdl::alter_table(ctx.as_ref(), plan)?;
92        Ok(convert_statement_result(result))
93    }
94
95    fn create_index(&self, plan: CreateIndexPlan) -> LlkvResult<Self::CreateIndexOutput> {
96        let ctx = self.context();
97        let result = CatalogDdl::create_index(ctx.as_ref(), plan)?;
98        Ok(convert_statement_result(result))
99    }
100
101    fn drop_index(&self, plan: DropIndexPlan) -> LlkvResult<Self::DropIndexOutput> {
102        CatalogDdl::drop_index(self.ctx.as_ref(), plan)
103    }
104}
105
106// Implement TransactionContext to integrate with llkv-transaction.
107impl<P> TransactionContext for RuntimeTransactionContext<P>
108where
109    P: Pager<Blob = EntryHandle> + Send + Sync + 'static,
110{
111    type Pager = P;
112    type Snapshot = llkv_table::catalog::TableCatalogSnapshot;
113
114    fn set_snapshot(&self, snapshot: TransactionSnapshot) {
115        self.update_snapshot(snapshot);
116    }
117
118    fn snapshot(&self) -> TransactionSnapshot {
119        self.current_snapshot()
120    }
121
122    fn table_column_specs(&self, table_name: &str) -> LlkvResult<Vec<PlanColumnSpec>> {
123        let (_, canonical_name) = llkv_table::canonical_table_name(table_name)?;
124        self.context().catalog().table_column_specs(&canonical_name)
125    }
126
127    fn export_table_rows(&self, table_name: &str) -> LlkvResult<ExecutorRowBatch> {
128        RuntimeContext::export_table_rows(self.context(), table_name)
129    }
130
131    fn get_batches_with_row_ids(
132        &self,
133        table_name: &str,
134        filter: Option<LlkvExpr<'static, String>>,
135    ) -> LlkvResult<Vec<RecordBatch>> {
136        self.context()
137            .get_batches_with_row_ids(table_name, filter, self.snapshot())
138    }
139
140    fn execute_select(&self, plan: SelectPlan) -> LlkvResult<SelectExecution<Self::Pager>> {
141        self.context().execute_select(plan, self.snapshot())
142    }
143
144    fn apply_create_table_plan(&self, plan: CreateTablePlan) -> LlkvResult<TransactionResult<P>> {
145        let ctx = self.context();
146        let result = CatalogDdl::create_table(ctx.as_ref(), plan)?;
147        Ok(convert_statement_result(result))
148    }
149
150    fn drop_table(&self, plan: DropTablePlan) -> LlkvResult<()> {
151        CatalogDdl::drop_table(self.ctx.as_ref(), plan)
152    }
153
154    fn insert(&self, plan: InsertPlan) -> LlkvResult<TransactionResult<P>> {
155        tracing::trace!(
156            "[TX_RUNTIME] TransactionContext::insert plan.table='{}', context_pager={:p}",
157            plan.table,
158            &*self.ctx.pager
159        );
160        let snapshot = self.current_snapshot();
161        let result = self.ctx().insert(plan, snapshot)?;
162        Ok(convert_statement_result(result))
163    }
164
165    fn update(&self, plan: UpdatePlan) -> LlkvResult<TransactionResult<P>> {
166        let snapshot = self.current_snapshot();
167        let result = self.ctx().update(plan, snapshot)?;
168        Ok(convert_statement_result(result))
169    }
170
171    fn delete(&self, plan: DeletePlan) -> LlkvResult<TransactionResult<P>> {
172        let snapshot = self.current_snapshot();
173        let result = self.ctx().delete(plan, snapshot)?;
174        Ok(convert_statement_result(result))
175    }
176
177    fn truncate(&self, plan: TruncatePlan) -> LlkvResult<TransactionResult<P>> {
178        let snapshot = self.current_snapshot();
179        let result = self.ctx().truncate(plan, snapshot)?;
180        Ok(convert_statement_result(result))
181    }
182
183    fn create_index(&self, plan: CreateIndexPlan) -> LlkvResult<TransactionResult<P>> {
184        let ctx = self.context();
185        let result = CatalogDdl::create_index(ctx.as_ref(), plan)?;
186        Ok(convert_statement_result(result))
187    }
188
189    fn append_batches_with_row_ids(
190        &self,
191        table_name: &str,
192        batches: Vec<RecordBatch>,
193    ) -> LlkvResult<usize> {
194        RuntimeContext::append_batches_with_row_ids(self.context(), table_name, batches)
195    }
196
197    fn table_names(&self) -> Vec<String> {
198        RuntimeContext::table_names(self.context())
199    }
200
201    fn table_id(&self, table_name: &str) -> LlkvResult<TableId> {
202        let ctx = self.context();
203        if ctx.is_table_marked_dropped(table_name) {
204            return Err(Error::InvalidArgumentError(format!(
205                "table '{}' has been dropped",
206                table_name
207            )));
208        }
209
210        let table = ctx.lookup_table(table_name)?;
211        Ok(table.table.table_id())
212    }
213
214    fn catalog_snapshot(&self) -> Self::Snapshot {
215        let ctx = self.context();
216        ctx.catalog.snapshot()
217    }
218
219    fn validate_commit_constraints(&self, txn_id: TxnId) -> LlkvResult<()> {
220        self.ctx.validate_primary_keys_for_commit(txn_id)
221    }
222
223    fn clear_transaction_state(&self, txn_id: TxnId) {
224        self.ctx.clear_transaction_state(txn_id);
225    }
226}
227
228fn convert_statement_result<P>(result: RuntimeStatementResult<P>) -> TransactionResult<P>
229where
230    P: Pager<Blob = EntryHandle> + Send + Sync + 'static,
231{
232    use llkv_transaction::TransactionResult as TxResult;
233    match result {
234        RuntimeStatementResult::CreateTable { table_name } => TxResult::CreateTable { table_name },
235        RuntimeStatementResult::CreateIndex {
236            table_name,
237            index_name,
238        } => TxResult::CreateIndex {
239            table_name,
240            index_name,
241        },
242        RuntimeStatementResult::Insert { rows_inserted, .. } => TxResult::Insert { rows_inserted },
243        RuntimeStatementResult::Update { rows_updated, .. } => TxResult::Update {
244            rows_matched: rows_updated,
245            rows_updated,
246        },
247        RuntimeStatementResult::Delete { rows_deleted, .. } => TxResult::Delete { rows_deleted },
248        RuntimeStatementResult::Transaction { kind } => TxResult::Transaction { kind },
249        RuntimeStatementResult::NoOp => TxResult::NoOp,
250        _ => panic!("unsupported StatementResult conversion"),
251    }
252}
llkv_runtime/runtime_transaction_context.rs

llkv_runtime/
runtime_transaction_context.rs
