linkrs_query/executor/streaming/
transaction_scope.rs

1//! Transaction scope for query execution instances.
2//!
3//! - Session-level `SessionTransactionController` holds explicit transaction
4//!   handles and manages state transitions.
5//! - `TransactionScope` is the execution view for a single query: it carries
6//!   pure metadata (transaction identity, mode, rollback flag).
7//! - Begin/Commit/Rollback commands trigger real state transitions through
8//!   the controller, not fake text results.
9//! - Client disconnect rolls back auto-commit transactions owned by the
10//!   instance, but does NOT end the session's explicit transaction.
11//!
12//! The `SessionTransactionController` is the single state machine that all
13//! transaction paths (SQL text, plan-based `TxnOperator`, embedded) flow through.
14//! The API layer (`GraphService`) owns the `TransactionManager` reference and
15//! performs the actual begin/commit/rollback; the controller tracks state.
16
17use crate::executor::base::ExecutionResult;
18use linkrs_core::error::QueryError;
19
20pub use linkrs_core::types::TransactionId;
21
22// ── CancelReason ───────────────────────────────────────────────────────────
23
24/// Typed reason for query cancellation.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub enum CancelReason {
27    UserKill,
28    Deadline,
29    ClientDisconnect,
30    MemoryLimit,
31    WorkerFailure,
32    Shutdown,
33    /// A write operation hit a transaction conflict (write-write conflict or
34    /// a rollback-only transaction); remaining pipeline stages of the same
35    /// transaction are cancelled to avoid wasted computation.
36    TransactionConflict,
37    Internal(String),
38}
39
40impl std::fmt::Display for CancelReason {
41    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
42        match self {
43            Self::UserKill => write!(f, "killed by user"),
44            Self::Deadline => write!(f, "deadline exceeded"),
45            Self::ClientDisconnect => write!(f, "client disconnected"),
46            Self::MemoryLimit => write!(f, "memory limit exceeded"),
47            Self::WorkerFailure => write!(f, "worker failure"),
48            Self::Shutdown => write!(f, "system shutdown"),
49            Self::TransactionConflict => write!(f, "transaction conflict"),
50            Self::Internal(msg) => write!(f, "{}", msg),
51        }
52    }
53}
54
55// ── TransactionScope ──────────────────────────────────────────────────────
56
57/// The transaction scope for a single query execution.
58///
59/// Pure metadata — no back-reference to the controller. The caller (session /
60/// query instance) holds both the scope and the controller separately.
61///
62/// - `ExplicitBorrowed`: the query runs within a session-level explicit
63///   transaction.
64/// - `AutoCommitOwned`: the instance owns an auto-commit transaction that is
65///   committed (or rolled back) when execution completes.
66/// - `ReadOnlySnapshot`: a read-only snapshot that does not need a write
67///   transaction handle.
68/// - `CommandScope`: transient scope for BEGIN / COMMIT / ROLLBACK state
69///   transitions.  No real data access happens in this scope.
70/// - `None`: no transaction is required (DDL, admin commands, etc.).
71#[derive(Debug, Clone)]
72pub enum TransactionScope {
73    ExplicitBorrowed {
74        transaction_id: TransactionId,
75        read_write: bool,
76    },
77    AutoCommitOwned {
78        transaction_id: TransactionId,
79        rollback_only: bool,
80    },
81    ReadOnlySnapshot {
82        transaction_id: TransactionId,
83    },
84    CommandScope,
85    None,
86}
87
88impl TransactionScope {
89    pub fn explicit(transaction_id: TransactionId, read_write: bool) -> Self {
90        Self::ExplicitBorrowed {
91            transaction_id,
92            read_write,
93        }
94    }
95
96    pub fn auto_commit(transaction_id: TransactionId) -> Self {
97        Self::AutoCommitOwned {
98            transaction_id,
99            rollback_only: false,
100        }
101    }
102
103    pub fn read_only(transaction_id: TransactionId) -> Self {
104        Self::ReadOnlySnapshot { transaction_id }
105    }
106
107    /// Whether a write operation is allowed in this scope.
108    pub fn allows_write(&self) -> bool {
109        match self {
110            Self::ExplicitBorrowed { read_write, .. } => *read_write,
111            Self::AutoCommitOwned { rollback_only, .. } => !*rollback_only,
112            _ => false,
113        }
114    }
115
116    /// Whether this is an explicit (session-level) transaction.
117    pub fn is_explicit(&self) -> bool {
118        matches!(self, Self::ExplicitBorrowed { .. })
119    }
120
121    /// Whether this is an auto-commit transaction.
122    pub fn is_auto_commit(&self) -> bool {
123        matches!(self, Self::AutoCommitOwned { .. })
124    }
125
126    /// Whether this scope expects at least one result row (command scopes do not).
127    pub fn produces_result(&self) -> bool {
128        !matches!(self, Self::CommandScope)
129    }
130
131    /// Return the transaction ID, if any.
132    pub fn transaction_id(&self) -> Option<TransactionId> {
133        match self {
134            Self::ExplicitBorrowed { transaction_id, .. } => Some(*transaction_id),
135            Self::AutoCommitOwned { transaction_id, .. } => Some(*transaction_id),
136            Self::ReadOnlySnapshot { transaction_id } => Some(*transaction_id),
137            _ => None,
138        }
139    }
140
141    /// Validate that writes are allowed.  Returns an error if this scope
142    /// forbids writes.
143    pub fn ensure_can_write(&self) -> Result<(), QueryError> {
144        if self.allows_write() {
145            Ok(())
146        } else {
147            Err(QueryError::execution(
148                "Current transaction scope does not allow writes".to_string(),
149            ))
150        }
151    }
152
153    /// Mark the transaction as rollback-only (auto-commit variants).
154    pub fn mark_rollback_only(&mut self) {
155        if let Self::AutoCommitOwned {
156            ref mut rollback_only,
157            ..
158        } = self
159        {
160            *rollback_only = true;
161        }
162    }
163
164    /// Whether the auto-commit transaction has been marked rollback-only.
165    pub fn is_rollback_only(&self) -> bool {
166        matches!(
167            self,
168            Self::AutoCommitOwned {
169                rollback_only: true,
170                ..
171            }
172        )
173    }
174}
175
176// ── Transaction state machine ─────────────────────────────────────────────
177
178/// Transaction state for session-level explicit transactions.
179#[derive(Debug, Clone, Copy, PartialEq, Eq)]
180pub enum TransactionState {
181    /// No transaction active.
182    None,
183    /// Transaction is active and accepting operations.
184    Active,
185    /// Transaction is committing (in-progress).
186    Committing,
187    /// Transaction has committed successfully.
188    Committed,
189    /// Transaction is rolling back (in-progress).
190    RollingBack,
191    /// Transaction has rolled back.
192    RolledBack,
193    /// Transaction has failed and can only be rolled back.
194    RollbackOnly,
195}
196
197impl TransactionState {
198    pub fn can_begin(&self) -> bool {
199        matches!(self, Self::None | Self::Committed | Self::RolledBack)
200    }
201
202    pub fn can_execute(&self) -> bool {
203        matches!(self, Self::Active)
204    }
205
206    pub fn can_commit(&self) -> bool {
207        matches!(self, Self::Active)
208    }
209
210    pub fn can_rollback(&self) -> bool {
211        matches!(self, Self::Active | Self::RollbackOnly)
212    }
213}
214
215// ── SessionTransactionController ──────────────────────────────────────────
216
217/// Session-level controller for explicit transactions.
218///
219/// Tracks the currently active explicit transaction for a session with
220/// a proper state machine.  The controller validates state transitions but
221/// does NOT own a [`TransactionManager`] reference — the caller (typically
222/// the API layer) performs the actual begin/commit/rollback operations
223/// on the manager and uses this controller for state tracking and scope
224/// creation.
225///
226/// Thread-safe: uses internal synchronization so it can be shared via `Arc`.
227#[derive(Debug, Default)]
228pub struct SessionTransactionController {
229    state: parking_lot::RwLock<SessionTransactionState>,
230}
231
232/// Internal state for a session's transaction tracking.
233#[derive(Debug, Clone)]
234struct SessionTransactionState {
235    active_transaction: Option<TransactionId>,
236    read_write: bool,
237    auto_commit: bool,
238    state: TransactionState,
239    rollback_only: bool,
240}
241
242impl Default for SessionTransactionState {
243    fn default() -> Self {
244        Self {
245            active_transaction: None,
246            read_write: true,
247            auto_commit: true,
248            state: TransactionState::None,
249            rollback_only: false,
250        }
251    }
252}
253
254impl SessionTransactionController {
255    pub fn new() -> Self {
256        Self {
257            state: parking_lot::RwLock::new(SessionTransactionState::default()),
258        }
259    }
260
261    // ── State queries ──
262
263    /// Current state of the session's explicit transaction.
264    pub fn state(&self) -> TransactionState {
265        self.state.read().state
266    }
267
268    /// Whether an explicit transaction is active and can execute queries.
269    pub fn is_active(&self) -> bool {
270        let s = self.state.read();
271        s.active_transaction.is_some() && s.state == TransactionState::Active
272    }
273
274    /// Whether auto-commit mode is enabled (outside explicit transactions).
275    pub fn auto_commit(&self) -> bool {
276        self.state.read().auto_commit
277    }
278
279    /// Whether the current transaction has been marked rollback-only.
280    pub fn rollback_only(&self) -> bool {
281        self.state.read().rollback_only
282    }
283
284    /// Return the active explicit transaction ID, if any.
285    pub fn current_transaction_id(&self) -> Option<TransactionId> {
286        self.state.read().active_transaction
287    }
288
289    // ── State transitions (tracking only — actual TM ops are caller's job) ──
290
291    /// Begin tracking a new explicit transaction.
292    ///
293    /// The caller should have already created the transaction on the
294    /// `TransactionManager` before calling this method.
295    pub fn begin_tracking(
296        &self,
297        transaction_id: TransactionId,
298        read_write: bool,
299    ) -> Result<(), QueryError> {
300        let mut state = self.state.write();
301        if !state.state.can_begin() {
302            return Err(QueryError::execution(format!(
303                "Cannot BEGIN in state {:?}",
304                state.state
305            )));
306        }
307        state.active_transaction = Some(transaction_id);
308        state.read_write = read_write;
309        state.state = TransactionState::Active;
310        state.rollback_only = false;
311        Ok(())
312    }
313
314    /// Transition to committing state.  Returns the active transaction ID.
315    ///
316    /// After calling this, the caller should perform the actual commit
317    /// on the `TransactionManager`, then call [`commit_finalize`] on success
318    /// or [`fail_commit`] on error.
319    pub fn begin_commit(&self) -> Result<TransactionId, QueryError> {
320        let mut state = self.state.write();
321        if !state.state.can_commit() {
322            return Err(QueryError::execution(format!(
323                "Cannot COMMIT in state {:?}",
324                state.state
325            )));
326        }
327        let txn_id = state
328            .active_transaction
329            .ok_or_else(|| QueryError::execution("No active transaction to commit".to_string()))?;
330        state.state = TransactionState::Committing;
331        Ok(txn_id)
332    }
333
334    /// Finalize a successful commit.
335    pub fn commit_finalize(&self) {
336        let mut state = self.state.write();
337        state.active_transaction = None;
338        state.state = TransactionState::Committed;
339        state.rollback_only = false;
340    }
341
342    /// Mark commit as failed — transition to RollbackOnly.
343    pub fn fail_commit(&self) {
344        let mut state = self.state.write();
345        state.state = TransactionState::RollbackOnly;
346        state.rollback_only = true;
347    }
348
349    /// Transition to rolling-back state. Returns the active transaction ID.
350    pub fn begin_rollback(&self) -> Result<TransactionId, QueryError> {
351        let mut state = self.state.write();
352        if !state.state.can_rollback() {
353            return Err(QueryError::execution(format!(
354                "Cannot ROLLBACK in state {:?}",
355                state.state
356            )));
357        }
358        let txn_id = state.active_transaction.ok_or_else(|| {
359            QueryError::execution("No active transaction to roll back".to_string())
360        })?;
361        state.state = TransactionState::RollingBack;
362        Ok(txn_id)
363    }
364
365    /// Finalize a successful rollback.
366    pub fn rollback_finalize(&self) {
367        let mut state = self.state.write();
368        state.active_transaction = None;
369        state.state = TransactionState::RolledBack;
370        state.rollback_only = false;
371    }
372
373    /// Mark the current transaction as rollback-only.
374    pub fn mark_rollback_only(&self) {
375        let mut state = self.state.write();
376        if state.state == TransactionState::Active {
377            state.state = TransactionState::RollbackOnly;
378            state.rollback_only = true;
379        }
380    }
381
382    /// Enable or disable auto-commit mode.
383    pub fn set_auto_commit(&self, enabled: bool) {
384        self.state.write().auto_commit = enabled;
385    }
386
387    // ── Scope creation ──
388
389    /// Get the current explicit `TransactionScope` for query binding.
390    pub fn current_scope(&self) -> TransactionScope {
391        let state = self.state.read();
392        match state.active_transaction {
393            Some(txn_id) if state.state == TransactionState::Active => {
394                TransactionScope::ExplicitBorrowed {
395                    transaction_id: txn_id,
396                    read_write: state.read_write,
397                }
398            }
399            _ => TransactionScope::None,
400        }
401    }
402
403    /// Create a [`TransactionScope::CommandScope`] for transaction control
404    /// statements (BEGIN / COMMIT / ROLLBACK).
405    pub fn command_scope() -> TransactionScope {
406        TransactionScope::CommandScope
407    }
408}
409
410// ── Transaction command result ─────────────────────────────────────────────
411
412/// Result of a transaction command (BEGIN, COMMIT, ROLLBACK).
413#[derive(Debug, Clone)]
414pub struct TransactionCommandResult {
415    pub command: &'static str,
416    pub transaction_id: Option<TransactionId>,
417    pub message: String,
418}
419
420impl TransactionCommandResult {
421    pub fn begin(transaction_id: TransactionId) -> Self {
422        Self {
423            command: "BEGIN",
424            transaction_id: Some(transaction_id),
425            message: "Transaction started".to_string(),
426        }
427    }
428
429    pub fn commit(transaction_id: TransactionId) -> Self {
430        Self {
431            command: "COMMIT",
432            transaction_id: Some(transaction_id),
433            message: "Transaction committed".to_string(),
434        }
435    }
436
437    pub fn rollback(transaction_id: TransactionId) -> Self {
438        Self {
439            command: "ROLLBACK",
440            transaction_id: Some(transaction_id),
441            message: "Transaction rolled back".to_string(),
442        }
443    }
444
445    pub fn rollback_to(savepoint_name: String) -> Self {
446        Self {
447            command: "ROLLBACK TO SAVEPOINT",
448            transaction_id: None,
449            message: format!("Rolled back to savepoint '{}'", savepoint_name),
450        }
451    }
452
453    pub fn savepoint(savepoint_name: String) -> Self {
454        Self {
455            command: "SAVEPOINT",
456            transaction_id: None,
457            message: format!("Created savepoint '{}'", savepoint_name),
458        }
459    }
460
461    pub fn release_savepoint(savepoint_name: String) -> Self {
462        Self {
463            command: "RELEASE SAVEPOINT",
464            transaction_id: None,
465            message: format!("Released savepoint '{}'", savepoint_name),
466        }
467    }
468
469    /// Convert to an [`ExecutionResult`].
470    pub fn into_execution_result(self) -> ExecutionResult {
471        let row = vec![
472            linkrs_core::Value::string(self.command),
473            linkrs_core::Value::string(self.message),
474        ];
475        let col_names = vec!["command".to_string(), "result".to_string()];
476        let dataset = linkrs_core::DataSet::from_rows(vec![row], col_names);
477        ExecutionResult::DataSet { data: dataset }
478    }
479}
480
481#[cfg(test)]
482mod tests {
483    use super::*;
484
485    #[test]
486    fn test_scope_allows_write() {
487        let scope = TransactionScope::explicit(TransactionId(1), true);
488        assert!(scope.allows_write());
489        assert!(scope.is_explicit());
490        assert_eq!(scope.transaction_id(), Some(TransactionId(1)));
491    }
492
493    #[test]
494    fn test_read_only_scope_does_not_allow_write() {
495        let scope = TransactionScope::read_only(TransactionId(2));
496        assert!(!scope.allows_write());
497        assert!(!scope.is_explicit());
498        assert!(!scope.is_auto_commit());
499    }
500
501    #[test]
502    fn test_auto_commit_scope_allows_write() {
503        let scope = TransactionScope::auto_commit(TransactionId(3));
504        assert!(scope.allows_write());
505        assert!(scope.is_auto_commit());
506    }
507
508    #[test]
509    fn test_rollback_only_prevents_write() {
510        let mut scope = TransactionScope::auto_commit(TransactionId(4));
511        assert!(scope.allows_write());
512        scope.mark_rollback_only();
513        assert!(!scope.allows_write());
514        assert!(scope.is_rollback_only());
515    }
516
517    #[test]
518    fn test_none_scope() {
519        let scope = TransactionScope::None;
520        assert!(!scope.allows_write());
521        assert!(!scope.is_explicit());
522        assert!(!scope.is_auto_commit());
523        assert!(scope.transaction_id().is_none());
524    }
525
526    #[test]
527    fn test_command_scope() {
528        let scope = TransactionScope::CommandScope;
529        assert!(!scope.produces_result());
530    }
531
532    #[test]
533    fn test_ensure_can_write() {
534        let write_scope = TransactionScope::explicit(TransactionId(1), true);
535        assert!(write_scope.ensure_can_write().is_ok());
536
537        let read_scope = TransactionScope::read_only(TransactionId(2));
538        assert!(read_scope.ensure_can_write().is_err());
539    }
540
541    #[test]
542    fn test_session_controller_state() {
543        let ctrl = SessionTransactionController::new();
544        assert_eq!(ctrl.state(), TransactionState::None);
545        assert!(!ctrl.is_active());
546        assert!(ctrl.auto_commit());
547    }
548
549    #[test]
550    fn test_session_controller_tracking() {
551        let ctrl = SessionTransactionController::new();
552        assert!(ctrl.begin_tracking(TransactionId(1), true).is_ok());
553        assert_eq!(ctrl.state(), TransactionState::Active);
554        assert!(ctrl.is_active());
555        assert_eq!(ctrl.current_transaction_id(), Some(TransactionId(1)));
556
557        let txn_id = ctrl.begin_commit().unwrap();
558        assert_eq!(txn_id, TransactionId(1));
559        assert_eq!(ctrl.state(), TransactionState::Committing);
560        ctrl.commit_finalize();
561        assert_eq!(ctrl.state(), TransactionState::Committed);
562        assert!(!ctrl.is_active());
563    }
564
565    #[test]
566    fn test_session_controller_rollback() {
567        let ctrl = SessionTransactionController::new();
568        ctrl.begin_tracking(TransactionId(2), false).unwrap();
569        assert!(ctrl.is_active());
570
571        let txn_id = ctrl.begin_rollback().unwrap();
572        assert_eq!(txn_id, TransactionId(2));
573        ctrl.rollback_finalize();
574        assert_eq!(ctrl.state(), TransactionState::RolledBack);
575    }
576
577    #[test]
578    fn test_session_controller_cannot_double_begin() {
579        let ctrl = SessionTransactionController::new();
580        ctrl.begin_tracking(TransactionId(1), true).unwrap();
581        let result = ctrl.begin_tracking(TransactionId(2), true);
582        assert!(result.is_err());
583    }
584
585    #[test]
586    fn test_session_controller_rollback_only() {
587        let ctrl = SessionTransactionController::new();
588        ctrl.begin_tracking(TransactionId(1), true).unwrap();
589        ctrl.mark_rollback_only();
590        assert!(ctrl.rollback_only());
591        assert_eq!(ctrl.state(), TransactionState::RollbackOnly);
592
593        // Cannot commit a rollback-only transaction
594        assert!(ctrl.begin_commit().is_err());
595        // Can only rollback
596        assert!(ctrl.begin_rollback().is_ok());
597    }
598
599    #[test]
600    fn test_current_scope() {
601        let ctrl = SessionTransactionController::new();
602        assert!(matches!(ctrl.current_scope(), TransactionScope::None));
603
604        ctrl.begin_tracking(TransactionId(42), true).unwrap();
605        let scope = ctrl.current_scope();
606        assert!(scope.is_explicit());
607        assert!(scope.allows_write());
608        assert_eq!(scope.transaction_id(), Some(TransactionId(42)));
609    }
610
611    #[test]
612    fn test_cancel_reason_display() {
613        assert_eq!(CancelReason::UserKill.to_string(), "killed by user");
614        assert_eq!(CancelReason::Deadline.to_string(), "deadline exceeded");
615        assert_eq!(CancelReason::Shutdown.to_string(), "system shutdown");
616        assert_eq!(
617            CancelReason::Internal("test error".to_string()).to_string(),
618            "test error"
619        );
620    }
621
622    #[test]
623    fn test_transaction_state_transitions() {
624        assert!(TransactionState::None.can_begin());
625        assert!(TransactionState::Committed.can_begin());
626        assert!(TransactionState::RolledBack.can_begin());
627        assert!(!TransactionState::Active.can_begin());
628
629        assert!(TransactionState::Active.can_execute());
630        assert!(!TransactionState::None.can_execute());
631        assert!(!TransactionState::Committed.can_execute());
632
633        assert!(TransactionState::Active.can_commit());
634        assert!(!TransactionState::None.can_commit());
635
636        assert!(TransactionState::Active.can_rollback());
637        assert!(TransactionState::RollbackOnly.can_rollback());
638        assert!(!TransactionState::None.can_rollback());
639    }
640}
linkrs_query/executor/streaming/transaction_scope.rs

linkrs_query/executor/streaming/
transaction_scope.rs
