linkrs_server/client/
client_session.rs

1use parking_lot::RwLock;
2use std::collections::HashMap;
3use std::sync::{Arc, Weak};
4use std::time::Instant;
5
6use linkrs_query::executor::streaming::runtime::ExecutionRuntime;
7use linkrs_query::executor::streaming::StreamingQueryResult;
8
9use super::cursor_context::{CursorContext, CursorPage};
10use super::query_context::QueryContext;
11use super::role_context::RoleContext;
12use super::session::Session;
13use super::space_context::SpaceContext;
14use super::statistics::StatisticsContext;
15use super::transaction_context::TransactionContext;
16use crate::api::session_variables::SessionVariables;
17use linkrs_core::error::QueryResult;
18use linkrs_core::types::SpaceSummary;
19use linkrs_core::Value;
20
21#[derive(Debug)]
22pub struct ClientSession {
23    session_context: super::session::SessionContext,
24    space_context: SpaceContext,
25    role_context: RoleContext,
26    query_context: QueryContext,
27    transaction_context: TransactionContext,
28    statistics_context: StatisticsContext,
29    cursor_context: CursorContext,
30    idle_start_time: Arc<RwLock<Instant>>,
31    /// Session-scoped user variables (`$name`) with transaction overlay.
32    session_variables: Arc<SessionVariables>,
33}
34
35impl ClientSession {
36    pub fn new(session: Session) -> Arc<Self> {
37        Arc::new(Self {
38            session_context: super::session::SessionContext::new(session),
39            space_context: SpaceContext::new(),
40            role_context: RoleContext::new(),
41            query_context: QueryContext::new(),
42            transaction_context: TransactionContext::new(),
43            statistics_context: StatisticsContext::new(),
44            cursor_context: CursorContext::new(),
45            idle_start_time: Arc::new(RwLock::new(Instant::now())),
46            session_variables: Arc::new(SessionVariables::new()),
47        })
48    }
49
50    pub fn id(&self) -> i64 {
51        self.session_context.id()
52    }
53
54    pub fn space(&self) -> Option<SpaceSummary> {
55        self.space_context.space()
56    }
57
58    pub fn set_space(&self, space: SpaceSummary) {
59        self.session_context.update_space_name(space.name.clone());
60        self.space_context.set_space(space);
61    }
62
63    pub fn clear_space(&self) {
64        self.space_context.clear_space();
65    }
66
67    pub fn space_name(&self) -> Option<String> {
68        self.session_context.space_name()
69    }
70
71    pub fn user(&self) -> String {
72        self.session_context.user()
73    }
74
75    pub fn roles(&self) -> std::collections::HashMap<i64, linkrs_core::RoleType> {
76        self.role_context.roles()
77    }
78
79    pub fn role_with_space(&self, space: i64) -> Option<linkrs_core::RoleType> {
80        self.role_context.role_with_space(space)
81    }
82
83    pub fn is_god(&self) -> bool {
84        self.role_context.is_god()
85    }
86
87    pub fn is_admin(&self) -> bool {
88        self.role_context.is_admin()
89    }
90
91    pub fn is_schema_admin(&self) -> bool {
92        self.role_context.is_schema_admin()
93    }
94
95    pub fn set_role(&self, space: i64, role: linkrs_core::RoleType) {
96        self.role_context.set_role(space, role);
97    }
98
99    pub fn idle_seconds(&self) -> u64 {
100        self.idle_start_time.read().elapsed().as_secs()
101    }
102
103    pub fn charge(&self) {
104        *self.idle_start_time.write() = Instant::now();
105    }
106
107    pub fn timezone(&self) -> Option<i32> {
108        self.session_context.timezone()
109    }
110
111    pub fn set_timezone(&self, timezone: i32) {
112        self.session_context.set_timezone(timezone);
113    }
114
115    pub fn graph_addr(&self) -> Option<String> {
116        self.session_context.graph_addr()
117    }
118
119    pub fn update_graph_addr(&self, host_addr: String) {
120        self.session_context.update_graph_addr(host_addr);
121    }
122
123    pub fn get_session(&self) -> Session {
124        self.session_context.get_session()
125    }
126
127    pub fn update_space_name(&self, space_name: String) {
128        self.session_context.update_space_name(space_name);
129    }
130
131    pub fn add_query(&self, ep_id: u32, query_context: String) {
132        self.query_context
133            .add_query(ep_id, query_context, self.id());
134    }
135
136    /// Register a streaming query with a weak runtime reference for KILL QUERY support.
137    pub fn register_streaming_query(
138        &self,
139        query_id: u32,
140        query_text: String,
141        runtime: Weak<ExecutionRuntime>,
142    ) {
143        self.query_context
144            .register_streaming_query(query_id, query_text, runtime, self.id());
145    }
146
147    /// Unregister a streaming query on completion.
148    pub fn unregister_streaming_query(&self, query_id: u32) {
149        self.query_context
150            .unregister_streaming_query(query_id, self.id());
151    }
152
153    pub fn delete_query(&self, ep_id: u32) {
154        self.query_context.delete_query(ep_id, self.id());
155    }
156
157    pub fn find_query(&self, ep_id: u32) -> bool {
158        self.query_context.find_query(ep_id)
159    }
160
161    pub fn mark_query_killed(&self, ep_id: u32) {
162        self.query_context.mark_query_killed(ep_id, self.id());
163    }
164
165    pub fn mark_all_queries_killed(&self) {
166        self.query_context.mark_all_queries_killed(self.id());
167    }
168
169    pub fn active_queries_count(&self) -> usize {
170        self.query_context.active_queries_count()
171    }
172
173    pub fn kill_query(&self, query_id: u32) -> QueryResult<()> {
174        self.query_context.kill_query(query_id, self.id())
175    }
176
177    pub fn kill_multiple_queries(&self, query_ids: &[u32]) -> Vec<QueryResult<()>> {
178        self.query_context
179            .kill_multiple_queries(query_ids, self.id())
180    }
181
182    /// Open a forward-only cursor over an already-started streaming result.
183    pub fn open_cursor(
184        &self,
185        query: String,
186        result: StreamingQueryResult,
187        columns: Vec<String>,
188    ) -> Result<u64, String> {
189        self.cursor_context.open(query, result, columns)
190    }
191
192    /// Fetch one page from a cursor (blocking; run off the async runtime).
193    pub fn fetch_cursor_page(
194        &self,
195        cursor_id: u64,
196        page_size: usize,
197    ) -> Result<CursorPage, String> {
198        self.cursor_context.fetch(cursor_id, page_size)
199    }
200
201    /// Release a cursor; returns false when the id was unknown.
202    pub fn close_cursor(&self, cursor_id: u64) -> bool {
203        self.cursor_context.close(cursor_id)
204    }
205
206    /// Release every cursor held by this session.
207    pub fn close_all_cursors(&self) {
208        self.cursor_context.close_all();
209    }
210
211    /// Number of currently open cursors.
212    pub fn cursor_count(&self) -> usize {
213        self.cursor_context.active_count()
214    }
215
216    pub fn current_transaction(&self) -> Option<linkrs_transaction::TransactionId> {
217        self.transaction_context.current_transaction()
218    }
219
220    pub fn bind_transaction(&self, txn_id: linkrs_transaction::TransactionId) {
221        self.transaction_context.bind_transaction(txn_id, self.id());
222    }
223
224    pub fn unbind_transaction(&self) {
225        self.transaction_context.unbind_transaction(self.id());
226    }
227
228    pub fn has_active_transaction(&self) -> bool {
229        self.transaction_context.has_active_transaction()
230    }
231
232    pub fn is_auto_commit(&self) -> bool {
233        self.transaction_context.is_auto_commit()
234    }
235
236    pub fn set_auto_commit(&self, auto_commit: bool) {
237        self.transaction_context
238            .set_auto_commit(auto_commit, self.id());
239    }
240
241    pub fn transaction_options(&self) -> linkrs_transaction::TransactionOptions {
242        self.transaction_context.transaction_options()
243    }
244
245    pub fn set_transaction_options(&self, options: linkrs_transaction::TransactionOptions) {
246        self.transaction_context.set_transaction_options(options);
247    }
248
249    pub fn push_savepoint(&self, savepoint_id: linkrs_transaction::SavepointId) {
250        self.transaction_context
251            .push_savepoint(savepoint_id, self.id());
252    }
253
254    pub fn savepoint_stack(&self) -> Vec<linkrs_transaction::SavepointId> {
255        self.transaction_context.savepoint_stack()
256    }
257
258    pub fn clear_savepoints(&self) {
259        self.transaction_context.clear_savepoints(self.id());
260    }
261
262    pub fn savepoint_count(&self) -> usize {
263        self.transaction_context.savepoint_count()
264    }
265
266    pub fn statistics(&self) -> &linkrs_core::SessionStatistics {
267        self.statistics_context.statistics()
268    }
269
270    // ── Session variables (`$name`) ───────────────────────────────────────
271
272    /// Assign a session variable, recording a transaction overlay operation
273    /// when an explicit transaction is active.
274    pub fn set_variable(&self, name: String, value: Value) {
275        self.session_variables
276            .set_variable(name, value, self.has_active_transaction());
277    }
278
279    /// Effective value of a session variable: transaction overlay first,
280    /// then the base store.
281    pub fn variable_value(&self, name: &str) -> Option<Value> {
282        self.session_variables.variable_value(name)
283    }
284
285    /// Snapshot of all session variables (base + overlay) for injection as
286    /// query inputs.
287    pub fn variables_snapshot(&self) -> HashMap<String, Value> {
288        self.session_variables.variables_snapshot()
289    }
290
291    /// COMMIT: apply overlay operations to the base store and clear the
292    /// overlay (the last assignment of each variable wins).
293    pub fn commit_variables(&self) {
294        self.session_variables.commit_variables();
295    }
296
297    /// Full ROLLBACK: restore pre-transaction values and clear the overlay.
298    pub fn rollback_variables(&self) {
299        self.session_variables.rollback_variables();
300    }
301
302    /// ROLLBACK TO SAVEPOINT: restore values assigned after the named
303    /// savepoint and truncate the overlay at the savepoint marker.
304    pub fn rollback_variables_to(&self, savepoint_name: &str) -> bool {
305        self.session_variables.rollback_variables_to(savepoint_name)
306    }
307
308    /// RELEASE SAVEPOINT: drop the marker (operations stay part of the
309    /// transaction; they are no longer individually rollback-able).
310    pub fn release_variable_savepoint(&self, savepoint_name: &str) {
311        self.session_variables
312            .release_variable_savepoint(savepoint_name);
313    }
314
315    /// SAVEPOINT: record a variable-overlay boundary so ROLLBACK TO can
316    /// restore assignments made after the savepoint.
317    pub fn push_variable_savepoint(&self, savepoint_name: &str) {
318        self.session_variables
319            .push_variable_savepoint(savepoint_name);
320    }
321}
322
323#[cfg(test)]
324mod tests {
325    use super::*;
326    use linkrs_core::types::DataType;
327
328    #[test]
329    fn test_client_session_creation() {
330        let session = Session {
331            session_id: 123,
332            user_name: "testuser".to_string(),
333            space_name: None,
334            graph_addr: None,
335            timezone: None,
336        };
337
338        let client_session = ClientSession::new(session);
339
340        assert_eq!(client_session.id(), 123);
341        assert_eq!(client_session.user(), "testuser");
342        assert_eq!(client_session.roles().len(), 0);
343        assert!(!client_session.is_admin());
344    }
345
346    #[test]
347    fn test_client_session_space_management() {
348        let session = Session {
349            session_id: 123,
350            user_name: "testuser".to_string(),
351            space_name: None,
352            graph_addr: None,
353            timezone: None,
354        };
355
356        let client_session = ClientSession::new(session);
357
358        assert!(client_session.space().is_none());
359        assert!(client_session.space_name().is_none());
360
361        let space_info = SpaceSummary::new(456, "test_space".to_string(), DataType::BigInt);
362        client_session.set_space(space_info.clone());
363
364        assert_eq!(client_session.space().expect("space should exist").id, 456);
365        assert_eq!(
366            client_session.space().expect("space should exist").name,
367            "test_space"
368        );
369        client_session.update_space_name("new_space".to_string());
370        assert_eq!(
371            client_session
372                .space_name()
373                .expect("space_name should exist"),
374            "new_space"
375        );
376
377        client_session.clear_space();
378        assert!(client_session.space().is_none());
379    }
380
381    #[test]
382    fn test_client_session_role_management() {
383        let session = Session {
384            session_id: 123,
385            user_name: "testuser".to_string(),
386            space_name: None,
387            graph_addr: None,
388            timezone: None,
389        };
390
391        let client_session = ClientSession::new(session);
392
393        assert!(client_session.role_with_space(1).is_none());
394        assert!(!client_session.is_admin());
395        assert!(!client_session.is_god());
396
397        client_session.set_role(1, linkrs_core::RoleType::Admin);
398        assert_eq!(
399            client_session
400                .role_with_space(1)
401                .expect("role should exist"),
402            linkrs_core::RoleType::Admin
403        );
404        assert!(client_session.is_admin());
405        assert!(!client_session.is_god());
406
407        client_session.set_role(2, linkrs_core::RoleType::God);
408        assert!(client_session.is_god());
409    }
410
411    #[test]
412    fn test_client_session_idle_time() {
413        let session = Session {
414            session_id: 123,
415            user_name: "testuser".to_string(),
416            space_name: None,
417            graph_addr: None,
418            timezone: None,
419        };
420
421        let client_session = ClientSession::new(session);
422
423        let idle1 = client_session.idle_seconds();
424        assert_eq!(idle1, 0);
425
426        std::thread::sleep(std::time::Duration::from_millis(1100));
427        let idle2 = client_session.idle_seconds();
428        assert!(idle2 > 0);
429
430        client_session.charge();
431        let idle3 = client_session.idle_seconds();
432        assert_eq!(idle3, 0);
433    }
434
435    #[test]
436    fn test_client_session_query_management() {
437        let session = Session {
438            session_id: 123,
439            user_name: "testuser".to_string(),
440            space_name: None,
441            graph_addr: None,
442            timezone: None,
443        };
444
445        let client_session = ClientSession::new(session);
446
447        assert_eq!(client_session.active_queries_count(), 0);
448        assert!(!client_session.find_query(1));
449
450        client_session.add_query(1, "SELECT * FROM user".to_string());
451        assert_eq!(client_session.active_queries_count(), 1);
452        assert!(client_session.find_query(1));
453
454        client_session.delete_query(1);
455        assert_eq!(client_session.active_queries_count(), 0);
456        assert!(!client_session.find_query(1));
457
458        client_session.add_query(2, "MATCH (n) RETURN n".to_string());
459        let result = client_session.kill_query(2);
460        assert!(result.is_ok());
461        assert!(!client_session.find_query(2));
462
463        let result = client_session.kill_query(999);
464        assert!(result.is_err());
465
466        client_session.add_query(3, "query 3".to_string());
467        client_session.add_query(4, "query 4".to_string());
468        let results = client_session.kill_multiple_queries(&[3, 4, 5]);
469        assert_eq!(results.len(), 3);
470        assert!(results[0].is_ok());
471        assert!(results[1].is_ok());
472        assert!(results[2].is_err());
473    }
474
475    #[test]
476    fn test_client_session_transaction_management() {
477        let session = Session {
478            session_id: 123,
479            user_name: "testuser".to_string(),
480            space_name: None,
481            graph_addr: None,
482            timezone: None,
483        };
484
485        let client_session = ClientSession::new(session);
486
487        assert!(client_session.current_transaction().is_none());
488        assert!(!client_session.has_active_transaction());
489        assert!(client_session.is_auto_commit());
490
491        client_session.bind_transaction(linkrs_transaction::TransactionId(1001));
492        assert_eq!(
493            client_session
494                .current_transaction()
495                .expect("current_transaction should exist"),
496            linkrs_transaction::TransactionId(1001)
497        );
498        assert!(client_session.has_active_transaction());
499
500        client_session.unbind_transaction();
501        assert!(client_session.current_transaction().is_none());
502
503        client_session.set_auto_commit(false);
504        assert!(!client_session.is_auto_commit());
505
506        let options = linkrs_transaction::TransactionOptions::default();
507        client_session.set_transaction_options(options.clone());
508        assert_eq!(client_session.transaction_options(), options);
509    }
510
511    #[test]
512    fn test_client_session_savepoint_management() {
513        let session = Session {
514            session_id: 123,
515            user_name: "testuser".to_string(),
516            space_name: None,
517            graph_addr: None,
518            timezone: None,
519        };
520
521        let client_session = ClientSession::new(session);
522
523        assert_eq!(client_session.savepoint_count(), 0);
524        assert!(client_session.savepoint_stack().is_empty());
525
526        client_session.push_savepoint(1);
527        client_session.push_savepoint(2);
528        assert_eq!(client_session.savepoint_count(), 2);
529        assert_eq!(client_session.savepoint_stack(), vec![1, 2]);
530
531        client_session.clear_savepoints();
532        assert_eq!(client_session.savepoint_count(), 0);
533    }
534
535    #[test]
536    fn test_session_variables_outside_transaction() {
537        let session = Session {
538            session_id: 123,
539            user_name: "testuser".to_string(),
540            space_name: None,
541            graph_addr: None,
542            timezone: None,
543        };
544
545        let client_session = ClientSession::new(session);
546
547        client_session.set_variable("x".into(), Value::Int(1));
548        assert_eq!(
549            client_session.variable_value("x"),
550            Some(Value::Int(1)),
551            "set outside a transaction writes the base store"
552        );
553
554        // Overwrite.
555        client_session.set_variable("x".into(), Value::Int(2));
556        assert_eq!(client_session.variable_value("x"), Some(Value::Int(2)));
557
558        let snapshot = client_session.variables_snapshot();
559        assert_eq!(snapshot.get("x"), Some(&Value::Int(2)));
560    }
561
562    #[test]
563    fn test_session_variables_transaction_rollback_restores() {
564        let session = Session {
565            session_id: 123,
566            user_name: "testuser".to_string(),
567            space_name: None,
568            graph_addr: None,
569            timezone: None,
570        };
571
572        let client_session = ClientSession::new(session);
573
574        // Pre-transaction value.
575        client_session.set_variable("x".into(), Value::Int(1));
576
577        client_session.bind_transaction(linkrs_transaction::TransactionId(7));
578        assert!(client_session.has_active_transaction());
579
580        // Assignments inside the transaction go to the overlay.
581        client_session.set_variable("x".into(), Value::Int(100));
582        client_session.set_variable("y".into(), Value::string("txn"));
583        assert_eq!(client_session.variable_value("x"), Some(Value::Int(100)));
584        assert_eq!(
585            client_session.variable_value("y"),
586            Some(Value::string("txn"))
587        );
588
589        // Full rollback restores pre-transaction values.
590        client_session.rollback_variables();
591        client_session.unbind_transaction();
592        assert_eq!(client_session.variable_value("x"), Some(Value::Int(1)));
593        assert_eq!(client_session.variable_value("y"), None);
594    }
595
596    #[test]
597    fn test_session_variables_transaction_commit_merges() {
598        let session = Session {
599            session_id: 123,
600            user_name: "testuser".to_string(),
601            space_name: None,
602            graph_addr: None,
603            timezone: None,
604        };
605
606        let client_session = ClientSession::new(session);
607
608        client_session.bind_transaction(linkrs_transaction::TransactionId(8));
609        client_session.set_variable("a".into(), Value::Int(5));
610
611        client_session.commit_variables();
612        client_session.unbind_transaction();
613        assert_eq!(client_session.variable_value("a"), Some(Value::Int(5)));
614    }
615
616    #[test]
617    fn test_session_variables_rollback_to_savepoint() {
618        let session = Session {
619            session_id: 123,
620            user_name: "testuser".to_string(),
621            space_name: None,
622            graph_addr: None,
623            timezone: None,
624        };
625
626        let client_session = ClientSession::new(session);
627
628        client_session.bind_transaction(linkrs_transaction::TransactionId(9));
629
630        client_session.set_variable("a".into(), Value::Int(1));
631        client_session.push_variable_savepoint("sp1");
632        client_session.set_variable("a".into(), Value::Int(2));
633        client_session.set_variable("b".into(), Value::Int(3));
634
635        assert!(
636            client_session.rollback_variables_to("sp1"),
637            "savepoint marker must be found"
638        );
639        assert_eq!(
640            client_session.variable_value("a"),
641            Some(Value::Int(1)),
642            "assignment after the savepoint is restored"
643        );
644        assert_eq!(client_session.variable_value("b"), None);
645
646        // Unknown savepoint: nothing changes.
647        client_session.set_variable("b".into(), Value::Int(4));
648        assert!(!client_session.rollback_variables_to("missing"));
649        assert_eq!(client_session.variable_value("b"), Some(Value::Int(4)));
650
651        // Released savepoints are no longer rollback targets.
652        client_session.push_variable_savepoint("sp2");
653        client_session.release_variable_savepoint("sp2");
654        assert!(!client_session.rollback_variables_to("sp2"));
655    }
656}
linkrs_server/client/client_session.rs

linkrs_server/client/
client_session.rs
