linkrs_server/client/
query_context.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Weak};
3
4use log::{info, warn};
5use parking_lot::RwLock;
6
7use linkrs_core::error::{QueryError, QueryResult};
8use linkrs_query::executor::streaming::runtime::ExecutionRuntime;
9
10#[derive(Debug)]
11pub struct QueryContext {
12    contexts: Arc<RwLock<HashMap<u32, String>>>,
13    runtime_registry: Arc<parking_lot::RwLock<HashMap<u32, Weak<ExecutionRuntime>>>>,
14}
15
16impl Default for QueryContext {
17    fn default() -> Self {
18        Self::new()
19    }
20}
21
22impl QueryContext {
23    pub fn new() -> Self {
24        Self {
25            contexts: Arc::new(RwLock::new(HashMap::new())),
26            runtime_registry: Arc::new(parking_lot::RwLock::new(HashMap::new())),
27        }
28    }
29
30    pub fn add_query(&self, ep_id: u32, query_context: String, session_id: i64) {
31        info!("Adding query {} to session {}", ep_id, session_id);
32        self.contexts.write().insert(ep_id, query_context);
33    }
34
35    /// Register a streaming query with its execution runtime for KILL QUERY support.
36    pub fn register_streaming_query(
37        &self,
38        query_id: u32,
39        query_text: String,
40        runtime: Weak<ExecutionRuntime>,
41        session_id: i64,
42    ) {
43        info!(
44            "Registering streaming query {} in session {}",
45            query_id, session_id
46        );
47        self.contexts.write().insert(query_id, query_text);
48        self.runtime_registry.write().insert(query_id, runtime);
49    }
50
51    /// Unregister a streaming query's runtime (call on completion).
52    pub fn unregister_streaming_query(&self, query_id: u32, session_id: i64) {
53        info!(
54            "Unregistering streaming query {} from session {}",
55            query_id, session_id
56        );
57        self.contexts.write().remove(&query_id);
58        self.runtime_registry.write().remove(&query_id);
59    }
60
61    pub fn delete_query(&self, ep_id: u32, session_id: i64) {
62        info!("Removing query {} from session {}", ep_id, session_id);
63        self.contexts.write().remove(&ep_id);
64    }
65
66    pub fn find_query(&self, ep_id: u32) -> bool {
67        self.contexts.read().contains_key(&ep_id)
68    }
69
70    pub fn mark_query_killed(&self, ep_id: u32, session_id: i64) {
71        info!(
72            "Marking query {} as killed in session {}",
73            ep_id, session_id
74        );
75
76        // Cancel the execution runtime if registered.
77        if let Some(weak) = self.runtime_registry.write().remove(&ep_id) {
78            if let Some(rt) = weak.upgrade() {
79                rt.cancel();
80                info!("Cancelled runtime for query {}", ep_id);
81            }
82        }
83
84        self.contexts.write().remove(&ep_id);
85    }
86
87    pub fn mark_all_queries_killed(&self, session_id: i64) {
88        let query_count = self.active_queries_count();
89        info!(
90            "Killing all {} queries in session {}",
91            query_count, session_id
92        );
93
94        // Cancel all registered runtimes.
95        for (_, weak) in self.runtime_registry.write().drain() {
96            if let Some(rt) = weak.upgrade() {
97                rt.cancel();
98            }
99        }
100
101        self.contexts.write().clear();
102    }
103
104    pub fn active_queries_count(&self) -> usize {
105        self.contexts.read().len()
106    }
107
108    pub fn kill_query(&self, query_id: u32, session_id: i64) -> QueryResult<()> {
109        info!(
110            "Attempting to kill query {} in session {}",
111            query_id, session_id
112        );
113
114        if !self.find_query(query_id) {
115            warn!("Query {} not found in session {}", query_id, session_id);
116            return Err(QueryError::execution(format!(
117                "Query not found: {}",
118                query_id
119            )));
120        }
121
122        self.mark_query_killed(query_id, session_id);
123
124        info!(
125            "Successfully killed query {} in session {}",
126            query_id, session_id
127        );
128        Ok(())
129    }
130
131    pub fn kill_multiple_queries(
132        &self,
133        query_ids: &[u32],
134        session_id: i64,
135    ) -> Vec<QueryResult<()>> {
136        query_ids
137            .iter()
138            .map(|&query_id| self.kill_query(query_id, session_id))
139            .collect()
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    use super::*;
146
147    #[test]
148    fn test_query_context() {
149        let context = QueryContext::new();
150        assert_eq!(context.active_queries_count(), 0);
151        assert!(!context.find_query(1));
152
153        context.add_query(1, "SELECT * FROM user".to_string(), 123);
154        assert_eq!(context.active_queries_count(), 1);
155        assert!(context.find_query(1));
156
157        context.delete_query(1, 123);
158        assert_eq!(context.active_queries_count(), 0);
159        assert!(!context.find_query(1));
160    }
161
162    #[test]
163    fn test_kill_query() {
164        let context = QueryContext::new();
165        context.add_query(2, "MATCH (n) RETURN n".to_string(), 123);
166        let result = context.kill_query(2, 123);
167        assert!(result.is_ok());
168        assert!(!context.find_query(2));
169
170        let result = context.kill_query(999, 123);
171        assert!(result.is_err());
172    }
173
174    #[test]
175    fn test_kill_multiple_queries() {
176        let context = QueryContext::new();
177        context.add_query(3, "query 3".to_string(), 123);
178        context.add_query(4, "query 4".to_string(), 123);
179        let results = context.kill_multiple_queries(&[3, 4, 5], 123);
180        assert_eq!(results.len(), 3);
181        assert!(results[0].is_ok());
182        assert!(results[1].is_ok());
183        assert!(results[2].is_err());
184    }
185}
linkrs_server/client/query_context.rs

linkrs_server/client/
query_context.rs
