linkrs_query/executor/streaming/operators/
txn_operator.rs

1use std::sync::Arc;
2
3use crate::executor::streaming::chunk::DataChunk;
4use crate::executor::streaming::executor::StreamingExecutor;
5use crate::executor::streaming::operators::source_operator::OperatorConfig;
6use crate::executor::streaming::runtime::ExecutionRuntime;
7use crate::executor::streaming::slot::SlotLayout;
8use crate::executor::streaming::transaction_scope::{
9    SessionTransactionController, TransactionCommandResult,
10};
11use linkrs_core::error::QueryError;
12use linkrs_core::Value;
13
14/// Transaction command operator kind.
15///
16/// Validates state transitions through the [`SessionTransactionController`]
17/// and produces a structured result chunk.  The actual TransactionManager
18/// operations (begin/commit/rollback/savepoint) are performed by the API
19/// layer before this operator runs.
20#[derive(Debug)]
21pub enum TxnOperatorKind {
22    BeginTransaction { emitted: bool },
23    Commit { emitted: bool },
24    Rollback { emitted: bool },
25    RollbackToSavepoint { name: String, emitted: bool },
26    Savepoint { name: String, emitted: bool },
27    ReleaseSavepoint { name: String, emitted: bool },
28}
29
30/// Transaction command operator.
31///
32/// Wraps [`TxnOperatorKind`] with the runtime context injected at `open()`.
33/// Lifecycle state is owned exclusively by the executor; operators never
34/// write it.
35#[derive(Debug)]
36pub struct TxnOperator {
37    pub kind: TxnOperatorKind,
38    pub runtime: Option<Arc<ExecutionRuntime>>,
39    pub output_layout: Arc<SlotLayout>,
40    pub config: OperatorConfig,
41}
42
43impl TxnOperator {
44    pub fn new(kind: TxnOperatorKind, output_layout: Arc<SlotLayout>) -> Self {
45        Self {
46            kind,
47            runtime: None,
48            output_layout,
49            config: OperatorConfig::default(),
50        }
51    }
52
53    pub fn from_spec(spec: &super::spec::TxnSpec, output_layout: Arc<SlotLayout>) -> Self {
54        let kind = match spec {
55            super::spec::TxnSpec::BeginTransaction => {
56                TxnOperatorKind::BeginTransaction { emitted: false }
57            }
58            super::spec::TxnSpec::Commit => TxnOperatorKind::Commit { emitted: false },
59            super::spec::TxnSpec::Rollback => TxnOperatorKind::Rollback { emitted: false },
60            super::spec::TxnSpec::RollbackToSavepoint { name } => {
61                TxnOperatorKind::RollbackToSavepoint {
62                    name: name.clone(),
63                    emitted: false,
64                }
65            }
66            super::spec::TxnSpec::Savepoint { name } => TxnOperatorKind::Savepoint {
67                name: name.clone(),
68                emitted: false,
69            },
70            super::spec::TxnSpec::ReleaseSavepoint { name } => TxnOperatorKind::ReleaseSavepoint {
71                name: name.clone(),
72                emitted: false,
73            },
74        };
75        Self::new(kind, output_layout)
76    }
77
78    /// Inject the runtime and execution config (called once by the executor
79    /// before this operator produces any data).
80    pub fn inject_context(
81        &mut self,
82        runtime: Option<&Arc<ExecutionRuntime>>,
83        config: OperatorConfig,
84    ) {
85        if let Some(rt) = runtime {
86            self.runtime = Some(rt.clone());
87        }
88        self.config = config;
89    }
90
91    fn controller(
92        runtime: &Option<Arc<ExecutionRuntime>>,
93    ) -> Result<Arc<SessionTransactionController>, QueryError> {
94        runtime
95            .as_ref()
96            .and_then(|rt| rt.session_controller())
97            .ok_or_else(|| {
98                QueryError::execution(
99                    "Transaction controller not available in execution runtime".to_string(),
100                )
101            })
102    }
103
104    pub fn open(&mut self, input: &mut StreamingExecutor) -> Result<(), QueryError> {
105        input.open()
106    }
107
108    pub fn next(
109        &mut self,
110        _input: &mut StreamingExecutor,
111    ) -> Result<Option<DataChunk>, QueryError> {
112        let Self {
113            kind,
114            runtime,
115            output_layout,
116            config: _,
117        } = self;
118        match kind {
119            TxnOperatorKind::BeginTransaction { emitted } => {
120                if *emitted {
121                    return Ok(None);
122                }
123                *emitted = true;
124                let ctrl = Self::controller(runtime)?;
125                // The API layer should have already created the transaction.
126                // We validate that begin_tracking succeeds (no state conflict).
127                let result = ctrl.current_scope();
128                match result {
129                    crate::executor::streaming::TransactionScope::ExplicitBorrowed {
130                        transaction_id, ..
131                    } => Ok(Some(
132                        TransactionCommandResult::begin(transaction_id)
133                            .into_data_chunk(Arc::clone(output_layout)),
134                    )),
135                    _ => Err(QueryError::execution(
136                        "BEGIN: no active transaction found; API layer must call begin before executing this plan".to_string(),
137                    )),
138                }
139            }
140            TxnOperatorKind::Commit { emitted } => {
141                if *emitted {
142                    return Ok(None);
143                }
144                *emitted = true;
145                let ctrl = Self::controller(runtime)?;
146                let txn_id = ctrl.begin_commit()?;
147                ctrl.commit_finalize();
148                Ok(Some(
149                    TransactionCommandResult::commit(txn_id)
150                        .into_data_chunk(Arc::clone(output_layout)),
151                ))
152            }
153            TxnOperatorKind::Rollback { emitted } => {
154                if *emitted {
155                    return Ok(None);
156                }
157                *emitted = true;
158                let ctrl = Self::controller(runtime)?;
159                let txn_id = ctrl.begin_rollback()?;
160                ctrl.rollback_finalize();
161                Ok(Some(
162                    TransactionCommandResult::rollback(txn_id)
163                        .into_data_chunk(Arc::clone(output_layout)),
164                ))
165            }
166            TxnOperatorKind::RollbackToSavepoint { name, emitted } => {
167                if *emitted {
168                    return Ok(None);
169                }
170                *emitted = true;
171                let ctrl = Self::controller(runtime)?;
172                // ROLLBACK TO SAVEPOINT does not leave the transaction: the
173                // controller must be Active (validated only, no transition).
174                if !ctrl.is_active() {
175                    return Err(QueryError::execution(format!(
176                        "Cannot ROLLBACK TO SAVEPOINT in state {:?}",
177                        ctrl.state()
178                    )));
179                }
180                Ok(Some(
181                    TransactionCommandResult::rollback_to(name.clone())
182                        .into_data_chunk(Arc::clone(output_layout)),
183                ))
184            }
185            TxnOperatorKind::Savepoint { name, emitted } => {
186                if *emitted {
187                    return Ok(None);
188                }
189                *emitted = true;
190                let ctrl = Self::controller(runtime)?;
191                if !ctrl.is_active() {
192                    return Err(QueryError::execution(format!(
193                        "Cannot SAVEPOINT in state {:?}",
194                        ctrl.state()
195                    )));
196                }
197                Ok(Some(
198                    TransactionCommandResult::savepoint(name.clone())
199                        .into_data_chunk(Arc::clone(output_layout)),
200                ))
201            }
202            TxnOperatorKind::ReleaseSavepoint { name, emitted } => {
203                if *emitted {
204                    return Ok(None);
205                }
206                *emitted = true;
207                let ctrl = Self::controller(runtime)?;
208                if !ctrl.is_active() {
209                    return Err(QueryError::execution(format!(
210                        "Cannot RELEASE SAVEPOINT in state {:?}",
211                        ctrl.state()
212                    )));
213                }
214                Ok(Some(
215                    TransactionCommandResult::release_savepoint(name.clone())
216                        .into_data_chunk(Arc::clone(output_layout)),
217                ))
218            }
219        }
220    }
221
222    pub fn stop(&mut self) -> Result<(), QueryError> {
223        Ok(())
224    }
225
226    pub fn close(&mut self) -> Result<(), QueryError> {
227        Ok(())
228    }
229}
230
231impl TransactionCommandResult {
232    fn into_data_chunk(self, output_layout: Arc<SlotLayout>) -> DataChunk {
233        let message = Value::string(self.message);
234        let command = Value::string(self.command);
235        DataChunk::new_with_layout(vec![vec![command, message]], output_layout)
236    }
237}
238
239#[cfg(test)]
240mod tests {
241    use super::*;
242    use crate::executor::streaming::operators::base::OperatorBase;
243    use crate::executor::streaming::operators::source_operator::SourceOperator;
244    use crate::executor::streaming::operators::source_operator::SourceOperatorKind;
245    use crate::executor::streaming::operators::spec::TxnSpec;
246    use crate::executor::streaming::runtime::ExecutionRuntime;
247
248    fn input() -> StreamingExecutor {
249        StreamingExecutor::Source(
250            OperatorBase::new(0),
251            SourceOperator::new(
252                SourceOperatorKind::ScanVertices {
253                    buffer: Vec::new(),
254                    current_index: 0,
255                    col_names: Vec::new(),
256                },
257                Arc::new(SlotLayout::new(Vec::new())),
258            ),
259        )
260    }
261
262    #[test]
263    fn transaction_command_requires_controller() {
264        let input = input();
265        let operator = TxnOperator::from_spec(
266            &TxnSpec::BeginTransaction,
267            Arc::new(SlotLayout::new(Vec::new())),
268        );
269        let mut executor = StreamingExecutor::Txn(OperatorBase::new(1), Box::new(input), operator);
270        executor.open().expect("should open");
271        let result = executor.advance();
272        assert!(result.is_err());
273        assert!(result
274            .unwrap_err()
275            .to_string()
276            .contains("controller not available"));
277        executor.close().expect("should close");
278    }
279
280    #[test]
281    fn transaction_command_emits_once() {
282        let input = input();
283        let operator = TxnOperator::from_spec(
284            &TxnSpec::BeginTransaction,
285            Arc::new(SlotLayout::new(Vec::new())),
286        );
287        let mut executor = StreamingExecutor::Txn(OperatorBase::new(1), Box::new(input), operator);
288
289        let ctrl = Arc::new(SessionTransactionController::new());
290        let rt = ExecutionRuntime::default_budget();
291        rt.set_session_controller(ctrl.clone());
292        executor.set_runtime(Some(Arc::new(rt)));
293
294        executor.open().expect("transaction command should open");
295
296        // First advance: without a pre-registered transaction, the BEGIN
297        // will fail because current_scope returns None.
298        let result = executor.advance();
299        assert!(result.is_err(), "expected error (no active transaction)");
300
301        // emitted is now true — second call should return None
302        let second = executor.advance().expect("second advance should not fail");
303        assert!(second.is_none());
304        executor.close().expect("transaction command should close");
305    }
306
307    #[test]
308    fn test_commit_without_active_transaction() {
309        let input = input();
310        let operator =
311            TxnOperator::from_spec(&TxnSpec::Commit, Arc::new(SlotLayout::new(Vec::new())));
312        let mut executor = StreamingExecutor::Txn(OperatorBase::new(1), Box::new(input), operator);
313
314        let ctrl = Arc::new(SessionTransactionController::new());
315        let rt = ExecutionRuntime::default_budget();
316        rt.set_session_controller(ctrl.clone());
317        executor.set_runtime(Some(Arc::new(rt)));
318
319        executor.open().expect("should open");
320        let result = executor.advance();
321        assert!(result.is_err());
322        assert!(result.unwrap_err().to_string().contains("Cannot COMMIT"));
323        executor.close().expect("should close");
324    }
325}
linkrs_query/executor/streaming/operators/txn_operator.rs

linkrs_query/executor/streaming/operators/
txn_operator.rs
