quill_sql/session/
mod.rs

1use crate::error::{QuillSQLError, QuillSQLResult};
2use crate::transaction::{IsolationLevel, Transaction};
3use sqlparser::ast::TransactionAccessMode;
4use std::sync::Arc;
5/// Session-level state used to manage transactions and defaults for a client connection.
6pub struct SessionContext {
7    default_isolation: IsolationLevel,
8    default_access_mode: TransactionAccessMode,
9    autocommit: bool,
10    active_txn: Option<Transaction>,
11    pending_session_isolation: Option<IsolationLevel>,
12    pending_session_access: Option<TransactionAccessMode>,
13}
14
15impl SessionContext {
16    pub fn new(default_isolation: IsolationLevel) -> Self {
17        Self {
18            default_isolation,
19            default_access_mode: TransactionAccessMode::ReadWrite,
20            autocommit: true,
21            active_txn: None,
22            pending_session_isolation: None,
23            pending_session_access: None,
24        }
25    }
26
27    pub fn default_isolation(&self) -> IsolationLevel {
28        self.default_isolation
29    }
30
31    pub fn set_default_isolation(&mut self, isolation: IsolationLevel) {
32        self.default_isolation = isolation;
33    }
34
35    pub fn default_access_mode(&self) -> TransactionAccessMode {
36        self.default_access_mode
37    }
38
39    pub fn set_default_access_mode(&mut self, mode: TransactionAccessMode) {
40        self.default_access_mode = mode;
41    }
42
43    /// Apply a new isolation level to the currently active transaction, if one exists.
44    pub fn set_active_isolation(&mut self, isolation: IsolationLevel) {
45        if let Some(txn) = self.active_txn.as_mut() {
46            txn.set_isolation_level(isolation);
47        }
48    }
49
50    pub fn autocommit(&self) -> bool {
51        self.autocommit
52    }
53
54    pub fn set_autocommit(&mut self, enabled: bool) {
55        self.autocommit = enabled;
56    }
57
58    /// Retrieve the pending session isolation override that will apply to the next transaction.
59    pub fn pending_session_isolation(&self) -> Option<IsolationLevel> {
60        self.pending_session_isolation
61    }
62
63    /// Record a pending session isolation override that should apply to the next transaction.
64    pub fn set_pending_session_isolation(&mut self, isolation: Option<IsolationLevel>) {
65        self.pending_session_isolation = isolation;
66    }
67
68    pub fn pending_session_access(&self) -> Option<TransactionAccessMode> {
69        self.pending_session_access
70    }
71
72    pub fn set_pending_session_access(&mut self, mode: Option<TransactionAccessMode>) {
73        self.pending_session_access = mode;
74    }
75
76    pub fn has_active_transaction(&self) -> bool {
77        self.active_txn.is_some()
78    }
79
80    pub fn active_txn(&self) -> Option<&Transaction> {
81        self.active_txn.as_ref()
82    }
83
84    pub fn active_txn_mut(&mut self) -> Option<&mut Transaction> {
85        self.active_txn.as_mut()
86    }
87
88    pub fn ensure_active_transaction(
89        &mut self,
90        txn_mgr: &Arc<crate::transaction::TransactionManager>,
91    ) -> QuillSQLResult<&mut Transaction> {
92        if self.active_txn.is_none() {
93            let isolation = self
94                .pending_session_isolation
95                .unwrap_or(self.default_isolation);
96            let access_mode = self
97                .pending_session_access
98                .unwrap_or(self.default_access_mode);
99            let mut txn = txn_mgr.begin(isolation, access_mode)?;
100            if let Some(isolation) = self.pending_session_isolation.take() {
101                self.default_isolation = isolation;
102                txn.set_isolation_level(isolation);
103            }
104            if let Some(mode) = self.pending_session_access.take() {
105                self.default_access_mode = mode;
106                txn.update_access_mode(mode);
107            }
108            self.active_txn = Some(txn);
109        }
110        self.active_txn
111            .as_mut()
112            .ok_or_else(|| QuillSQLError::Execution("failed to start transaction".to_string()))
113    }
114
115    pub fn set_active_transaction(&mut self, mut txn: Transaction) -> QuillSQLResult<()> {
116        if self.active_txn.is_some() {
117            return Err(QuillSQLError::Execution(
118                "transaction already active in session".to_string(),
119            ));
120        }
121        if let Some(isolation) = self.pending_session_isolation.take() {
122            self.default_isolation = isolation;
123            txn.set_isolation_level(isolation);
124        } else {
125            self.default_isolation = txn.isolation_level();
126        }
127
128        if let Some(mode) = self.pending_session_access.take() {
129            self.default_access_mode = mode;
130            txn.update_access_mode(mode);
131        } else {
132            self.default_access_mode = txn.access_mode();
133        }
134
135        self.active_txn = Some(txn);
136        Ok(())
137    }
138
139    pub fn take_active_transaction(&mut self) -> Option<Transaction> {
140        self.active_txn.take()
141    }
142
143    pub fn clear_active_transaction(&mut self) {
144        self.active_txn = None;
145    }
146
147    /// Apply isolation override when a transaction-scoped SET TRANSACTION is issued.
148    pub fn apply_transaction_modes(&mut self, modes: &crate::plan::logical_plan::TransactionModes) {
149        if let Some(level) = modes.isolation_level {
150            if let Some(txn) = self.active_txn.as_mut() {
151                txn.set_isolation_level(level);
152            }
153        }
154        if let Some(mode) = modes.access_mode {
155            if let Some(txn) = self.active_txn.as_mut() {
156                txn.update_access_mode(mode);
157            }
158        }
159    }
160
161    /// Merge `SET SESSION TRANSACTION` modes into session defaults.
162    pub fn apply_session_modes(&mut self, modes: &crate::plan::logical_plan::TransactionModes) {
163        if let Some(level) = modes.isolation_level {
164            self.default_isolation = level;
165            self.pending_session_isolation = Some(level);
166        }
167        if let Some(mode) = modes.access_mode {
168            self.default_access_mode = mode;
169            self.pending_session_access = Some(mode);
170        }
171    }
172}
quill_sql/session/mod.rs

quill_sql/session/
mod.rs
