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 skip_wal(mut self, skip_wal: bool) -> Self {
157        self.mutable_session_data
158            .get_or_insert_default()
159            .write()
160            .unwrap()
161            .skip_wal = skip_wal;
162        self
163    }
164
165    pub fn read_preference(mut self, read_preference: ReadPreference) -> Self {
166        self.mutable_session_data
167            .get_or_insert_default()
168            .write()
169            .unwrap()
170            .read_preference = read_preference;
171        self
172    }
173}
174
175impl From<&RegionRequestHeader> for QueryContext {
176    fn from(value: &RegionRequestHeader) -> Self {
177        if let Some(ctx) = &value.query_context {
178            ctx.clone().into()
179        } else {
180            QueryContextBuilder::default()
181                .set_extension(
182                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
183                    generate_remote_query_id(),
184                )
185                .build()
186        }
187    }
188}
189
190impl From<api::v1::QueryContext> for QueryContext {
191    fn from(ctx: api::v1::QueryContext) -> Self {
192        let sequences = ctx.snapshot_seqs.as_ref();
193        QueryContextBuilder::default()
194            .current_catalog(ctx.current_catalog)
195            .current_schema(ctx.current_schema)
196            .timezone(parse_timezone(Some(&ctx.timezone)))
197            .extensions(ctx.extensions)
198            .channel(ctx.channel.into())
199            .snapshot_seqs(Arc::new(RwLock::new(
200                sequences
201                    .map(|x| x.snapshot_seqs.clone())
202                    .unwrap_or_default(),
203            )))
204            .sst_min_sequences(Arc::new(RwLock::new(
205                sequences
206                    .map(|x| x.sst_min_sequences.clone())
207                    .unwrap_or_default(),
208            )))
209            .explain_options(ctx.explain)
210            .build()
211    }
212}
213
214impl From<QueryContext> for api::v1::QueryContext {
215    fn from(
216        QueryContext {
217            current_catalog,
218            mutable_session_data: mutable_inner,
219            extensions,
220            channel,
221            snapshot_seqs,
222            sst_min_sequences,
223            mutable_query_context_data,
224            ..
225        }: QueryContext,
226    ) -> Self {
227        let explain = mutable_query_context_data.read().unwrap().explain_options;
228        let mutable_inner = mutable_inner.read().unwrap();
229        api::v1::QueryContext {
230            current_catalog,
231            current_schema: mutable_inner.schema.clone(),
232            timezone: mutable_inner.timezone.to_string(),
233            extensions,
234            channel: channel as u32,
235            snapshot_seqs: Some(api::v1::SnapshotSequences {
236                snapshot_seqs: snapshot_seqs.read().unwrap().clone(),
237                sst_min_sequences: sst_min_sequences.read().unwrap().clone(),
238            }),
239            explain,
240        }
241    }
242}
243
244impl From<&QueryContext> for api::v1::QueryContext {
245    fn from(ctx: &QueryContext) -> Self {
246        ctx.clone().into()
247    }
248}
249
250impl QueryContext {
251    /// Forks this context with an independent snapshot of mutable session data.
252    ///
253    /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields
254    /// held in mutable session data do not affect this context.
255    pub fn fork(&self) -> Self {
256        let mut fork = self.clone();
257        fork.mutable_session_data = Arc::new(RwLock::new(
258            self.mutable_session_data.read().unwrap().clone(),
259        ));
260        fork
261    }
262
263    pub fn arc() -> QueryContextRef {
264        Arc::new(
265            QueryContextBuilder::default()
266                .set_extension(
267                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
268                    generate_remote_query_id(),
269                )
270                .build(),
271        )
272    }
273
274    /// Create a new  datafusion's ConfigOptions instance based on the current QueryContext.
275    pub fn create_config_options(&self) -> ConfigOptions {
276        let mut config = ConfigOptions::default();
277        config.execution.time_zone = Some(self.timezone().to_string());
278        config
279    }
280
281    pub fn with(catalog: &str, schema: &str) -> QueryContext {
282        QueryContextBuilder::default()
283            .current_catalog(catalog.to_string())
284            .current_schema(schema.to_string())
285            .set_extension(
286                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
287                generate_remote_query_id(),
288            )
289            .build()
290    }
291
292    pub fn with_channel(catalog: &str, schema: &str, channel: Channel) -> QueryContext {
293        QueryContextBuilder::default()
294            .current_catalog(catalog.to_string())
295            .current_schema(schema.to_string())
296            .channel(channel)
297            .set_extension(
298                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
299                generate_remote_query_id(),
300            )
301            .build()
302    }
303
304    pub fn with_db_name(db_name: Option<&str>) -> QueryContext {
305        let (catalog, schema) = db_name
306            .map(|db| {
307                let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
308                (catalog, schema)
309            })
310            .unwrap_or_else(|| {
311                (
312                    DEFAULT_CATALOG_NAME.to_string(),
313                    DEFAULT_SCHEMA_NAME.to_string(),
314                )
315            });
316        QueryContextBuilder::default()
317            .current_catalog(catalog)
318            .current_schema(schema.clone())
319            .set_extension(
320                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
321                generate_remote_query_id(),
322            )
323            .build()
324    }
325
326    pub fn current_schema(&self) -> String {
327        self.mutable_session_data.read().unwrap().schema.clone()
328    }
329
330    pub fn set_current_schema(&self, new_schema: &str) {
331        self.mutable_session_data.write().unwrap().schema = new_schema.to_string();
332    }
333
334    pub fn current_catalog(&self) -> &str {
335        &self.current_catalog
336    }
337
338    pub fn set_current_catalog(&mut self, new_catalog: &str) {
339        self.current_catalog = new_catalog.to_string();
340    }
341
342    pub fn sql_dialect(&self) -> &(dyn Dialect + Send + Sync) {
343        &*self.sql_dialect
344    }
345
346    pub fn get_db_string(&self) -> String {
347        let catalog = self.current_catalog();
348        let schema = self.current_schema();
349        build_db_string(catalog, &schema)
350    }
351
352    pub fn timezone(&self) -> Timezone {
353        self.mutable_session_data.read().unwrap().timezone.clone()
354    }
355
356    pub fn set_timezone(&self, timezone: Timezone) {
357        self.mutable_session_data.write().unwrap().timezone = timezone;
358    }
359
360    /// Returns whether ordinary inserts in this request should skip WAL.
361    pub fn skip_wal(&self) -> bool {
362        self.mutable_session_data.read().unwrap().skip_wal
363    }
364
365    pub fn set_skip_wal(&self, skip_wal: bool) {
366        self.mutable_session_data.write().unwrap().skip_wal = skip_wal;
367    }
368
369    pub fn read_preference(&self) -> ReadPreference {
370        self.mutable_session_data.read().unwrap().read_preference
371    }
372
373    pub fn set_read_preference(&self, read_preference: ReadPreference) {
374        self.mutable_session_data.write().unwrap().read_preference = read_preference;
375    }
376
377    pub fn current_user(&self) -> UserInfoRef {
378        self.mutable_session_data.read().unwrap().user_info.clone()
379    }
380
381    pub fn set_current_user(&self, user: UserInfoRef) {
382        self.mutable_session_data.write().unwrap().user_info = user;
383    }
384
385    pub fn set_extension<S1: Into<String>, S2: Into<String>>(&mut self, key: S1, value: S2) {
386        self.extensions.insert(key.into(), value.into());
387    }
388
389    pub fn extension<S: AsRef<str>>(&self, key: S) -> Option<&str> {
390        self.extensions.get(key.as_ref()).map(|v| v.as_str())
391    }
392
393    pub fn remote_query_id(&self) -> Option<&str> {
394        self.extension(REMOTE_QUERY_ID_EXTENSION_KEY)
395    }
396
397    pub fn remote_query_id_value(&self) -> Option<QueryId> {
398        self.remote_query_id()
399            .and_then(|query_id| query_id.parse().ok())
400    }
401
402    pub fn enable_live_analyze_metrics(&mut self) {
403        if let Some(remote_query_id) = self.remote_query_id().map(str::to_string) {
404            self.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, remote_query_id);
405        }
406    }
407
408    pub fn live_analyze_metrics_enabled(&self) -> bool {
409        self.remote_query_id()
410            .zip(self.extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY))
411            .is_some_and(|(remote_query_id, value)| value == remote_query_id)
412    }
413
414    pub fn extensions(&self) -> HashMap<String, String> {
415        self.extensions.clone()
416    }
417
418    /// Default to double quote and fallback to back quote
419    pub fn quote_style(&self) -> char {
420        if self.sql_dialect().is_delimited_identifier_start('"') {
421            '"'
422        } else if self.sql_dialect().is_delimited_identifier_start('\'') {
423            '\''
424        } else {
425            '`'
426        }
427    }
428
429    pub fn configuration_parameter(&self) -> &ConfigurationVariables {
430        &self.configuration_parameter
431    }
432
433    pub fn channel(&self) -> Channel {
434        self.channel
435    }
436
437    pub fn set_channel(&mut self, channel: Channel) {
438        self.channel = channel;
439    }
440
441    pub fn warning(&self) -> Option<String> {
442        self.mutable_query_context_data
443            .read()
444            .unwrap()
445            .warning
446            .clone()
447    }
448
449    pub fn set_warning(&self, msg: String) {
450        self.mutable_query_context_data.write().unwrap().warning = Some(msg);
451    }
452
453    pub fn explain_format(&self) -> Option<String> {
454        self.mutable_query_context_data
455            .read()
456            .unwrap()
457            .explain_format
458            .clone()
459    }
460
461    pub fn set_explain_format(&self, format: String) {
462        self.mutable_query_context_data
463            .write()
464            .unwrap()
465            .explain_format = Some(format);
466    }
467
468    pub fn explain_verbose(&self) -> bool {
469        self.mutable_query_context_data
470            .read()
471            .unwrap()
472            .explain_options
473            .map(|opts| opts.verbose)
474            .unwrap_or(false)
475    }
476
477    pub fn set_explain_verbose(&self, verbose: bool) {
478        self.mutable_query_context_data
479            .write()
480            .unwrap()
481            .explain_options
482            .get_or_insert_default()
483            .verbose = verbose;
484    }
485
486    pub fn query_timeout(&self) -> Option<Duration> {
487        self.mutable_session_data.read().unwrap().query_timeout
488    }
489
490    pub fn query_timeout_as_millis(&self) -> u128 {
491        let timeout = self.mutable_session_data.read().unwrap().query_timeout;
492        if let Some(t) = timeout {
493            return t.as_millis();
494        }
495        0
496    }
497
498    pub fn set_query_timeout(&self, timeout: Duration) {
499        self.mutable_session_data.write().unwrap().query_timeout = Some(timeout);
500    }
501
502    pub fn insert_cursor(&self, name: String, rb: RecordBatchStreamCursor) {
503        let mut guard = self.mutable_session_data.write().unwrap();
504        guard.cursors.insert(name, Arc::new(rb));
505
506        let cursor_count = guard.cursors.len();
507        if cursor_count > CURSOR_COUNT_WARNING_LIMIT {
508            warn!("Current connection has {} open cursors", cursor_count);
509        }
510    }
511
512    pub fn remove_cursor(&self, name: &str) {
513        let mut guard = self.mutable_session_data.write().unwrap();
514        guard.cursors.remove(name);
515    }
516
517    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
518        let guard = self.mutable_session_data.read().unwrap();
519        let rb = guard.cursors.get(name);
520        rb.cloned()
521    }
522
523    pub fn snapshots(&self) -> HashMap<u64, u64> {
524        self.snapshot_seqs.read().unwrap().clone()
525    }
526
527    pub fn sst_min_sequences(&self) -> HashMap<u64, u64> {
528        self.sst_min_sequences.read().unwrap().clone()
529    }
530
531    pub fn get_snapshot(&self, region_id: u64) -> Option<u64> {
532        self.snapshot_seqs.read().unwrap().get(&region_id).cloned()
533    }
534
535    pub fn set_snapshot(&self, region_id: u64, sequence: u64) {
536        self.snapshot_seqs
537            .write()
538            .unwrap()
539            .insert(region_id, sequence);
540    }
541
542    /// Returns `true` if the session can cast strings to numbers in MySQL style.
543    pub fn auto_string_to_numeric(&self) -> bool {
544        matches!(self.channel, Channel::Mysql)
545    }
546
547    /// Finds the minimal sequence of SST files to scan of a Region.
548    pub fn sst_min_sequence(&self, region_id: u64) -> Option<u64> {
549        self.sst_min_sequences
550            .read()
551            .unwrap()
552            .get(&region_id)
553            .copied()
554    }
555
556    pub fn process_id(&self) -> u32 {
557        self.process_id
558    }
559
560    /// Get client information
561    pub fn conn_info(&self) -> &ConnInfo {
562        &self.conn_info
563    }
564
565    pub fn protocol_ctx(&self) -> &ProtocolCtx {
566        &self.protocol_ctx
567    }
568
569    pub fn set_protocol_ctx(&mut self, protocol_ctx: ProtocolCtx) {
570        self.protocol_ctx = protocol_ctx;
571    }
572}
573
574impl QueryContextBuilder {
575    pub fn build(self) -> QueryContext {
576        let channel = self.channel.unwrap_or_default();
577        let mut extensions = self.extensions.unwrap_or_default();
578        extensions
579            .entry(REMOTE_QUERY_ID_EXTENSION_KEY.to_string())
580            .or_insert_with(generate_remote_query_id);
581        QueryContext {
582            current_catalog: self
583                .current_catalog
584                .unwrap_or_else(|| DEFAULT_CATALOG_NAME.to_string()),
585            snapshot_seqs: self.snapshot_seqs.unwrap_or_default(),
586            sst_min_sequences: self.sst_min_sequences.unwrap_or_default(),
587            mutable_session_data: self.mutable_session_data.unwrap_or_default(),
588            mutable_query_context_data: self.mutable_query_context_data.unwrap_or_default(),
589            sql_dialect: self
590                .sql_dialect
591                .unwrap_or_else(|| Arc::new(GreptimeDbDialect {})),
592            extensions,
593            configuration_parameter: self
594                .configuration_parameter
595                .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())),
596            channel,
597            process_id: self.process_id.unwrap_or_default(),
598            conn_info: self.conn_info.unwrap_or_default(),
599            protocol_ctx: self.protocol_ctx.unwrap_or_default(),
600        }
601    }
602
603    pub fn set_extension(mut self, key: String, value: String) -> Self {
604        self.extensions
605            .get_or_insert_with(HashMap::new)
606            .insert(key, value);
607        self
608    }
609}
610
611#[derive(Debug, Clone, Default)]
612pub struct ConnInfo {
613    pub client_addr: Option<SocketAddr>,
614    pub channel: Channel,
615}
616
617impl Display for ConnInfo {
618    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
619        write!(
620            f,
621            "{}[{}]",
622            self.channel,
623            self.client_addr
624                .map(|addr| addr.to_string())
625                .as_deref()
626                .unwrap_or("unknown client addr")
627        )
628    }
629}
630
631impl ConnInfo {
632    pub fn new(client_addr: Option<SocketAddr>, channel: Channel) -> Self {
633        Self {
634            client_addr,
635            channel,
636        }
637    }
638}
639
640/// Returns the SQL dialect for the given query channel.
641pub fn dialect_for_channel(channel: Channel) -> Arc<dyn Dialect + Send + Sync> {
642    match channel {
643        Channel::Mysql => Arc::new(MySqlDialect {}),
644        Channel::Postgres => Arc::new(PostgreSqlDialect {}),
645        _ => Arc::new(GenericDialect {}),
646    }
647}
648
649#[derive(Default, Debug)]
650pub struct ConfigurationVariables {
651    postgres_bytea_output: ArcSwap<PGByteaOutputValue>,
652    pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>,
653    pg_intervalstyle_format: ArcSwap<PGIntervalStyle>,
654    allow_query_fallback: ArcSwap<bool>,
655}
656
657impl Clone for ConfigurationVariables {
658    fn clone(&self) -> Self {
659        Self {
660            postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()),
661            pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()),
662            pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()),
663            allow_query_fallback: ArcSwap::new(self.allow_query_fallback.load().clone()),
664        }
665    }
666}
667
668impl ConfigurationVariables {
669    pub fn new() -> Self {
670        Self::default()
671    }
672
673    pub fn set_postgres_bytea_output(&self, value: PGByteaOutputValue) {
674        let _ = self.postgres_bytea_output.swap(Arc::new(value));
675    }
676
677    pub fn postgres_bytea_output(&self) -> Arc<PGByteaOutputValue> {
678        self.postgres_bytea_output.load().clone()
679    }
680
681    pub fn pg_datetime_style(&self) -> Arc<(PGDateTimeStyle, PGDateOrder)> {
682        self.pg_datestyle_format.load().clone()
683    }
684
685    pub fn set_pg_datetime_style(&self, style: PGDateTimeStyle, order: PGDateOrder) {
686        self.pg_datestyle_format.swap(Arc::new((style, order)));
687    }
688
689    pub fn pg_intervalstyle_format(&self) -> Arc<PGIntervalStyle> {
690        self.pg_intervalstyle_format.load().clone()
691    }
692
693    pub fn set_pg_intervalstyle_format(&self, value: PGIntervalStyle) {
694        self.pg_intervalstyle_format.swap(Arc::new(value));
695    }
696
697    pub fn allow_query_fallback(&self) -> bool {
698        **self.allow_query_fallback.load()
699    }
700
701    pub fn set_allow_query_fallback(&self, allow: bool) {
702        self.allow_query_fallback.swap(Arc::new(allow));
703    }
704}
705
706#[cfg(test)]
707mod test {
708    use std::collections::HashMap;
709
710    use common_catalog::consts::DEFAULT_CATALOG_NAME;
711
712    use super::*;
713    use crate::Session;
714    use crate::context::Channel;
715
716    #[test]
717    fn test_session() {
718        let session = Session::new(
719            Some("127.0.0.1:9000".parse().unwrap()),
720            Channel::Mysql,
721            Default::default(),
722            100,
723        );
724        // test user_info
725        assert_eq!(session.user_info().username(), "greptime");
726
727        // test channel
728        assert_eq!(session.conn_info().channel, Channel::Mysql);
729        let client_addr = session.conn_info().client_addr.as_ref().unwrap();
730        assert_eq!(client_addr.ip().to_string(), "127.0.0.1");
731        assert_eq!(client_addr.port(), 9000);
732
733        assert_eq!("mysql[127.0.0.1:9000]", session.conn_info().to_string());
734        assert_eq!(100, session.process_id());
735
736        let query_ctx = session.new_query_context();
737        assert!(query_ctx.remote_query_id().is_some());
738    }
739
740    #[test]
741    fn test_context_db_string() {
742        let context = QueryContext::with("a0b1c2d3", "test");
743        assert_eq!("a0b1c2d3-test", context.get_db_string());
744
745        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "test");
746        assert_eq!("test", context.get_db_string());
747    }
748
749    #[test]
750    fn test_fork_has_independent_mutable_session_data() {
751        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
752        let fork = context.fork();
753
754        fork.set_current_schema("private");
755
756        assert_eq!(context.current_schema(), "public");
757        assert_eq!(fork.current_schema(), "private");
758    }
759
760    #[test]
761    fn test_skip_wal_default_builder_and_fork() {
762        let default_context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
763        assert!(!default_context.skip_wal());
764        assert!(!QueryContextBuilder::default().build().skip_wal());
765        let context = QueryContextBuilder::default().skip_wal(true).build();
766        assert!(context.skip_wal());
767        let fork = context.fork();
768        assert!(fork.skip_wal());
769        fork.set_skip_wal(false);
770        assert!(context.skip_wal());
771        assert!(!fork.skip_wal());
772        context.set_skip_wal(false);
773        fork.set_skip_wal(true);
774        assert!(!context.skip_wal());
775        assert!(fork.skip_wal());
776    }
777
778    #[test]
779    fn test_skip_wal_is_not_serialized_in_query_context() {
780        let context = QueryContextBuilder::default().skip_wal(true).build();
781        let api_context: api::v1::QueryContext = context.into();
782        assert!(
783            !api_context
784                .extensions
785                .contains_key(crate::hints::INSERT_SKIP_WAL_HINT)
786        );
787        let restored: QueryContext = api_context.into();
788        assert!(!restored.skip_wal());
789    }
790
791    #[test]
792    fn test_api_query_context_roundtrip_with_sequences() {
793        let api_ctx = api::v1::QueryContext {
794            current_catalog: "c1".to_string(),
795            current_schema: "s1".to_string(),
796            timezone: "UTC".to_string(),
797            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
798            channel: Channel::Grpc as u32,
799            snapshot_seqs: Some(api::v1::SnapshotSequences {
800                snapshot_seqs: HashMap::from([(1, 100)]),
801                sst_min_sequences: HashMap::from([(1, 90)]),
802            }),
803            explain: None,
804        };
805
806        let session_ctx: QueryContext = api_ctx.clone().into();
807        let roundtrip_api: api::v1::QueryContext = session_ctx.into();
808
809        assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
810        assert_eq!(roundtrip_api.current_schema, api_ctx.current_schema);
811        assert_eq!(roundtrip_api.timezone, api_ctx.timezone);
812        assert_eq!(
813            roundtrip_api.extensions.get("flow.return_region_seq"),
814            Some(&"true".to_string())
815        );
816        assert!(
817            roundtrip_api
818                .extensions
819                .contains_key(REMOTE_QUERY_ID_EXTENSION_KEY)
820        );
821        assert_eq!(roundtrip_api.channel, api_ctx.channel);
822        assert_eq!(roundtrip_api.snapshot_seqs, api_ctx.snapshot_seqs);
823    }
824
825    #[test]
826    fn test_query_context_remote_query_id_round_trip() {
827        let query_id = "0195f4fd-c503-7c54-8b8f-7dfb8f6f9c4a";
828        let ctx = QueryContextBuilder::default()
829            .current_catalog(DEFAULT_CATALOG_NAME.to_string())
830            .current_schema("public".to_string())
831            .set_extension(
832                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
833                query_id.to_string(),
834            )
835            .build();
836
837        assert_eq!(ctx.remote_query_id(), Some(query_id));
838        assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
839
840        let proto: api::v1::QueryContext = (&ctx).into();
841        let restored = QueryContext::from(proto);
842        assert_eq!(restored.remote_query_id(), Some(query_id));
843        assert_eq!(
844            restored.remote_query_id_value().unwrap().to_string(),
845            query_id
846        );
847    }
848
849    #[test]
850    fn test_live_analyze_metrics_requires_matching_remote_query_id() {
851        let mut ctx = QueryContext::arc().as_ref().clone();
852        assert!(!ctx.live_analyze_metrics_enabled());
853
854        ctx.enable_live_analyze_metrics();
855        assert!(ctx.live_analyze_metrics_enabled());
856
857        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true");
858        assert!(!ctx.live_analyze_metrics_enabled());
859
860        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id");
861        assert!(!ctx.live_analyze_metrics_enabled());
862    }
863}
session/context.rs

session/
context.rs
