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

session/
context.rs
