catalog/
process_manager.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::collections::HashMap;
16use std::collections::hash_map::Entry;
17use std::fmt::{Debug, Display, Formatter};
18use std::sync::atomic::{AtomicU32, Ordering};
19use std::sync::{Arc, RwLock};
20use std::time::{Duration, Instant, UNIX_EPOCH};
21
22use api::v1::frontend::{KillProcessRequest, ListProcessRequest, ProcessInfo};
23use common_base::cancellation::CancellationHandle;
24use common_event_recorder::EventRecorderRef;
25use common_frontend::selector::{FrontendSelector, MetaClientSelector};
26use common_frontend::slow_query_event::SlowQueryEvent;
27use common_telemetry::logging::SlowQueriesRecordType;
28use common_telemetry::{debug, info, slow};
29use common_time::util::current_time_millis;
30use meta_client::MetaClientRef;
31use promql_parser::parser::EvalStmt;
32use rand::random;
33use snafu::{OptionExt, ResultExt, ensure};
34use sql::statements::statement::Statement;
35
36use crate::error;
37use crate::metrics::{PROCESS_KILL_COUNT, PROCESS_LIST_COUNT};
38
39pub type ProcessId = u32;
40pub type ProcessManagerRef = Arc<ProcessManager>;
41
42/// Query process manager.
43pub struct ProcessManager {
44    /// Local frontend server address,
45    server_addr: String,
46    /// Next process id for local queries.
47    next_id: AtomicU32,
48    /// Running process per catalog.
49    catalogs: RwLock<HashMap<String, HashMap<ProcessId, CancellableProcess>>>,
50    /// Frontend selector to locate frontend nodes.
51    frontend_selector: Option<MetaClientSelector>,
52}
53
54/// Represents a parsed query statement, functionally equivalent to [query::parser::QueryStatement].
55/// This enum is defined here to avoid cyclic dependencies with the query parser module.
56#[derive(Debug, Clone)]
57pub enum QueryStatement {
58    Sql(Statement),
59    // The optional string is the alias of the PromQL query.
60    Promql(EvalStmt, Option<String>),
61    /// Logical plan with original query string
62    Plan(String),
63}
64
65impl Display for QueryStatement {
66    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
67        match self {
68            QueryStatement::Sql(stmt) => write!(f, "{}", stmt),
69            QueryStatement::Promql(eval_stmt, alias) => {
70                if let Some(alias) = alias {
71                    write!(f, "{} AS {}", eval_stmt, alias)
72                } else {
73                    write!(f, "{}", eval_stmt)
74                }
75            }
76            QueryStatement::Plan(query) => write!(f, "{}", query),
77        }
78    }
79}
80
81impl ProcessManager {
82    /// Create a [ProcessManager] instance with server address and kv client.
83    pub fn new(server_addr: String, meta_client: Option<MetaClientRef>) -> Self {
84        let frontend_selector = meta_client.map(MetaClientSelector::new);
85        Self {
86            server_addr,
87            next_id: Default::default(),
88            catalogs: Default::default(),
89            frontend_selector,
90        }
91    }
92}
93
94impl ProcessManager {
95    /// Registers a submitted query. Use the provided id if present.
96    #[must_use]
97    pub fn register_query(
98        self: &Arc<Self>,
99        catalog: String,
100        schemas: Vec<String>,
101        query: String,
102        client: String,
103        query_id: Option<ProcessId>,
104        _slow_query_timer: Option<SlowQueryTimer>,
105    ) -> Ticket {
106        let id = query_id.unwrap_or_else(|| self.next_id.fetch_add(1, Ordering::Relaxed));
107        let process = ProcessInfo {
108            id,
109            catalog: catalog.clone(),
110            schemas,
111            query,
112            start_timestamp: current_time_millis(),
113            client,
114            frontend: self.server_addr.clone(),
115        };
116        let cancellation_handle = Arc::new(CancellationHandle::default());
117        let cancellable_process = CancellableProcess::new(cancellation_handle.clone(), process);
118
119        self.catalogs
120            .write()
121            .unwrap()
122            .entry(catalog.clone())
123            .or_default()
124            .insert(id, cancellable_process);
125
126        Ticket {
127            catalog,
128            manager: self.clone(),
129            id,
130            cancellation_handle,
131            _slow_query_timer,
132        }
133    }
134
135    /// Generates the next process id.
136    pub fn next_id(&self) -> u32 {
137        self.next_id.fetch_add(1, Ordering::Relaxed)
138    }
139
140    /// De-register a query from process list.
141    pub fn deregister_query(&self, catalog: String, id: ProcessId) {
142        if let Entry::Occupied(mut o) = self.catalogs.write().unwrap().entry(catalog) {
143            let process = o.get_mut().remove(&id);
144            debug!("Deregister process: {:?}", process);
145            if o.get().is_empty() {
146                o.remove();
147            }
148        }
149    }
150
151    /// List local running processes in given catalog.
152    pub fn local_processes(&self, catalog: Option<&str>) -> error::Result<Vec<ProcessInfo>> {
153        let catalogs = self.catalogs.read().unwrap();
154        let result = if let Some(catalog) = catalog {
155            if let Some(catalogs) = catalogs.get(catalog) {
156                catalogs.values().map(|p| p.process.clone()).collect()
157            } else {
158                vec![]
159            }
160        } else {
161            catalogs
162                .values()
163                .flat_map(|v| v.values().map(|p| p.process.clone()))
164                .collect()
165        };
166        Ok(result)
167    }
168
169    /// Lists running processes from every selected frontend and the local registry.
170    ///
171    /// Returns an error if a remote frontend cannot provide its process list.
172    pub async fn list_all_processes(
173        &self,
174        catalog: Option<&str>,
175    ) -> error::Result<Vec<ProcessInfo>> {
176        let mut processes = vec![];
177        if let Some(remote_frontend_selector) = self.frontend_selector.as_ref() {
178            let frontends = remote_frontend_selector
179                .select(|peer| peer.addr != self.server_addr)
180                .await
181                .context(error::InvokeFrontendSnafu)?;
182            for mut f in frontends {
183                let resp = f
184                    .list_process(ListProcessRequest {
185                        catalog: catalog.unwrap_or_default().to_string(),
186                    })
187                    .await
188                    .context(error::InvokeFrontendSnafu)?;
189                processes.extend(resp.processes);
190            }
191        }
192        processes.extend(self.local_processes(catalog)?);
193        Ok(processes)
194    }
195
196    /// Kills query with provided catalog and id.
197    pub async fn kill_process(
198        &self,
199        server_addr: String,
200        catalog: String,
201        id: ProcessId,
202    ) -> error::Result<bool> {
203        if server_addr == self.server_addr {
204            self.kill_local_process(catalog, id).await
205        } else {
206            let mut nodes = self
207                .frontend_selector
208                .as_ref()
209                .context(error::MetaClientMissingSnafu)?
210                .select(|peer| peer.addr == server_addr)
211                .await
212                .context(error::InvokeFrontendSnafu)?;
213            ensure!(
214                !nodes.is_empty(),
215                error::FrontendNotFoundSnafu { addr: server_addr }
216            );
217
218            let request = KillProcessRequest {
219                server_addr,
220                catalog,
221                process_id: id,
222            };
223            nodes[0]
224                .kill_process(request)
225                .await
226                .context(error::InvokeFrontendSnafu)?;
227            Ok(true)
228        }
229    }
230
231    /// Kills local query with provided catalog and id.
232    pub async fn kill_local_process(&self, catalog: String, id: ProcessId) -> error::Result<bool> {
233        if let Some(catalogs) = self.catalogs.write().unwrap().get_mut(&catalog) {
234            if let Some(process) = catalogs.remove(&id) {
235                process.handle.cancel();
236                info!(
237                    "Killed process, catalog: {}, id: {:?}",
238                    process.process.catalog, process.process.id
239                );
240                PROCESS_KILL_COUNT.with_label_values(&[&catalog]).inc();
241                Ok(true)
242            } else {
243                debug!("Failed to kill process, id not found: {}", id);
244                Ok(false)
245            }
246        } else {
247            debug!("Failed to kill process, catalog not found: {}", catalog);
248            Ok(false)
249        }
250    }
251}
252
253pub struct Ticket {
254    pub(crate) catalog: String,
255    pub(crate) manager: ProcessManagerRef,
256    pub(crate) id: ProcessId,
257    pub cancellation_handle: Arc<CancellationHandle>,
258
259    // Keep the handle of the slow query timer to ensure it will trigger the event recording when dropped.
260    _slow_query_timer: Option<SlowQueryTimer>,
261}
262
263impl Drop for Ticket {
264    fn drop(&mut self) {
265        self.manager
266            .deregister_query(std::mem::take(&mut self.catalog), self.id);
267    }
268}
269
270struct CancellableProcess {
271    handle: Arc<CancellationHandle>,
272    process: ProcessInfo,
273}
274
275impl Drop for CancellableProcess {
276    fn drop(&mut self) {
277        PROCESS_LIST_COUNT
278            .with_label_values(&[&self.process.catalog])
279            .dec();
280    }
281}
282
283impl CancellableProcess {
284    fn new(handle: Arc<CancellationHandle>, process: ProcessInfo) -> Self {
285        PROCESS_LIST_COUNT
286            .with_label_values(&[&process.catalog])
287            .inc();
288        Self { handle, process }
289    }
290}
291
292impl Debug for CancellableProcess {
293    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
294        f.debug_struct("CancellableProcess")
295            .field("cancelled", &self.handle.is_cancelled())
296            .field("process", &self.process)
297            .finish()
298    }
299}
300
301/// SlowQueryTimer is used to log slow query when it's dropped.
302/// In drop(), it will check if the query is slow and send the slow query event to the handler.
303pub struct SlowQueryTimer {
304    start: Instant,
305    stmt: QueryStatement,
306    schema_name: String,
307    threshold: Duration,
308    sample_ratio: f64,
309    record_type: SlowQueriesRecordType,
310    recorder: EventRecorderRef,
311    record_state: Arc<RwLock<SlowQueryRecordState>>,
312}
313
314#[derive(Default)]
315struct SlowQueryRecordState {
316    force_record: bool,
317    payload: serde_json::Value,
318}
319
320/// A handle to attach diagnostics to a slow query and force it to be recorded.
321#[derive(Clone)]
322pub struct SlowQueryRecorder {
323    record_state: Arc<RwLock<SlowQueryRecordState>>,
324}
325
326impl SlowQueryRecorder {
327    /// Attaches the payload and marks the query to bypass the threshold and sampling checks.
328    pub fn force_record_with_payload(&self, payload: serde_json::Value) {
329        let mut state = self.record_state.write().unwrap();
330        state.payload = payload;
331        state.force_record = true;
332    }
333}
334
335impl SlowQueryTimer {
336    pub fn new(
337        stmt: QueryStatement,
338        schema_name: String,
339        threshold: Duration,
340        sample_ratio: f64,
341        record_type: SlowQueriesRecordType,
342        recorder: EventRecorderRef,
343    ) -> Self {
344        Self {
345            start: Instant::now(),
346            stmt,
347            schema_name,
348            threshold,
349            sample_ratio,
350            record_type,
351            recorder,
352            record_state: Arc::default(),
353        }
354    }
355
356    /// Returns a handle that can attach diagnostics and force this query to be recorded.
357    pub fn recorder(&self) -> SlowQueryRecorder {
358        SlowQueryRecorder {
359            record_state: self.record_state.clone(),
360        }
361    }
362}
363
364impl SlowQueryTimer {
365    fn send_slow_query_event(&self, elapsed: Duration, payload: serde_json::Value) {
366        let mut slow_query_event = SlowQueryEvent {
367            cost: elapsed.as_millis() as u64,
368            threshold: self.threshold.as_millis() as u64,
369            query: "".to_string(),
370            schema_name: self.schema_name.clone(),
371
372            // The following fields are only used for PromQL queries.
373            is_promql: false,
374            promql_range: None,
375            promql_step: None,
376            promql_start: None,
377            promql_end: None,
378            payload,
379        };
380
381        match &self.stmt {
382            QueryStatement::Promql(stmt, _alias) => {
383                slow_query_event.is_promql = true;
384                slow_query_event.query = self.stmt.to_string();
385                slow_query_event.promql_step = Some(stmt.interval.as_millis() as u64);
386
387                let start = stmt
388                    .start
389                    .duration_since(UNIX_EPOCH)
390                    .unwrap_or_default()
391                    .as_millis() as i64;
392
393                let end = stmt
394                    .end
395                    .duration_since(UNIX_EPOCH)
396                    .unwrap_or_default()
397                    .as_millis() as i64;
398
399                slow_query_event.promql_range = Some((end - start) as u64);
400                slow_query_event.promql_start = Some(start);
401                slow_query_event.promql_end = Some(end);
402            }
403            QueryStatement::Sql(stmt) => {
404                slow_query_event.query = stmt.to_string();
405            }
406            QueryStatement::Plan(query) => {
407                slow_query_event.query = query.clone();
408            }
409        }
410
411        match self.record_type {
412            // Send the slow query event to the event recorder to persist it as the system table.
413            SlowQueriesRecordType::SystemTable => {
414                self.recorder.record(Box::new(slow_query_event));
415            }
416            // Record the slow query through its dedicated logger.
417            SlowQueriesRecordType::Log => {
418                slow!(
419                    cost = slow_query_event.cost,
420                    threshold = slow_query_event.threshold,
421                    query = slow_query_event.query,
422                    schema_name = slow_query_event.schema_name,
423                    is_promql = slow_query_event.is_promql,
424                    promql_range = slow_query_event.promql_range,
425                    promql_step = slow_query_event.promql_step,
426                    promql_start = slow_query_event.promql_start,
427                    promql_end = slow_query_event.promql_end,
428                    payload = slow_query_event.payload.to_string(),
429                );
430            }
431        }
432    }
433}
434
435impl Drop for SlowQueryTimer {
436    fn drop(&mut self) {
437        let elapsed = self.start.elapsed();
438        let (force_record, payload) = {
439            let state = self.record_state.read().unwrap();
440            (state.force_record, state.payload.clone())
441        };
442
443        if force_record
444            || (elapsed > self.threshold
445                && (self.sample_ratio >= 1.0 || random::<f64>() <= self.sample_ratio))
446        {
447            self.send_slow_query_event(elapsed, payload);
448        }
449    }
450}
451
452#[cfg(test)]
453mod tests {
454    use std::sync::{Arc, Mutex};
455    use std::time::Duration;
456
457    use common_event_recorder::{Event, EventRecorder, EventTypeFilter, EventTypeFilterRef};
458    use common_frontend::slow_query_event::SlowQueryEvent;
459    use common_telemetry::logging::SlowQueriesRecordType;
460    use serde_json::Value;
461
462    use crate::process_manager::{ProcessManager, QueryStatement, SlowQueryTimer};
463
464    #[derive(Debug, Default)]
465    struct RecordingEventRecorder {
466        events: Mutex<Vec<(String, Value)>>,
467    }
468
469    impl EventRecorder for RecordingEventRecorder {
470        fn record(&self, event: Box<dyn Event>) {
471            let event = event
472                .as_any()
473                .downcast_ref::<SlowQueryEvent>()
474                .expect("expected a slow query event");
475            self.events
476                .lock()
477                .unwrap()
478                .push((event.query.clone(), event.payload.clone()));
479        }
480
481        fn event_type_filter(&self) -> EventTypeFilterRef {
482            Arc::new(EventTypeFilter::All)
483        }
484
485        fn close(&self) {}
486    }
487
488    #[test]
489    fn test_forced_slow_query_bypasses_threshold_and_sampling() {
490        let event_recorder = Arc::new(RecordingEventRecorder::default());
491        let timer = SlowQueryTimer::new(
492            QueryStatement::Plan("EXPLAIN ANALYZE VERBOSE SELECT 1".to_string()),
493            "public".to_string(),
494            Duration::from_secs(3600),
495            0.0,
496            SlowQueriesRecordType::SystemTable,
497            event_recorder.clone(),
498        );
499        let payload = serde_json::json!({
500            "timed_out": true,
501            "metrics": [{"stage": 0}],
502        });
503        timer.recorder().force_record_with_payload(payload.clone());
504
505        drop(timer);
506
507        let events = event_recorder.events.lock().unwrap();
508        assert_eq!(events.len(), 1);
509        assert_eq!(events[0].0, "EXPLAIN ANALYZE VERBOSE SELECT 1");
510        assert_eq!(events[0].1, payload);
511    }
512
513    #[test]
514    fn test_unforced_fast_query_is_not_recorded() {
515        let event_recorder = Arc::new(RecordingEventRecorder::default());
516        let timer = SlowQueryTimer::new(
517            QueryStatement::Plan("SELECT 1".to_string()),
518            "public".to_string(),
519            Duration::from_secs(3600),
520            0.0,
521            SlowQueriesRecordType::SystemTable,
522            event_recorder.clone(),
523        );
524
525        drop(timer);
526
527        assert!(event_recorder.events.lock().unwrap().is_empty());
528    }
529
530    #[tokio::test]
531    async fn test_register_query() {
532        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
533        let ticket = process_manager.clone().register_query(
534            "public".to_string(),
535            vec!["test".to_string()],
536            "SELECT * FROM table".to_string(),
537            "".to_string(),
538            None,
539            None,
540        );
541
542        let running_processes = process_manager.local_processes(None).unwrap();
543        assert_eq!(running_processes.len(), 1);
544        assert_eq!(&running_processes[0].frontend, "127.0.0.1:8000");
545        assert_eq!(running_processes[0].id, ticket.id);
546        assert_eq!(&running_processes[0].query, "SELECT * FROM table");
547
548        drop(ticket);
549        assert_eq!(process_manager.local_processes(None).unwrap().len(), 0);
550    }
551
552    #[tokio::test]
553    async fn test_register_query_with_custom_id() {
554        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
555        let custom_id = 12345;
556
557        let ticket = process_manager.clone().register_query(
558            "public".to_string(),
559            vec!["test".to_string()],
560            "SELECT * FROM table".to_string(),
561            "client1".to_string(),
562            Some(custom_id),
563            None,
564        );
565
566        assert_eq!(ticket.id, custom_id);
567
568        let running_processes = process_manager.local_processes(None).unwrap();
569        assert_eq!(running_processes.len(), 1);
570        assert_eq!(running_processes[0].id, custom_id);
571        assert_eq!(&running_processes[0].client, "client1");
572    }
573
574    #[tokio::test]
575    async fn test_multiple_queries_same_catalog() {
576        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
577
578        let ticket1 = process_manager.clone().register_query(
579            "public".to_string(),
580            vec!["schema1".to_string()],
581            "SELECT * FROM table1".to_string(),
582            "client1".to_string(),
583            None,
584            None,
585        );
586
587        let ticket2 = process_manager.clone().register_query(
588            "public".to_string(),
589            vec!["schema2".to_string()],
590            "SELECT * FROM table2".to_string(),
591            "client2".to_string(),
592            None,
593            None,
594        );
595
596        let running_processes = process_manager.local_processes(Some("public")).unwrap();
597        assert_eq!(running_processes.len(), 2);
598
599        // Verify both processes are present
600        let ids: Vec<u32> = running_processes.iter().map(|p| p.id).collect();
601        assert!(ids.contains(&ticket1.id));
602        assert!(ids.contains(&ticket2.id));
603    }
604
605    #[tokio::test]
606    async fn test_multiple_catalogs() {
607        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
608
609        let _ticket1 = process_manager.clone().register_query(
610            "catalog1".to_string(),
611            vec!["schema1".to_string()],
612            "SELECT * FROM table1".to_string(),
613            "client1".to_string(),
614            None,
615            None,
616        );
617
618        let _ticket2 = process_manager.clone().register_query(
619            "catalog2".to_string(),
620            vec!["schema2".to_string()],
621            "SELECT * FROM table2".to_string(),
622            "client2".to_string(),
623            None,
624            None,
625        );
626
627        // Test listing processes for specific catalog
628        let catalog1_processes = process_manager.local_processes(Some("catalog1")).unwrap();
629        assert_eq!(catalog1_processes.len(), 1);
630        assert_eq!(&catalog1_processes[0].catalog, "catalog1");
631
632        let catalog2_processes = process_manager.local_processes(Some("catalog2")).unwrap();
633        assert_eq!(catalog2_processes.len(), 1);
634        assert_eq!(&catalog2_processes[0].catalog, "catalog2");
635
636        // Test listing all processes
637        let all_processes = process_manager.local_processes(None).unwrap();
638        assert_eq!(all_processes.len(), 2);
639    }
640
641    #[tokio::test]
642    async fn test_deregister_query() {
643        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
644
645        let ticket = process_manager.clone().register_query(
646            "public".to_string(),
647            vec!["test".to_string()],
648            "SELECT * FROM table".to_string(),
649            "client1".to_string(),
650            None,
651            None,
652        );
653        assert_eq!(process_manager.local_processes(None).unwrap().len(), 1);
654        process_manager.deregister_query("public".to_string(), ticket.id);
655        assert_eq!(process_manager.local_processes(None).unwrap().len(), 0);
656    }
657
658    #[tokio::test]
659    async fn test_cancellation_handle() {
660        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
661
662        let ticket = process_manager.clone().register_query(
663            "public".to_string(),
664            vec!["test".to_string()],
665            "SELECT * FROM table".to_string(),
666            "client1".to_string(),
667            None,
668            None,
669        );
670
671        assert!(!ticket.cancellation_handle.is_cancelled());
672        ticket.cancellation_handle.cancel();
673        assert!(ticket.cancellation_handle.is_cancelled());
674    }
675
676    #[tokio::test]
677    async fn test_kill_local_process() {
678        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
679
680        let ticket = process_manager.clone().register_query(
681            "public".to_string(),
682            vec!["test".to_string()],
683            "SELECT * FROM table".to_string(),
684            "client1".to_string(),
685            None,
686            None,
687        );
688        assert!(!ticket.cancellation_handle.is_cancelled());
689        let killed = process_manager
690            .kill_process(
691                "127.0.0.1:8000".to_string(),
692                "public".to_string(),
693                ticket.id,
694            )
695            .await
696            .unwrap();
697
698        assert!(killed);
699        assert_eq!(process_manager.local_processes(None).unwrap().len(), 0);
700    }
701
702    #[tokio::test]
703    async fn test_kill_nonexistent_process() {
704        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
705        let killed = process_manager
706            .kill_process("127.0.0.1:8000".to_string(), "public".to_string(), 999)
707            .await
708            .unwrap();
709        assert!(!killed);
710    }
711
712    #[tokio::test]
713    async fn test_kill_process_nonexistent_catalog() {
714        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
715        let killed = process_manager
716            .kill_process("127.0.0.1:8000".to_string(), "nonexistent".to_string(), 1)
717            .await
718            .unwrap();
719        assert!(!killed);
720    }
721
722    #[tokio::test]
723    async fn test_process_info_fields() {
724        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
725
726        let _ticket = process_manager.clone().register_query(
727            "test_catalog".to_string(),
728            vec!["schema1".to_string(), "schema2".to_string()],
729            "SELECT COUNT(*) FROM users WHERE age > 18".to_string(),
730            "test_client".to_string(),
731            Some(42),
732            None,
733        );
734
735        let processes = process_manager.local_processes(None).unwrap();
736        assert_eq!(processes.len(), 1);
737
738        let process = &processes[0];
739        assert_eq!(process.id, 42);
740        assert_eq!(&process.catalog, "test_catalog");
741        assert_eq!(process.schemas, vec!["schema1", "schema2"]);
742        assert_eq!(&process.query, "SELECT COUNT(*) FROM users WHERE age > 18");
743        assert_eq!(&process.client, "test_client");
744        assert_eq!(&process.frontend, "127.0.0.1:8000");
745        assert!(process.start_timestamp > 0);
746    }
747
748    #[tokio::test]
749    async fn test_ticket_drop_deregisters_process() {
750        let process_manager = Arc::new(ProcessManager::new("127.0.0.1:8000".to_string(), None));
751
752        {
753            let _ticket = process_manager.clone().register_query(
754                "public".to_string(),
755                vec!["test".to_string()],
756                "SELECT * FROM table".to_string(),
757                "client1".to_string(),
758                None,
759                None,
760            );
761
762            // Process should be registered
763            assert_eq!(process_manager.local_processes(None).unwrap().len(), 1);
764        } // ticket goes out of scope here
765
766        // Process should be automatically deregistered
767        assert_eq!(process_manager.local_processes(None).unwrap().len(), 0);
768    }
769}
catalog/process_manager.rs

catalog/
process_manager.rs
