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

session/
context.rs
