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    /// Local-only write batching selection; never transported in protobuf extensions.
83    #[builder(default)]
84    batching_enabled: bool,
85    /// Local-only opt-in for logical metric writes, independent of ordinary-table batching.
86    #[builder(default)]
87    logical_batching_enabled: bool,
88    /// Track which protocol the query comes from.
89    #[builder(default)]
90    channel: Channel,
91    /// Process-local admission for one database in this request. Never sent over the wire.
92    #[builder(setter(skip))]
93    admitted_write: Option<(String, String)>,
94    /// Process id for managing on-going queries
95    #[builder(default)]
96    process_id: u32,
97    /// Connection information
98    #[builder(default)]
99    conn_info: ConnInfo,
100    /// Protocol specific context
101    #[builder(default)]
102    protocol_ctx: ProtocolCtx,
103}
104
105/// This fields hold data that is only valid to current query context
106#[derive(Debug, Builder, Clone, Default)]
107pub struct QueryContextMutableFields {
108    warning: Option<String>,
109    // TODO: remove this when format is supported in datafusion
110    explain_format: Option<String>,
111    /// Explain options to control the verbose analyze output.
112    explain_options: Option<ExplainOptions>,
113}
114
115impl Display for QueryContext {
116    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
117        write!(
118            f,
119            "QueryContext{{catalog: {}, schema: {}}}",
120            self.current_catalog(),
121            self.current_schema()
122        )
123    }
124}
125
126impl QueryContextBuilder {
127    pub fn current_schema(mut self, schema: String) -> Self {
128        if self.mutable_session_data.is_none() {
129            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
130        }
131
132        // safe for unwrap because previous none check
133        self.mutable_session_data
134            .as_mut()
135            .unwrap()
136            .write()
137            .unwrap()
138            .schema = schema;
139        self
140    }
141
142    pub fn timezone(mut self, timezone: Timezone) -> Self {
143        if self.mutable_session_data.is_none() {
144            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
145        }
146
147        self.mutable_session_data
148            .as_mut()
149            .unwrap()
150            .write()
151            .unwrap()
152            .timezone = timezone;
153        self
154    }
155
156    pub fn explain_options(mut self, explain_options: Option<ExplainOptions>) -> Self {
157        self.mutable_query_context_data
158            .get_or_insert_default()
159            .write()
160            .unwrap()
161            .explain_options = explain_options;
162        self
163    }
164
165    pub fn skip_wal(mut self, skip_wal: bool) -> Self {
166        self.mutable_session_data
167            .get_or_insert_default()
168            .write()
169            .unwrap()
170            .skip_wal = skip_wal;
171        self
172    }
173
174    pub fn read_preference(mut self, read_preference: ReadPreference) -> Self {
175        self.mutable_session_data
176            .get_or_insert_default()
177            .write()
178            .unwrap()
179            .read_preference = read_preference;
180        self
181    }
182}
183
184impl From<&RegionRequestHeader> for QueryContext {
185    fn from(value: &RegionRequestHeader) -> Self {
186        if let Some(ctx) = &value.query_context {
187            ctx.clone().into()
188        } else {
189            QueryContextBuilder::default()
190                .set_extension(
191                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
192                    generate_remote_query_id(),
193                )
194                .build()
195        }
196    }
197}
198
199impl From<api::v1::QueryContext> for QueryContext {
200    fn from(ctx: api::v1::QueryContext) -> Self {
201        let sequences = ctx.snapshot_seqs.as_ref();
202        QueryContextBuilder::default()
203            .current_catalog(ctx.current_catalog)
204            .current_schema(ctx.current_schema)
205            .timezone(parse_timezone(Some(&ctx.timezone)))
206            .extensions(ctx.extensions)
207            .channel(ctx.channel.into())
208            .snapshot_seqs(Arc::new(RwLock::new(
209                sequences
210                    .map(|x| x.snapshot_seqs.clone())
211                    .unwrap_or_default(),
212            )))
213            .sst_min_sequences(Arc::new(RwLock::new(
214                sequences
215                    .map(|x| x.sst_min_sequences.clone())
216                    .unwrap_or_default(),
217            )))
218            .explain_options(ctx.explain)
219            .build()
220    }
221}
222
223impl From<QueryContext> for api::v1::QueryContext {
224    fn from(
225        QueryContext {
226            current_catalog,
227            mutable_session_data: mutable_inner,
228            extensions,
229            channel,
230            snapshot_seqs,
231            sst_min_sequences,
232            mutable_query_context_data,
233            ..
234        }: QueryContext,
235    ) -> Self {
236        let explain = mutable_query_context_data.read().unwrap().explain_options;
237        let mutable_inner = mutable_inner.read().unwrap();
238        api::v1::QueryContext {
239            current_catalog,
240            current_schema: mutable_inner.schema.clone(),
241            timezone: mutable_inner.timezone.to_string(),
242            extensions,
243            channel: channel as u32,
244            snapshot_seqs: Some(api::v1::SnapshotSequences {
245                snapshot_seqs: snapshot_seqs.read().unwrap().clone(),
246                sst_min_sequences: sst_min_sequences.read().unwrap().clone(),
247            }),
248            explain,
249        }
250    }
251}
252
253impl From<&QueryContext> for api::v1::QueryContext {
254    fn from(ctx: &QueryContext) -> Self {
255        ctx.clone().into()
256    }
257}
258
259impl QueryContext {
260    /// Forks this context with an independent snapshot of mutable session data.
261    ///
262    /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields
263    /// held in mutable session data do not affect this context.
264    pub fn fork(&self) -> Self {
265        let mut fork = self.clone();
266        fork.mutable_session_data = Arc::new(RwLock::new(
267            self.mutable_session_data.read().unwrap().clone(),
268        ));
269        fork
270    }
271
272    /// Forks a context for internal writes after admitting the complete request.
273    /// Call only after write admission succeeds for the current database.
274    pub fn with_write_admission(&self) -> Self {
275        let mut ctx = self.fork();
276        ctx.admitted_write = Some((ctx.current_catalog().to_string(), ctx.current_schema()));
277        ctx
278    }
279
280    /// Returns zero for writes covered by this request's database admission.
281    /// Usage accounting and the original protocol channel remain unchanged.
282    pub fn write_rows_to_admit(&self, catalog: &str, schema: &str, rows: u64) -> u64 {
283        if self
284            .admitted_write
285            .as_ref()
286            .is_some_and(|(c, s)| c == catalog && s == schema)
287        {
288            0
289        } else {
290            rows
291        }
292    }
293
294    pub fn arc() -> QueryContextRef {
295        Arc::new(
296            QueryContextBuilder::default()
297                .set_extension(
298                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
299                    generate_remote_query_id(),
300                )
301                .build(),
302        )
303    }
304
305    /// Create a new  datafusion's ConfigOptions instance based on the current QueryContext.
306    pub fn create_config_options(&self) -> ConfigOptions {
307        let mut config = ConfigOptions::default();
308        config.execution.time_zone = Some(self.timezone().to_string());
309        config
310    }
311
312    pub fn with(catalog: &str, schema: &str) -> QueryContext {
313        QueryContextBuilder::default()
314            .current_catalog(catalog.to_string())
315            .current_schema(schema.to_string())
316            .set_extension(
317                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
318                generate_remote_query_id(),
319            )
320            .build()
321    }
322
323    pub fn with_channel(catalog: &str, schema: &str, channel: Channel) -> QueryContext {
324        QueryContextBuilder::default()
325            .current_catalog(catalog.to_string())
326            .current_schema(schema.to_string())
327            .channel(channel)
328            .set_extension(
329                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
330                generate_remote_query_id(),
331            )
332            .build()
333    }
334
335    pub fn with_db_name(db_name: Option<&str>) -> QueryContext {
336        let (catalog, schema) = db_name
337            .map(|db| {
338                let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
339                (catalog, schema)
340            })
341            .unwrap_or_else(|| {
342                (
343                    DEFAULT_CATALOG_NAME.to_string(),
344                    DEFAULT_SCHEMA_NAME.to_string(),
345                )
346            });
347        QueryContextBuilder::default()
348            .current_catalog(catalog)
349            .current_schema(schema.clone())
350            .set_extension(
351                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
352                generate_remote_query_id(),
353            )
354            .build()
355    }
356
357    pub fn current_schema(&self) -> String {
358        self.mutable_session_data.read().unwrap().schema.clone()
359    }
360
361    pub fn set_current_schema(&self, new_schema: &str) {
362        self.mutable_session_data.write().unwrap().schema = new_schema.to_string();
363    }
364
365    pub fn current_catalog(&self) -> &str {
366        &self.current_catalog
367    }
368
369    pub fn set_current_catalog(&mut self, new_catalog: &str) {
370        self.current_catalog = new_catalog.to_string();
371    }
372
373    pub fn sql_dialect(&self) -> &(dyn Dialect + Send + Sync) {
374        &*self.sql_dialect
375    }
376
377    pub fn get_db_string(&self) -> String {
378        let catalog = self.current_catalog();
379        let schema = self.current_schema();
380        build_db_string(catalog, &schema)
381    }
382
383    pub fn timezone(&self) -> Timezone {
384        self.mutable_session_data.read().unwrap().timezone.clone()
385    }
386
387    pub fn set_timezone(&self, timezone: Timezone) {
388        self.mutable_session_data.write().unwrap().timezone = timezone;
389    }
390
391    /// Returns whether ordinary inserts in this request should skip WAL.
392    pub fn skip_wal(&self) -> bool {
393        self.mutable_session_data.read().unwrap().skip_wal
394    }
395
396    pub fn set_skip_wal(&self, skip_wal: bool) {
397        self.mutable_session_data.write().unwrap().skip_wal = skip_wal;
398    }
399
400    pub fn read_preference(&self) -> ReadPreference {
401        self.mutable_session_data.read().unwrap().read_preference
402    }
403
404    pub fn set_read_preference(&self, read_preference: ReadPreference) {
405        self.mutable_session_data.write().unwrap().read_preference = read_preference;
406    }
407
408    pub fn current_user(&self) -> UserInfoRef {
409        self.mutable_session_data.read().unwrap().user_info.clone()
410    }
411
412    pub fn set_current_user(&self, user: UserInfoRef) {
413        self.mutable_session_data.write().unwrap().user_info = user;
414    }
415
416    pub fn set_extension<S1: Into<String>, S2: Into<String>>(&mut self, key: S1, value: S2) {
417        self.extensions.insert(key.into(), value.into());
418    }
419
420    pub fn extension<S: AsRef<str>>(&self, key: S) -> Option<&str> {
421        self.extensions.get(key.as_ref()).map(|v| v.as_str())
422    }
423
424    pub fn remote_query_id(&self) -> Option<&str> {
425        self.extension(REMOTE_QUERY_ID_EXTENSION_KEY)
426    }
427
428    pub fn remote_query_id_value(&self) -> Option<QueryId> {
429        self.remote_query_id()
430            .and_then(|query_id| query_id.parse().ok())
431    }
432
433    pub fn enable_live_analyze_metrics(&mut self) {
434        if let Some(remote_query_id) = self.remote_query_id().map(str::to_string) {
435            self.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, remote_query_id);
436        }
437    }
438
439    pub fn live_analyze_metrics_enabled(&self) -> bool {
440        self.remote_query_id()
441            .zip(self.extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY))
442            .is_some_and(|(remote_query_id, value)| value == remote_query_id)
443    }
444
445    pub fn extensions(&self) -> HashMap<String, String> {
446        self.extensions.clone()
447    }
448
449    /// Default to double quote and fallback to back quote
450    pub fn quote_style(&self) -> char {
451        if self.sql_dialect().is_delimited_identifier_start('"') {
452            '"'
453        } else if self.sql_dialect().is_delimited_identifier_start('\'') {
454            '\''
455        } else {
456            '`'
457        }
458    }
459
460    pub fn configuration_parameter(&self) -> &ConfigurationVariables {
461        &self.configuration_parameter
462    }
463
464    /// Whether the local HTTP entry point selected logical-table batching.
465    pub fn logical_batching_enabled(&self) -> bool {
466        self.logical_batching_enabled
467    }
468
469    /// Sets local logical-table batching selection without adding a wire-visible extension.
470    pub fn set_logical_batching_enabled(&mut self, enabled: bool) {
471        self.logical_batching_enabled = enabled;
472    }
473
474    /// Whether the local protocol entry point selected ordinary-table batching.
475    pub fn batching_enabled(&self) -> bool {
476        self.batching_enabled
477    }
478
479    /// Sets local write batching selection without adding a wire-visible extension.
480    pub fn set_batching_enabled(&mut self, enabled: bool) {
481        self.batching_enabled = enabled;
482    }
483
484    pub fn channel(&self) -> Channel {
485        self.channel
486    }
487
488    pub fn set_channel(&mut self, channel: Channel) {
489        self.channel = channel;
490    }
491
492    pub fn warning(&self) -> Option<String> {
493        self.mutable_query_context_data
494            .read()
495            .unwrap()
496            .warning
497            .clone()
498    }
499
500    pub fn set_warning(&self, msg: String) {
501        self.mutable_query_context_data.write().unwrap().warning = Some(msg);
502    }
503
504    pub fn explain_format(&self) -> Option<String> {
505        self.mutable_query_context_data
506            .read()
507            .unwrap()
508            .explain_format
509            .clone()
510    }
511
512    pub fn set_explain_format(&self, format: String) {
513        self.mutable_query_context_data
514            .write()
515            .unwrap()
516            .explain_format = Some(format);
517    }
518
519    pub fn explain_verbose(&self) -> bool {
520        self.mutable_query_context_data
521            .read()
522            .unwrap()
523            .explain_options
524            .map(|opts| opts.verbose)
525            .unwrap_or(false)
526    }
527
528    pub fn set_explain_verbose(&self, verbose: bool) {
529        self.mutable_query_context_data
530            .write()
531            .unwrap()
532            .explain_options
533            .get_or_insert_default()
534            .verbose = verbose;
535    }
536
537    pub fn query_timeout(&self) -> Option<Duration> {
538        self.mutable_session_data.read().unwrap().query_timeout
539    }
540
541    pub fn query_timeout_as_millis(&self) -> u128 {
542        let timeout = self.mutable_session_data.read().unwrap().query_timeout;
543        if let Some(t) = timeout {
544            return t.as_millis();
545        }
546        0
547    }
548
549    pub fn set_query_timeout(&self, timeout: Duration) {
550        self.mutable_session_data.write().unwrap().query_timeout = Some(timeout);
551    }
552
553    pub fn insert_cursor(&self, name: String, rb: RecordBatchStreamCursor) {
554        let mut guard = self.mutable_session_data.write().unwrap();
555        guard.cursors.insert(name, Arc::new(rb));
556
557        let cursor_count = guard.cursors.len();
558        if cursor_count > CURSOR_COUNT_WARNING_LIMIT {
559            warn!("Current connection has {} open cursors", cursor_count);
560        }
561    }
562
563    pub fn remove_cursor(&self, name: &str) {
564        let mut guard = self.mutable_session_data.write().unwrap();
565        guard.cursors.remove(name);
566    }
567
568    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
569        let guard = self.mutable_session_data.read().unwrap();
570        let rb = guard.cursors.get(name);
571        rb.cloned()
572    }
573
574    pub fn snapshots(&self) -> HashMap<u64, u64> {
575        self.snapshot_seqs.read().unwrap().clone()
576    }
577
578    pub fn sst_min_sequences(&self) -> HashMap<u64, u64> {
579        self.sst_min_sequences.read().unwrap().clone()
580    }
581
582    pub fn get_snapshot(&self, region_id: u64) -> Option<u64> {
583        self.snapshot_seqs.read().unwrap().get(&region_id).cloned()
584    }
585
586    pub fn set_snapshot(&self, region_id: u64, sequence: u64) {
587        self.snapshot_seqs
588            .write()
589            .unwrap()
590            .insert(region_id, sequence);
591    }
592
593    /// Returns `true` if the session can cast strings to numbers in MySQL style.
594    pub fn auto_string_to_numeric(&self) -> bool {
595        matches!(self.channel, Channel::Mysql)
596    }
597
598    /// Finds the minimal sequence of SST files to scan of a Region.
599    pub fn sst_min_sequence(&self, region_id: u64) -> Option<u64> {
600        self.sst_min_sequences
601            .read()
602            .unwrap()
603            .get(&region_id)
604            .copied()
605    }
606
607    pub fn process_id(&self) -> u32 {
608        self.process_id
609    }
610
611    /// Get client information
612    pub fn conn_info(&self) -> &ConnInfo {
613        &self.conn_info
614    }
615
616    pub fn protocol_ctx(&self) -> &ProtocolCtx {
617        &self.protocol_ctx
618    }
619
620    pub fn set_protocol_ctx(&mut self, protocol_ctx: ProtocolCtx) {
621        self.protocol_ctx = protocol_ctx;
622    }
623}
624
625impl QueryContextBuilder {
626    pub fn build(self) -> QueryContext {
627        let channel = self.channel.unwrap_or_default();
628        let mut extensions = self.extensions.unwrap_or_default();
629        extensions
630            .entry(REMOTE_QUERY_ID_EXTENSION_KEY.to_string())
631            .or_insert_with(generate_remote_query_id);
632        QueryContext {
633            current_catalog: self
634                .current_catalog
635                .unwrap_or_else(|| DEFAULT_CATALOG_NAME.to_string()),
636            snapshot_seqs: self.snapshot_seqs.unwrap_or_default(),
637            sst_min_sequences: self.sst_min_sequences.unwrap_or_default(),
638            mutable_session_data: self.mutable_session_data.unwrap_or_default(),
639            mutable_query_context_data: self.mutable_query_context_data.unwrap_or_default(),
640            sql_dialect: self
641                .sql_dialect
642                .unwrap_or_else(|| Arc::new(GreptimeDbDialect {})),
643            extensions,
644            configuration_parameter: self
645                .configuration_parameter
646                .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())),
647            channel,
648            batching_enabled: self.batching_enabled.unwrap_or_default(),
649            admitted_write: None,
650            logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(),
651            process_id: self.process_id.unwrap_or_default(),
652            conn_info: self.conn_info.unwrap_or_default(),
653            protocol_ctx: self.protocol_ctx.unwrap_or_default(),
654        }
655    }
656
657    pub fn set_extension(mut self, key: String, value: String) -> Self {
658        self.extensions
659            .get_or_insert_with(HashMap::new)
660            .insert(key, value);
661        self
662    }
663}
664
665#[derive(Debug, Clone, Default)]
666pub struct ConnInfo {
667    pub client_addr: Option<SocketAddr>,
668    pub channel: Channel,
669}
670
671impl Display for ConnInfo {
672    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
673        write!(
674            f,
675            "{}[{}]",
676            self.channel,
677            self.client_addr
678                .map(|addr| addr.to_string())
679                .as_deref()
680                .unwrap_or("unknown client addr")
681        )
682    }
683}
684
685impl ConnInfo {
686    pub fn new(client_addr: Option<SocketAddr>, channel: Channel) -> Self {
687        Self {
688            client_addr,
689            channel,
690        }
691    }
692}
693
694/// Returns the SQL dialect for the given query channel.
695pub fn dialect_for_channel(channel: Channel) -> Arc<dyn Dialect + Send + Sync> {
696    match channel {
697        Channel::Mysql => Arc::new(MySqlDialect {}),
698        Channel::Postgres => Arc::new(PostgreSqlDialect {}),
699        _ => Arc::new(GenericDialect {}),
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 crate::Session;
767    use crate::context::{Channel, *};
768
769    #[test]
770    fn test_session() {
771        let session = Session::new(
772            Some("127.0.0.1:9000".parse().unwrap()),
773            Channel::Mysql,
774            Default::default(),
775            100,
776        );
777        // test user_info
778        assert_eq!(session.user_info().username(), "greptime");
779
780        // test channel
781        assert_eq!(session.conn_info().channel, Channel::Mysql);
782        let client_addr = session.conn_info().client_addr.as_ref().unwrap();
783        assert_eq!(client_addr.ip().to_string(), "127.0.0.1");
784        assert_eq!(client_addr.port(), 9000);
785
786        assert_eq!("mysql[127.0.0.1:9000]", session.conn_info().to_string());
787        assert_eq!(100, session.process_id());
788
789        let query_ctx = session.new_query_context();
790        assert!(query_ctx.remote_query_id().is_some());
791    }
792
793    #[test]
794    fn test_context_db_string() {
795        let context = QueryContext::with("a0b1c2d3", "test");
796        assert_eq!("a0b1c2d3-test", context.get_db_string());
797
798        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "test");
799        assert_eq!("test", context.get_db_string());
800    }
801
802    #[test]
803    fn test_fork_has_independent_mutable_session_data() {
804        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
805        let fork = context.fork();
806
807        fork.set_current_schema("private");
808
809        assert_eq!(context.current_schema(), "public");
810        assert_eq!(fork.current_schema(), "private");
811    }
812
813    #[test]
814    fn test_skip_wal_default_builder_and_fork() {
815        let default_context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
816        assert!(!default_context.skip_wal());
817        assert!(!QueryContextBuilder::default().build().skip_wal());
818        let context = QueryContextBuilder::default().skip_wal(true).build();
819        assert!(context.skip_wal());
820        let fork = context.fork();
821        assert!(fork.skip_wal());
822        fork.set_skip_wal(false);
823        assert!(context.skip_wal());
824        assert!(!fork.skip_wal());
825        context.set_skip_wal(false);
826        fork.set_skip_wal(true);
827        assert!(!context.skip_wal());
828        assert!(fork.skip_wal());
829    }
830
831    #[test]
832    fn test_write_admission_is_local_and_database_scoped() {
833        let ctx = QueryContext::with_channel("greptime", "public", Channel::Otlp);
834        let admitted = ctx.with_write_admission();
835        assert_eq!(ctx.write_rows_to_admit("greptime", "public", 100), 100);
836        assert_eq!(admitted.write_rows_to_admit("greptime", "public", 100), 0);
837        assert_eq!(admitted.write_rows_to_admit("greptime", "other", 100), 100);
838        assert_eq!(admitted.write_rows_to_admit("other", "public", 100), 100);
839        assert_eq!(admitted.channel(), Channel::Otlp);
840        assert_eq!(
841            admitted
842                .fork()
843                .write_rows_to_admit("greptime", "public", 100),
844            0
845        );
846        let wire: api::v1::QueryContext = admitted.into();
847        let restored = QueryContext::from(wire);
848        assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100);
849    }
850
851    #[test]
852    fn test_skip_wal_is_not_serialized_in_query_context() {
853        let context = QueryContextBuilder::default().skip_wal(true).build();
854        let api_context: api::v1::QueryContext = context.into();
855        assert!(
856            !api_context
857                .extensions
858                .contains_key(crate::hints::INSERT_SKIP_WAL_HINT)
859        );
860        let restored: QueryContext = api_context.into();
861        assert!(!restored.skip_wal());
862    }
863
864    #[test]
865    fn test_api_query_context_roundtrip_with_sequences() {
866        let api_ctx = api::v1::QueryContext {
867            current_catalog: "c1".to_string(),
868            current_schema: "s1".to_string(),
869            timezone: "UTC".to_string(),
870            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
871            channel: Channel::Internal as u32,
872            snapshot_seqs: Some(api::v1::SnapshotSequences {
873                snapshot_seqs: HashMap::from([(1, 100)]),
874                sst_min_sequences: HashMap::from([(1, 90)]),
875            }),
876            explain: None,
877        };
878
879        let session_ctx: QueryContext = api_ctx.clone().into();
880        let roundtrip_api: api::v1::QueryContext = session_ctx.into();
881
882        assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
883        assert_eq!(roundtrip_api.current_schema, api_ctx.current_schema);
884        assert_eq!(roundtrip_api.timezone, api_ctx.timezone);
885        assert_eq!(
886            roundtrip_api.extensions.get("flow.return_region_seq"),
887            Some(&"true".to_string())
888        );
889        assert!(
890            roundtrip_api
891                .extensions
892                .contains_key(REMOTE_QUERY_ID_EXTENSION_KEY)
893        );
894        assert_eq!(roundtrip_api.channel, api_ctx.channel);
895        assert_eq!(roundtrip_api.snapshot_seqs, api_ctx.snapshot_seqs);
896    }
897
898    #[test]
899    fn test_query_context_remote_query_id_round_trip() {
900        let query_id = "0195f4fd-c503-7c54-8b8f-7dfb8f6f9c4a";
901        let ctx = QueryContextBuilder::default()
902            .current_catalog(DEFAULT_CATALOG_NAME.to_string())
903            .current_schema("public".to_string())
904            .set_extension(
905                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
906                query_id.to_string(),
907            )
908            .build();
909
910        assert_eq!(ctx.remote_query_id(), Some(query_id));
911        assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
912
913        let proto: api::v1::QueryContext = (&ctx).into();
914        let restored = QueryContext::from(proto);
915        assert_eq!(restored.remote_query_id(), Some(query_id));
916        assert_eq!(
917            restored.remote_query_id_value().unwrap().to_string(),
918            query_id
919        );
920    }
921
922    #[test]
923    fn test_live_analyze_metrics_requires_matching_remote_query_id() {
924        let mut ctx = QueryContext::arc().as_ref().clone();
925        assert!(!ctx.live_analyze_metrics_enabled());
926
927        ctx.enable_live_analyze_metrics();
928        assert!(ctx.live_analyze_metrics_enabled());
929
930        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true");
931        assert!(!ctx.live_analyze_metrics_enabled());
932
933        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id");
934        assert!(!ctx.live_analyze_metrics_enabled());
935    }
936    #[test]
937    fn test_batching_selection_is_local_only() {
938        let mut ctx = QueryContextBuilder::default().build();
939        assert!(!ctx.batching_enabled());
940        assert!(!ctx.logical_batching_enabled());
941        ctx.set_logical_batching_enabled(true);
942        assert!(ctx.clone().logical_batching_enabled());
943        assert!(ctx.fork().logical_batching_enabled());
944        ctx.set_batching_enabled(true);
945        assert!(ctx.clone().batching_enabled());
946        assert!(ctx.fork().batching_enabled());
947        let wire: api::v1::QueryContext = ctx.into();
948        let restored = QueryContext::from(wire);
949        assert!(!restored.batching_enabled());
950        assert!(!restored.logical_batching_enabled());
951        let ctx = QueryContextBuilder::default()
952            .set_extension("batching_enabled".to_string(), "true".to_string())
953            .build();
954        assert!(!ctx.batching_enabled());
955    }
956}
session/context.rs

session/
context.rs
