alopex_server/
session.rs

1use std::sync::Arc;
2use std::time::{Duration, SystemTime};
3
4use alopex_cluster::TableLifecycleEffect;
5use alopex_core::async_runtime::{BoxFuture, BoxStream};
6use alopex_sql::catalog::TableMetadata;
7use alopex_sql::executor::{ExecutionResult, ExecutorError, Row};
8use alopex_sql::planner::PlannedStatement;
9use alopex_sql::storage::erased::ErasedAsyncSqlTransaction;
10use dashmap::DashMap;
11use futures::StreamExt;
12use tokio::sync::mpsc;
13use tokio_stream::wrappers::ReceiverStream;
14use uuid::Uuid;
15
16use crate::error::{Result, ServerError};
17
18/// Session identifier.
19#[derive(Clone, Debug, Eq, Hash, PartialEq, serde::Serialize, serde::Deserialize)]
20pub struct SessionId(Uuid);
21
22impl SessionId {
23    pub fn new() -> Self {
24        Self(Uuid::new_v4())
25    }
26}
27
28impl Default for SessionId {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl std::fmt::Display for SessionId {
35    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
36        write!(f, "{}", self.0)
37    }
38}
39
40impl std::str::FromStr for SessionId {
41    type Err = uuid::Error;
42
43    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
44        Ok(Self(Uuid::parse_str(s)?))
45    }
46}
47
48/// Session lifecycle state.
49#[derive(Clone, Copy, Debug, Eq, PartialEq, serde::Serialize, serde::Deserialize)]
50pub enum SessionState {
51    Idle,
52    InTransaction,
53    Committing,
54    RollingBack,
55}
56
57/// Snapshot of a session for safe sharing.
58#[derive(Clone, Debug, serde::Serialize)]
59pub struct SessionSnapshot {
60    pub id: SessionId,
61    pub has_transaction: bool,
62    pub created_at: SystemTime,
63    pub last_active: SystemTime,
64    pub expires_at: SystemTime,
65    pub state: SessionState,
66}
67
68/// Transaction handle for a session.
69#[derive(Clone)]
70pub struct TxnHandle {
71    inner: Arc<TxnHandleInner>,
72}
73
74#[derive(Clone)]
75pub enum CatalogRollbackEffect {
76    DropTable { table_name: String },
77    CreateTable { table: Box<TableMetadata> },
78}
79
80struct TxnHandleInner {
81    txn: tokio::sync::Mutex<Option<Box<dyn ErasedAsyncSqlTransaction>>>,
82    pending_table_lifecycle_effects: tokio::sync::Mutex<Vec<TableLifecycleEffect>>,
83    pending_catalog_rollback_effects: tokio::sync::Mutex<Vec<CatalogRollbackEffect>>,
84    created_at: SystemTime,
85}
86
87impl TxnHandle {
88    pub fn new(txn: Box<dyn ErasedAsyncSqlTransaction>) -> Self {
89        Self {
90            inner: Arc::new(TxnHandleInner {
91                txn: tokio::sync::Mutex::new(Some(txn)),
92                pending_table_lifecycle_effects: tokio::sync::Mutex::new(Vec::new()),
93                pending_catalog_rollback_effects: tokio::sync::Mutex::new(Vec::new()),
94                created_at: SystemTime::now(),
95            }),
96        }
97    }
98
99    pub fn created_at(&self) -> SystemTime {
100        self.inner.created_at
101    }
102
103    pub fn execute<'a>(
104        &'a self,
105        sql: &'a str,
106    ) -> BoxFuture<'a, alopex_sql::executor::Result<ExecutionResult>> {
107        Box::pin(async move {
108            let mut guard = self.inner.txn.lock().await;
109            let txn = guard
110                .as_mut()
111                .ok_or_else(|| ExecutorError::InvalidOperation {
112                    operation: "execute".into(),
113                    reason: "transaction is closed".into(),
114                })?;
115            txn.execute(sql).await
116        })
117    }
118
119    pub fn execute_multi<'a>(
120        &'a self,
121        sql: &'a str,
122    ) -> BoxFuture<'a, alopex_sql::executor::Result<Vec<ExecutionResult>>> {
123        Box::pin(async move {
124            let mut guard = self.inner.txn.lock().await;
125            let txn = guard
126                .as_mut()
127                .ok_or_else(|| ExecutorError::InvalidOperation {
128                    operation: "execute_multi".into(),
129                    reason: "transaction is closed".into(),
130                })?;
131            txn.execute_multi(sql).await
132        })
133    }
134
135    pub fn query<'a>(&'a self, sql: &'a str) -> BoxStream<'a, alopex_sql::executor::Result<Row>> {
136        let (sender, receiver) = mpsc::channel(32);
137        let sql = sql.to_string();
138        let inner = Arc::clone(&self.inner);
139
140        tokio::spawn(async move {
141            let guard = inner.txn.lock().await;
142            let Some(txn) = guard.as_ref() else {
143                let _ = sender
144                    .send(Err(ExecutorError::InvalidOperation {
145                        operation: "query".into(),
146                        reason: "transaction is closed".into(),
147                    }))
148                    .await;
149                return;
150            };
151            let mut stream = txn.query(&sql);
152            while let Some(item) = stream.next().await {
153                if sender.send(item).await.is_err() {
154                    break;
155                }
156            }
157        });
158
159        Box::pin(ReceiverStream::new(receiver))
160    }
161
162    pub fn plan_for_routing<'a>(
163        &'a self,
164        sql: &'a str,
165    ) -> BoxFuture<'a, alopex_sql::executor::Result<Vec<PlannedStatement>>> {
166        Box::pin(async move {
167            let guard = self.inner.txn.lock().await;
168            let txn = guard
169                .as_ref()
170                .ok_or_else(|| ExecutorError::InvalidOperation {
171                    operation: "plan_for_routing".into(),
172                    reason: "transaction is closed".into(),
173                })?;
174            txn.plan_for_routing(sql).await
175        })
176    }
177
178    pub async fn buffer_table_lifecycle_effects(&self, effects: Vec<TableLifecycleEffect>) {
179        if effects.is_empty() {
180            return;
181        }
182        self.inner
183            .pending_table_lifecycle_effects
184            .lock()
185            .await
186            .extend(effects);
187    }
188
189    pub async fn buffer_catalog_rollback_effects(&self, effects: Vec<CatalogRollbackEffect>) {
190        if effects.is_empty() {
191            return;
192        }
193        self.inner
194            .pending_catalog_rollback_effects
195            .lock()
196            .await
197            .extend(effects);
198    }
199
200    pub async fn commit(self) -> alopex_sql::executor::Result<Vec<TableLifecycleEffect>> {
201        let mut guard = self.inner.txn.lock().await;
202        let txn = guard
203            .take()
204            .ok_or_else(|| ExecutorError::InvalidOperation {
205                operation: "commit".into(),
206                reason: "transaction is closed".into(),
207            })?;
208        txn.commit_boxed().await?;
209        let mut effects = self.inner.pending_table_lifecycle_effects.lock().await;
210        Ok(std::mem::take(&mut *effects))
211    }
212
213    pub async fn rollback(self) -> alopex_sql::executor::Result<Vec<CatalogRollbackEffect>> {
214        let mut guard = self.inner.txn.lock().await;
215        let txn = guard
216            .take()
217            .ok_or_else(|| ExecutorError::InvalidOperation {
218                operation: "rollback".into(),
219                reason: "transaction is closed".into(),
220            })?;
221        let result = txn.rollback_boxed().await;
222        self.inner
223            .pending_table_lifecycle_effects
224            .lock()
225            .await
226            .clear();
227        result?;
228        let mut effects = self.inner.pending_catalog_rollback_effects.lock().await;
229        Ok(std::mem::take(&mut *effects))
230    }
231}
232
233/// Session configuration.
234#[derive(Clone, Copy, Debug)]
235pub struct SessionConfig {
236    pub ttl: Duration,
237}
238
239/// Transaction factory for session manager.
240pub type TransactionFactory =
241    Arc<dyn Fn() -> BoxFuture<'static, Result<Box<dyn ErasedAsyncSqlTransaction>>> + Send + Sync>;
242
243/// Session manager for server.
244pub struct SessionManager {
245    sessions: DashMap<SessionId, Session>,
246    config: SessionConfig,
247    txn_factory: TransactionFactory,
248}
249
250struct Session {
251    id: SessionId,
252    txn_handle: Option<TxnHandle>,
253    created_at: SystemTime,
254    last_active: SystemTime,
255    expires_at: SystemTime,
256    state: SessionState,
257}
258
259impl SessionManager {
260    pub fn new(config: SessionConfig, txn_factory: TransactionFactory) -> Self {
261        Self {
262            sessions: DashMap::new(),
263            config,
264            txn_factory,
265        }
266    }
267
268    pub async fn create_session(&self) -> Result<SessionId> {
269        let now = SystemTime::now();
270        let id = SessionId::new();
271        let session = Session {
272            id: id.clone(),
273            txn_handle: None,
274            created_at: now,
275            last_active: now,
276            expires_at: now + self.config.ttl,
277            state: SessionState::Idle,
278        };
279        self.sessions.insert(id.clone(), session);
280        Ok(id)
281    }
282
283    pub async fn get_session(&self, id: &SessionId) -> Result<SessionSnapshot> {
284        let entry = self
285            .sessions
286            .get(id)
287            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
288        if entry.expires_at <= SystemTime::now() {
289            drop(entry);
290            self.sessions.remove(id);
291            return Err(ServerError::SessionExpired("session expired".into()));
292        }
293        Ok(SessionSnapshot {
294            id: entry.id.clone(),
295            has_transaction: entry.txn_handle.is_some(),
296            created_at: entry.created_at,
297            last_active: entry.last_active,
298            expires_at: entry.expires_at,
299            state: entry.state,
300        })
301    }
302
303    pub async fn begin_transaction(&self, id: &SessionId) -> Result<TxnHandle> {
304        let mut entry = self
305            .sessions
306            .get_mut(id)
307            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
308        if entry.expires_at <= SystemTime::now() {
309            drop(entry);
310            self.sessions.remove(id);
311            return Err(ServerError::SessionExpired("session expired".into()));
312        }
313        if entry.txn_handle.is_some() {
314            return Err(ServerError::Conflict("transaction already active".into()));
315        }
316        let txn = (self.txn_factory)().await?;
317        let handle = TxnHandle::new(txn);
318        entry.txn_handle = Some(handle.clone());
319        entry.last_active = SystemTime::now();
320        entry.state = SessionState::InTransaction;
321        Ok(handle)
322    }
323
324    pub async fn get_transaction(&self, id: &SessionId) -> Result<TxnHandle> {
325        let mut entry = self
326            .sessions
327            .get_mut(id)
328            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
329        if entry.expires_at <= SystemTime::now() {
330            drop(entry);
331            self.sessions.remove(id);
332            return Err(ServerError::SessionExpired("session expired".into()));
333        }
334        let handle = entry
335            .txn_handle
336            .clone()
337            .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
338        entry.last_active = SystemTime::now();
339        entry.state = SessionState::InTransaction;
340        Ok(handle)
341    }
342
343    pub async fn execute_in_session(&self, id: &SessionId, sql: &str) -> Result<ExecutionResult> {
344        let handle = {
345            let mut entry = self
346                .sessions
347                .get_mut(id)
348                .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
349            if entry.expires_at <= SystemTime::now() {
350                drop(entry);
351                self.sessions.remove(id);
352                return Err(ServerError::SessionExpired("session expired".into()));
353            }
354            let handle = entry
355                .txn_handle
356                .clone()
357                .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
358            entry.last_active = SystemTime::now();
359            handle
360        };
361
362        handle
363            .execute(sql)
364            .await
365            .map_err(|err| ServerError::Sql(err.into()))
366    }
367
368    pub async fn commit(&self, id: &SessionId) -> Result<Vec<TableLifecycleEffect>> {
369        let handle = self.take_handle(id, SessionState::Committing)?;
370        let effects = handle
371            .commit()
372            .await
373            .map_err(|err| ServerError::Sql(err.into()))?;
374        Ok(effects)
375    }
376
377    pub async fn rollback(&self, id: &SessionId) -> Result<Vec<CatalogRollbackEffect>> {
378        let handle = self.take_handle(id, SessionState::RollingBack)?;
379        let effects = handle
380            .rollback()
381            .await
382            .map_err(|err| ServerError::Sql(err.into()))?;
383        Ok(effects)
384    }
385
386    pub fn cleanup_expired(&self) {
387        let now = SystemTime::now();
388        let expired: Vec<SessionId> = self
389            .sessions
390            .iter()
391            .filter(|entry| entry.expires_at <= now)
392            .map(|entry| entry.id.clone())
393            .collect();
394        for id in expired {
395            self.sessions.remove(&id);
396        }
397    }
398
399    fn take_handle(&self, id: &SessionId, state: SessionState) -> Result<TxnHandle> {
400        let mut entry = self
401            .sessions
402            .get_mut(id)
403            .ok_or_else(|| ServerError::NotFound("session not found".into()))?;
404        if entry.expires_at <= SystemTime::now() {
405            drop(entry);
406            self.sessions.remove(id);
407            return Err(ServerError::SessionExpired("session expired".into()));
408        }
409        let handle = entry
410            .txn_handle
411            .take()
412            .ok_or_else(|| ServerError::BadRequest("transaction not started".into()))?;
413        entry.state = state;
414        entry.last_active = SystemTime::now();
415        Ok(handle)
416    }
417}
alopex_server/session.rs

alopex_server/
session.rs
