session/
context.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::fmt::{Display, Formatter};
17use std::net::SocketAddr;
18use std::sync::{Arc, RwLock};
19use std::time::Duration;
20
21use api::v1::ExplainOptions;
22use api::v1::region::RegionRequestHeader;
23use arc_swap::ArcSwap;
24use auth::UserInfoRef;
25pub use common_base::protocol::Channel;
26use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
27use common_catalog::{build_db_string, parse_catalog_and_schema_from_db_string};
28use common_recordbatch::cursor::RecordBatchStreamCursor;
29use common_telemetry::warn;
30use common_time::Timezone;
31use common_time::timezone::parse_timezone;
32use datafusion_common::config::ConfigOptions;
33use derive_builder::Builder;
34use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
35
36pub use crate::hints::{
37    LIVE_ANALYZE_METRICS_EXTENSION_KEY, REMOTE_QUERY_ID_EXTENSION_KEY,
38    SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
39};
40use crate::protocol_ctx::ProtocolCtx;
41use crate::query_id::QueryId;
42use crate::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
43use crate::{MutableInner, ReadPreference};
44
45pub type QueryContextRef = Arc<QueryContext>;
46pub type ConnInfoRef = Arc<ConnInfo>;
47
48pub const FLIGHT_METRICS_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1);
49
50const CURSOR_COUNT_WARNING_LIMIT: usize = 10;
51
52pub fn generate_remote_query_id() -> String {
53    generate_remote_query_id_value().to_string()
54}
55
56pub fn generate_remote_query_id_value() -> QueryId {
57    QueryId::new()
58}
59
60#[derive(Debug, Builder, Clone)]
61#[builder(pattern = "owned")]
62#[builder(build_fn(skip))]
63pub struct QueryContext {
64    current_catalog: String,
65    /// mapping of RegionId to SequenceNumber, for snapshot read, meaning that the read should only
66    /// container data that was committed before(and include) the given sequence number
67    /// this field will only be filled if extensions contains a pair of "snapshot_read" and "true"
68    snapshot_seqs: Arc<RwLock<HashMap<u64, u64>>>,
69    /// Mappings of the RegionId to the minimal sequence of SST file to scan.
70    sst_min_sequences: Arc<RwLock<HashMap<u64, u64>>>,
71    // we use Arc<RwLock>> for modifiable fields
72    #[builder(default)]
73    mutable_session_data: Arc<RwLock<MutableInner>>,
74    #[builder(default)]
75    mutable_query_context_data: Arc<RwLock<QueryContextMutableFields>>,
76    sql_dialect: Arc<dyn Dialect + Send + Sync>,
77    #[builder(default)]
78    extensions: HashMap<String, String>,
79    /// The configuration parameter are used to store the parameters that are set by the user
80    #[builder(default)]
81    configuration_parameter: Arc<ConfigurationVariables>,
82    /// Track which protocol the query comes from.
83    #[builder(default)]
84    channel: Channel,
85    /// Process id for managing on-going queries
86    #[builder(default)]
87    process_id: u32,
88    /// Connection information
89    #[builder(default)]
90    conn_info: ConnInfo,
91    /// Protocol specific context
92    #[builder(default)]
93    protocol_ctx: ProtocolCtx,
94}
95
96/// This fields hold data that is only valid to current query context
97#[derive(Debug, Builder, Clone, Default)]
98pub struct QueryContextMutableFields {
99    warning: Option<String>,
100    // TODO: remove this when format is supported in datafusion
101    explain_format: Option<String>,
102    /// Explain options to control the verbose analyze output.
103    explain_options: Option<ExplainOptions>,
104}
105
106impl Display for QueryContext {
107    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
108        write!(
109            f,
110            "QueryContext{{catalog: {}, schema: {}}}",
111            self.current_catalog(),
112            self.current_schema()
113        )
114    }
115}
116
117impl QueryContextBuilder {
118    pub fn current_schema(mut self, schema: String) -> Self {
119        if self.mutable_session_data.is_none() {
120            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
121        }
122
123        // safe for unwrap because previous none check
124        self.mutable_session_data
125            .as_mut()
126            .unwrap()
127            .write()
128            .unwrap()
129            .schema = schema;
130        self
131    }
132
133    pub fn timezone(mut self, timezone: Timezone) -> Self {
134        if self.mutable_session_data.is_none() {
135            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
136        }
137
138        self.mutable_session_data
139            .as_mut()
140            .unwrap()
141            .write()
142            .unwrap()
143            .timezone = timezone;
144        self
145    }
146
147    pub fn explain_options(mut self, explain_options: Option<ExplainOptions>) -> Self {
148        self.mutable_query_context_data
149            .get_or_insert_default()
150            .write()
151            .unwrap()
152            .explain_options = explain_options;
153        self
154    }
155
156    pub fn read_preference(mut self, read_preference: ReadPreference) -> Self {
157        self.mutable_session_data
158            .get_or_insert_default()
159            .write()
160            .unwrap()
161            .read_preference = read_preference;
162        self
163    }
164}
165
166impl From<&RegionRequestHeader> for QueryContext {
167    fn from(value: &RegionRequestHeader) -> Self {
168        if let Some(ctx) = &value.query_context {
169            ctx.clone().into()
170        } else {
171            QueryContextBuilder::default()
172                .set_extension(
173                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
174                    generate_remote_query_id(),
175                )
176                .build()
177        }
178    }
179}
180
181impl From<api::v1::QueryContext> for QueryContext {
182    fn from(ctx: api::v1::QueryContext) -> Self {
183        let sequences = ctx.snapshot_seqs.as_ref();
184        QueryContextBuilder::default()
185            .current_catalog(ctx.current_catalog)
186            .current_schema(ctx.current_schema)
187            .timezone(parse_timezone(Some(&ctx.timezone)))
188            .extensions(ctx.extensions)
189            .channel(ctx.channel.into())
190            .snapshot_seqs(Arc::new(RwLock::new(
191                sequences
192                    .map(|x| x.snapshot_seqs.clone())
193                    .unwrap_or_default(),
194            )))
195            .sst_min_sequences(Arc::new(RwLock::new(
196                sequences
197                    .map(|x| x.sst_min_sequences.clone())
198                    .unwrap_or_default(),
199            )))
200            .explain_options(ctx.explain)
201            .build()
202    }
203}
204
205impl From<QueryContext> for api::v1::QueryContext {
206    fn from(
207        QueryContext {
208            current_catalog,
209            mutable_session_data: mutable_inner,
210            extensions,
211            channel,
212            snapshot_seqs,
213            sst_min_sequences,
214            mutable_query_context_data,
215            ..
216        }: QueryContext,
217    ) -> Self {
218        let explain = mutable_query_context_data.read().unwrap().explain_options;
219        let mutable_inner = mutable_inner.read().unwrap();
220        api::v1::QueryContext {
221            current_catalog,
222            current_schema: mutable_inner.schema.clone(),
223            timezone: mutable_inner.timezone.to_string(),
224            extensions,
225            channel: channel as u32,
226            snapshot_seqs: Some(api::v1::SnapshotSequences {
227                snapshot_seqs: snapshot_seqs.read().unwrap().clone(),
228                sst_min_sequences: sst_min_sequences.read().unwrap().clone(),
229            }),
230            explain,
231        }
232    }
233}
234
235impl From<&QueryContext> for api::v1::QueryContext {
236    fn from(ctx: &QueryContext) -> Self {
237        ctx.clone().into()
238    }
239}
240
241impl QueryContext {
242    /// Forks this context with an independent snapshot of mutable session data.
243    ///
244    /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields
245    /// held in mutable session data do not affect this context.
246    pub fn fork(&self) -> Self {
247        let mut fork = self.clone();
248        fork.mutable_session_data = Arc::new(RwLock::new(
249            self.mutable_session_data.read().unwrap().clone(),
250        ));
251        fork
252    }
253
254    pub fn arc() -> QueryContextRef {
255        Arc::new(
256            QueryContextBuilder::default()
257                .set_extension(
258                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
259                    generate_remote_query_id(),
260                )
261                .build(),
262        )
263    }
264
265    /// Create a new  datafusion's ConfigOptions instance based on the current QueryContext.
266    pub fn create_config_options(&self) -> ConfigOptions {
267        let mut config = ConfigOptions::default();
268        config.execution.time_zone = Some(self.timezone().to_string());
269        config
270    }
271
272    pub fn with(catalog: &str, schema: &str) -> QueryContext {
273        QueryContextBuilder::default()
274            .current_catalog(catalog.to_string())
275            .current_schema(schema.to_string())
276            .set_extension(
277                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
278                generate_remote_query_id(),
279            )
280            .build()
281    }
282
283    pub fn with_channel(catalog: &str, schema: &str, channel: Channel) -> QueryContext {
284        QueryContextBuilder::default()
285            .current_catalog(catalog.to_string())
286            .current_schema(schema.to_string())
287            .channel(channel)
288            .set_extension(
289                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
290                generate_remote_query_id(),
291            )
292            .build()
293    }
294
295    pub fn with_db_name(db_name: Option<&str>) -> QueryContext {
296        let (catalog, schema) = db_name
297            .map(|db| {
298                let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
299                (catalog, schema)
300            })
301            .unwrap_or_else(|| {
302                (
303                    DEFAULT_CATALOG_NAME.to_string(),
304                    DEFAULT_SCHEMA_NAME.to_string(),
305                )
306            });
307        QueryContextBuilder::default()
308            .current_catalog(catalog)
309            .current_schema(schema.clone())
310            .set_extension(
311                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
312                generate_remote_query_id(),
313            )
314            .build()
315    }
316
317    pub fn current_schema(&self) -> String {
318        self.mutable_session_data.read().unwrap().schema.clone()
319    }
320
321    pub fn set_current_schema(&self, new_schema: &str) {
322        self.mutable_session_data.write().unwrap().schema = new_schema.to_string();
323    }
324
325    pub fn current_catalog(&self) -> &str {
326        &self.current_catalog
327    }
328
329    pub fn set_current_catalog(&mut self, new_catalog: &str) {
330        self.current_catalog = new_catalog.to_string();
331    }
332
333    pub fn sql_dialect(&self) -> &(dyn Dialect + Send + Sync) {
334        &*self.sql_dialect
335    }
336
337    pub fn get_db_string(&self) -> String {
338        let catalog = self.current_catalog();
339        let schema = self.current_schema();
340        build_db_string(catalog, &schema)
341    }
342
343    pub fn timezone(&self) -> Timezone {
344        self.mutable_session_data.read().unwrap().timezone.clone()
345    }
346
347    pub fn set_timezone(&self, timezone: Timezone) {
348        self.mutable_session_data.write().unwrap().timezone = timezone;
349    }
350
351    pub fn read_preference(&self) -> ReadPreference {
352        self.mutable_session_data.read().unwrap().read_preference
353    }
354
355    pub fn set_read_preference(&self, read_preference: ReadPreference) {
356        self.mutable_session_data.write().unwrap().read_preference = read_preference;
357    }
358
359    pub fn current_user(&self) -> UserInfoRef {
360        self.mutable_session_data.read().unwrap().user_info.clone()
361    }
362
363    pub fn set_current_user(&self, user: UserInfoRef) {
364        self.mutable_session_data.write().unwrap().user_info = user;
365    }
366
367    pub fn set_extension<S1: Into<String>, S2: Into<String>>(&mut self, key: S1, value: S2) {
368        self.extensions.insert(key.into(), value.into());
369    }
370
371    pub fn extension<S: AsRef<str>>(&self, key: S) -> Option<&str> {
372        self.extensions.get(key.as_ref()).map(|v| v.as_str())
373    }
374
375    pub fn remote_query_id(&self) -> Option<&str> {
376        self.extension(REMOTE_QUERY_ID_EXTENSION_KEY)
377    }
378
379    pub fn remote_query_id_value(&self) -> Option<QueryId> {
380        self.remote_query_id()
381            .and_then(|query_id| query_id.parse().ok())
382    }
383
384    pub fn enable_live_analyze_metrics(&mut self) {
385        if let Some(remote_query_id) = self.remote_query_id().map(str::to_string) {
386            self.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, remote_query_id);
387        }
388    }
389
390    pub fn live_analyze_metrics_enabled(&self) -> bool {
391        self.remote_query_id()
392            .zip(self.extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY))
393            .is_some_and(|(remote_query_id, value)| value == remote_query_id)
394    }
395
396    pub fn extensions(&self) -> HashMap<String, String> {
397        self.extensions.clone()
398    }
399
400    /// Default to double quote and fallback to back quote
401    pub fn quote_style(&self) -> char {
402        if self.sql_dialect().is_delimited_identifier_start('"') {
403            '"'
404        } else if self.sql_dialect().is_delimited_identifier_start('\'') {
405            '\''
406        } else {
407            '`'
408        }
409    }
410
411    pub fn configuration_parameter(&self) -> &ConfigurationVariables {
412        &self.configuration_parameter
413    }
414
415    pub fn channel(&self) -> Channel {
416        self.channel
417    }
418
419    pub fn set_channel(&mut self, channel: Channel) {
420        self.channel = channel;
421    }
422
423    pub fn warning(&self) -> Option<String> {
424        self.mutable_query_context_data
425            .read()
426            .unwrap()
427            .warning
428            .clone()
429    }
430
431    pub fn set_warning(&self, msg: String) {
432        self.mutable_query_context_data.write().unwrap().warning = Some(msg);
433    }
434
435    pub fn explain_format(&self) -> Option<String> {
436        self.mutable_query_context_data
437            .read()
438            .unwrap()
439            .explain_format
440            .clone()
441    }
442
443    pub fn set_explain_format(&self, format: String) {
444        self.mutable_query_context_data
445            .write()
446            .unwrap()
447            .explain_format = Some(format);
448    }
449
450    pub fn explain_verbose(&self) -> bool {
451        self.mutable_query_context_data
452            .read()
453            .unwrap()
454            .explain_options
455            .map(|opts| opts.verbose)
456            .unwrap_or(false)
457    }
458
459    pub fn set_explain_verbose(&self, verbose: bool) {
460        self.mutable_query_context_data
461            .write()
462            .unwrap()
463            .explain_options
464            .get_or_insert_default()
465            .verbose = verbose;
466    }
467
468    pub fn query_timeout(&self) -> Option<Duration> {
469        self.mutable_session_data.read().unwrap().query_timeout
470    }
471
472    pub fn query_timeout_as_millis(&self) -> u128 {
473        let timeout = self.mutable_session_data.read().unwrap().query_timeout;
474        if let Some(t) = timeout {
475            return t.as_millis();
476        }
477        0
478    }
479
480    pub fn set_query_timeout(&self, timeout: Duration) {
481        self.mutable_session_data.write().unwrap().query_timeout = Some(timeout);
482    }
483
484    pub fn insert_cursor(&self, name: String, rb: RecordBatchStreamCursor) {
485        let mut guard = self.mutable_session_data.write().unwrap();
486        guard.cursors.insert(name, Arc::new(rb));
487
488        let cursor_count = guard.cursors.len();
489        if cursor_count > CURSOR_COUNT_WARNING_LIMIT {
490            warn!("Current connection has {} open cursors", cursor_count);
491        }
492    }
493
494    pub fn remove_cursor(&self, name: &str) {
495        let mut guard = self.mutable_session_data.write().unwrap();
496        guard.cursors.remove(name);
497    }
498
499    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
500        let guard = self.mutable_session_data.read().unwrap();
501        let rb = guard.cursors.get(name);
502        rb.cloned()
503    }
504
505    pub fn snapshots(&self) -> HashMap<u64, u64> {
506        self.snapshot_seqs.read().unwrap().clone()
507    }
508
509    pub fn sst_min_sequences(&self) -> HashMap<u64, u64> {
510        self.sst_min_sequences.read().unwrap().clone()
511    }
512
513    pub fn get_snapshot(&self, region_id: u64) -> Option<u64> {
514        self.snapshot_seqs.read().unwrap().get(&region_id).cloned()
515    }
516
517    pub fn set_snapshot(&self, region_id: u64, sequence: u64) {
518        self.snapshot_seqs
519            .write()
520            .unwrap()
521            .insert(region_id, sequence);
522    }
523
524    /// Returns `true` if the session can cast strings to numbers in MySQL style.
525    pub fn auto_string_to_numeric(&self) -> bool {
526        matches!(self.channel, Channel::Mysql)
527    }
528
529    /// Finds the minimal sequence of SST files to scan of a Region.
530    pub fn sst_min_sequence(&self, region_id: u64) -> Option<u64> {
531        self.sst_min_sequences
532            .read()
533            .unwrap()
534            .get(&region_id)
535            .copied()
536    }
537
538    pub fn process_id(&self) -> u32 {
539        self.process_id
540    }
541
542    /// Get client information
543    pub fn conn_info(&self) -> &ConnInfo {
544        &self.conn_info
545    }
546
547    pub fn protocol_ctx(&self) -> &ProtocolCtx {
548        &self.protocol_ctx
549    }
550
551    pub fn set_protocol_ctx(&mut self, protocol_ctx: ProtocolCtx) {
552        self.protocol_ctx = protocol_ctx;
553    }
554}
555
556impl QueryContextBuilder {
557    pub fn build(self) -> QueryContext {
558        let channel = self.channel.unwrap_or_default();
559        let mut extensions = self.extensions.unwrap_or_default();
560        extensions
561            .entry(REMOTE_QUERY_ID_EXTENSION_KEY.to_string())
562            .or_insert_with(generate_remote_query_id);
563        QueryContext {
564            current_catalog: self
565                .current_catalog
566                .unwrap_or_else(|| DEFAULT_CATALOG_NAME.to_string()),
567            snapshot_seqs: self.snapshot_seqs.unwrap_or_default(),
568            sst_min_sequences: self.sst_min_sequences.unwrap_or_default(),
569            mutable_session_data: self.mutable_session_data.unwrap_or_default(),
570            mutable_query_context_data: self.mutable_query_context_data.unwrap_or_default(),
571            sql_dialect: self
572                .sql_dialect
573                .unwrap_or_else(|| Arc::new(GreptimeDbDialect {})),
574            extensions,
575            configuration_parameter: self
576                .configuration_parameter
577                .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())),
578            channel,
579            process_id: self.process_id.unwrap_or_default(),
580            conn_info: self.conn_info.unwrap_or_default(),
581            protocol_ctx: self.protocol_ctx.unwrap_or_default(),
582        }
583    }
584
585    pub fn set_extension(mut self, key: String, value: String) -> Self {
586        self.extensions
587            .get_or_insert_with(HashMap::new)
588            .insert(key, value);
589        self
590    }
591}
592
593#[derive(Debug, Clone, Default)]
594pub struct ConnInfo {
595    pub client_addr: Option<SocketAddr>,
596    pub channel: Channel,
597}
598
599impl Display for ConnInfo {
600    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
601        write!(
602            f,
603            "{}[{}]",
604            self.channel,
605            self.client_addr
606                .map(|addr| addr.to_string())
607                .as_deref()
608                .unwrap_or("unknown client addr")
609        )
610    }
611}
612
613impl ConnInfo {
614    pub fn new(client_addr: Option<SocketAddr>, channel: Channel) -> Self {
615        Self {
616            client_addr,
617            channel,
618        }
619    }
620}
621
622/// Returns the SQL dialect for the given query channel.
623pub fn dialect_for_channel(channel: Channel) -> Arc<dyn Dialect + Send + Sync> {
624    match channel {
625        Channel::Mysql => Arc::new(MySqlDialect {}),
626        Channel::Postgres => Arc::new(PostgreSqlDialect {}),
627        _ => Arc::new(GenericDialect {}),
628    }
629}
630
631#[derive(Default, Debug)]
632pub struct ConfigurationVariables {
633    postgres_bytea_output: ArcSwap<PGByteaOutputValue>,
634    pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>,
635    pg_intervalstyle_format: ArcSwap<PGIntervalStyle>,
636    allow_query_fallback: ArcSwap<bool>,
637}
638
639impl Clone for ConfigurationVariables {
640    fn clone(&self) -> Self {
641        Self {
642            postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()),
643            pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()),
644            pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()),
645            allow_query_fallback: ArcSwap::new(self.allow_query_fallback.load().clone()),
646        }
647    }
648}
649
650impl ConfigurationVariables {
651    pub fn new() -> Self {
652        Self::default()
653    }
654
655    pub fn set_postgres_bytea_output(&self, value: PGByteaOutputValue) {
656        let _ = self.postgres_bytea_output.swap(Arc::new(value));
657    }
658
659    pub fn postgres_bytea_output(&self) -> Arc<PGByteaOutputValue> {
660        self.postgres_bytea_output.load().clone()
661    }
662
663    pub fn pg_datetime_style(&self) -> Arc<(PGDateTimeStyle, PGDateOrder)> {
664        self.pg_datestyle_format.load().clone()
665    }
666
667    pub fn set_pg_datetime_style(&self, style: PGDateTimeStyle, order: PGDateOrder) {
668        self.pg_datestyle_format.swap(Arc::new((style, order)));
669    }
670
671    pub fn pg_intervalstyle_format(&self) -> Arc<PGIntervalStyle> {
672        self.pg_intervalstyle_format.load().clone()
673    }
674
675    pub fn set_pg_intervalstyle_format(&self, value: PGIntervalStyle) {
676        self.pg_intervalstyle_format.swap(Arc::new(value));
677    }
678
679    pub fn allow_query_fallback(&self) -> bool {
680        **self.allow_query_fallback.load()
681    }
682
683    pub fn set_allow_query_fallback(&self, allow: bool) {
684        self.allow_query_fallback.swap(Arc::new(allow));
685    }
686}
687
688#[cfg(test)]
689mod test {
690    use std::collections::HashMap;
691
692    use common_catalog::consts::DEFAULT_CATALOG_NAME;
693
694    use super::*;
695    use crate::Session;
696    use crate::context::Channel;
697
698    #[test]
699    fn test_session() {
700        let session = Session::new(
701            Some("127.0.0.1:9000".parse().unwrap()),
702            Channel::Mysql,
703            Default::default(),
704            100,
705        );
706        // test user_info
707        assert_eq!(session.user_info().username(), "greptime");
708
709        // test channel
710        assert_eq!(session.conn_info().channel, Channel::Mysql);
711        let client_addr = session.conn_info().client_addr.as_ref().unwrap();
712        assert_eq!(client_addr.ip().to_string(), "127.0.0.1");
713        assert_eq!(client_addr.port(), 9000);
714
715        assert_eq!("mysql[127.0.0.1:9000]", session.conn_info().to_string());
716        assert_eq!(100, session.process_id());
717
718        let query_ctx = session.new_query_context();
719        assert!(query_ctx.remote_query_id().is_some());
720    }
721
722    #[test]
723    fn test_context_db_string() {
724        let context = QueryContext::with("a0b1c2d3", "test");
725        assert_eq!("a0b1c2d3-test", context.get_db_string());
726
727        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "test");
728        assert_eq!("test", context.get_db_string());
729    }
730
731    #[test]
732    fn test_fork_has_independent_mutable_session_data() {
733        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
734        let fork = context.fork();
735
736        fork.set_current_schema("private");
737
738        assert_eq!(context.current_schema(), "public");
739        assert_eq!(fork.current_schema(), "private");
740    }
741
742    #[test]
743    fn test_api_query_context_roundtrip_with_sequences() {
744        let api_ctx = api::v1::QueryContext {
745            current_catalog: "c1".to_string(),
746            current_schema: "s1".to_string(),
747            timezone: "UTC".to_string(),
748            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
749            channel: Channel::Grpc as u32,
750            snapshot_seqs: Some(api::v1::SnapshotSequences {
751                snapshot_seqs: HashMap::from([(1, 100)]),
752                sst_min_sequences: HashMap::from([(1, 90)]),
753            }),
754            explain: None,
755        };
756
757        let session_ctx: QueryContext = api_ctx.clone().into();
758        let roundtrip_api: api::v1::QueryContext = session_ctx.into();
759
760        assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
761        assert_eq!(roundtrip_api.current_schema, api_ctx.current_schema);
762        assert_eq!(roundtrip_api.timezone, api_ctx.timezone);
763        assert_eq!(
764            roundtrip_api.extensions.get("flow.return_region_seq"),
765            Some(&"true".to_string())
766        );
767        assert!(
768            roundtrip_api
769                .extensions
770                .contains_key(REMOTE_QUERY_ID_EXTENSION_KEY)
771        );
772        assert_eq!(roundtrip_api.channel, api_ctx.channel);
773        assert_eq!(roundtrip_api.snapshot_seqs, api_ctx.snapshot_seqs);
774    }
775
776    #[test]
777    fn test_query_context_remote_query_id_round_trip() {
778        let query_id = "0195f4fd-c503-7c54-8b8f-7dfb8f6f9c4a";
779        let ctx = QueryContextBuilder::default()
780            .current_catalog(DEFAULT_CATALOG_NAME.to_string())
781            .current_schema("public".to_string())
782            .set_extension(
783                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
784                query_id.to_string(),
785            )
786            .build();
787
788        assert_eq!(ctx.remote_query_id(), Some(query_id));
789        assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
790
791        let proto: api::v1::QueryContext = (&ctx).into();
792        let restored = QueryContext::from(proto);
793        assert_eq!(restored.remote_query_id(), Some(query_id));
794        assert_eq!(
795            restored.remote_query_id_value().unwrap().to_string(),
796            query_id
797        );
798    }
799
800    #[test]
801    fn test_live_analyze_metrics_requires_matching_remote_query_id() {
802        let mut ctx = QueryContext::arc().as_ref().clone();
803        assert!(!ctx.live_analyze_metrics_enabled());
804
805        ctx.enable_live_analyze_metrics();
806        assert!(ctx.live_analyze_metrics_enabled());
807
808        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true");
809        assert!(!ctx.live_analyze_metrics_enabled());
810
811        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id");
812        assert!(!ctx.live_analyze_metrics_enabled());
813    }
814}
session/context.rs

session/
context.rs
