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::str::FromStr;
19use std::sync::{Arc, RwLock};
20use std::time::Duration;
21
22use api::v1::ExplainOptions;
23use api::v1::region::RegionRequestHeader;
24use arc_swap::ArcSwap;
25use auth::UserInfoRef;
26pub use common_base::protocol::Channel;
27use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
28use common_catalog::{build_db_string, parse_catalog_and_schema_from_db_string};
29use common_recordbatch::cursor::RecordBatchStreamCursor;
30use common_telemetry::warn;
31use common_time::Timezone;
32use common_time::timezone::parse_timezone;
33use datafusion_common::config::ConfigOptions;
34use derive_builder::Builder;
35use sql::dialect::{Dialect, GenericDialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
36
37pub use crate::hints::{
38    LIVE_ANALYZE_METRICS_EXTENSION_KEY, READ_PREFERENCE_EXTENSION_KEY,
39    REMOTE_QUERY_ID_EXTENSION_KEY, SUPPORT_FLIGHT_METRICS_BEFORE_BATCH_EXTENSION_KEY,
40};
41use crate::protocol_ctx::ProtocolCtx;
42use crate::query_id::QueryId;
43use crate::session_config::{
44    InvalidConfigValueSnafu, PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle,
45};
46use crate::{MutableInner, ReadPreference};
47
48pub type QueryContextRef = Arc<QueryContext>;
49pub type ConnInfoRef = Arc<ConnInfo>;
50
51pub const FLIGHT_METRICS_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1);
52
53const CURSOR_COUNT_WARNING_LIMIT: usize = 10;
54
55pub fn generate_remote_query_id() -> String {
56    generate_remote_query_id_value().to_string()
57}
58
59pub fn generate_remote_query_id_value() -> QueryId {
60    QueryId::new()
61}
62
63#[derive(Debug, Builder, Clone)]
64#[builder(pattern = "owned")]
65#[builder(build_fn(skip))]
66pub struct QueryContext {
67    current_catalog: String,
68    /// mapping of RegionId to SequenceNumber, for snapshot read, meaning that the read should only
69    /// container data that was committed before(and include) the given sequence number
70    snapshot_seqs: Arc<RwLock<HashMap<u64, u64>>>,
71    /// Mappings of the RegionId to the minimal sequence of SST file to scan.
72    sst_min_sequences: Arc<RwLock<HashMap<u64, u64>>>,
73    // we use Arc<RwLock>> for modifiable fields
74    #[builder(default)]
75    mutable_session_data: Arc<RwLock<MutableInner>>,
76    #[builder(default)]
77    mutable_query_context_data: Arc<RwLock<QueryContextMutableFields>>,
78    sql_dialect: Arc<dyn Dialect + Send + Sync>,
79    #[builder(default)]
80    extensions: HashMap<String, String>,
81    #[builder(setter(skip))]
82    query_options_snapshotted: bool,
83    /// The configuration parameter are used to store the parameters that are set by the user
84    #[builder(default)]
85    configuration_parameter: Arc<ConfigurationVariables>,
86    /// Local-only write batching selection; never transported in protobuf extensions.
87    #[builder(default)]
88    batching_enabled: bool,
89    /// Local-only opt-in for logical metric writes, independent of ordinary-table batching.
90    #[builder(default)]
91    logical_batching_enabled: bool,
92    /// Track which protocol the query comes from.
93    #[builder(default)]
94    channel: Channel,
95    /// Process-local admission for one database in this request. Never sent over the wire.
96    #[builder(setter(skip))]
97    admitted_write: Option<(String, String)>,
98    /// Process id for managing on-going queries
99    #[builder(default)]
100    process_id: u32,
101    /// Connection information
102    #[builder(default)]
103    conn_info: ConnInfo,
104    /// Protocol specific context
105    #[builder(default)]
106    protocol_ctx: ProtocolCtx,
107}
108
109/// This fields hold data that is only valid to current query context
110#[derive(Debug, Builder, Clone, Default)]
111pub struct QueryContextMutableFields {
112    warning: Option<String>,
113    // TODO: remove this when format is supported in datafusion
114    explain_format: Option<String>,
115    /// Explain options to control the verbose analyze output.
116    explain_options: Option<ExplainOptions>,
117}
118
119impl Display for QueryContext {
120    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
121        write!(
122            f,
123            "QueryContext{{catalog: {}, schema: {}}}",
124            self.current_catalog(),
125            self.current_schema()
126        )
127    }
128}
129
130impl QueryContextBuilder {
131    pub fn current_schema(mut self, schema: String) -> Self {
132        if self.mutable_session_data.is_none() {
133            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
134        }
135
136        // safe for unwrap because previous none check
137        self.mutable_session_data
138            .as_mut()
139            .unwrap()
140            .write()
141            .unwrap()
142            .schema = schema;
143        self
144    }
145
146    pub fn timezone(mut self, timezone: Timezone) -> Self {
147        if self.mutable_session_data.is_none() {
148            self.mutable_session_data = Some(Arc::new(RwLock::new(MutableInner::default())));
149        }
150
151        self.mutable_session_data
152            .as_mut()
153            .unwrap()
154            .write()
155            .unwrap()
156            .timezone = timezone;
157        self
158    }
159
160    pub fn explain_options(mut self, explain_options: Option<ExplainOptions>) -> Self {
161        self.mutable_query_context_data
162            .get_or_insert_default()
163            .write()
164            .unwrap()
165            .explain_options = explain_options;
166        self
167    }
168
169    pub fn skip_wal(mut self, skip_wal: bool) -> Self {
170        self.mutable_session_data
171            .get_or_insert_default()
172            .write()
173            .unwrap()
174            .skip_wal = skip_wal;
175        self
176    }
177
178    pub fn read_preference(mut self, read_preference: ReadPreference) -> Self {
179        self.mutable_session_data
180            .get_or_insert_default()
181            .write()
182            .unwrap()
183            .read_preference = read_preference;
184        self
185    }
186}
187
188impl TryFrom<&RegionRequestHeader> for QueryContext {
189    type Error = crate::session_config::Error;
190
191    fn try_from(value: &RegionRequestHeader) -> Result<Self, Self::Error> {
192        if let Some(ctx) = &value.query_context {
193            ctx.clone().try_into()
194        } else {
195            Ok(QueryContextBuilder::default()
196                .set_extension(
197                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
198                    generate_remote_query_id(),
199                )
200                .build())
201        }
202    }
203}
204
205impl TryFrom<api::v1::QueryContext> for QueryContext {
206    type Error = crate::session_config::Error;
207
208    fn try_from(mut ctx: api::v1::QueryContext) -> Result<Self, Self::Error> {
209        let read_preference = ctx.extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
210        let read_preference = match read_preference {
211            Some(value) => ReadPreference::from_str(&value).map_err(|_| {
212                InvalidConfigValueSnafu {
213                    name: "read_preference",
214                    value,
215                    hint: "Expected a supported read preference",
216                }
217                .build()
218            })?,
219            None => ReadPreference::default(),
220        };
221        let sequences = ctx.snapshot_seqs.as_ref();
222        Ok(QueryContextBuilder::default()
223            .current_catalog(ctx.current_catalog)
224            .current_schema(ctx.current_schema)
225            .timezone(parse_timezone(Some(&ctx.timezone)))
226            .extensions(ctx.extensions)
227            .channel(ctx.channel.into())
228            .snapshot_seqs(Arc::new(RwLock::new(
229                sequences
230                    .map(|x| x.snapshot_seqs.clone())
231                    .unwrap_or_default(),
232            )))
233            .sst_min_sequences(Arc::new(RwLock::new(
234                sequences
235                    .map(|x| x.sst_min_sequences.clone())
236                    .unwrap_or_default(),
237            )))
238            .explain_options(ctx.explain)
239            .read_preference(read_preference)
240            .build())
241    }
242}
243
244impl From<QueryContext> for api::v1::QueryContext {
245    fn from(
246        QueryContext {
247            current_catalog,
248            mutable_session_data: mutable_inner,
249            extensions,
250            channel,
251            snapshot_seqs,
252            sst_min_sequences,
253            mutable_query_context_data,
254            ..
255        }: QueryContext,
256    ) -> Self {
257        let explain = mutable_query_context_data.read().unwrap().explain_options;
258        let mutable_inner = mutable_inner.read().unwrap();
259        let mut extensions = extensions;
260        extensions.remove(READ_PREFERENCE_EXTENSION_KEY);
261        if mutable_inner.read_preference != ReadPreference::Leader {
262            extensions.insert(
263                READ_PREFERENCE_EXTENSION_KEY.to_string(),
264                mutable_inner
265                    .read_preference
266                    .to_string()
267                    .to_ascii_lowercase(),
268            );
269        }
270        api::v1::QueryContext {
271            current_catalog,
272            current_schema: mutable_inner.schema.clone(),
273            timezone: mutable_inner.timezone.to_string(),
274            extensions,
275            channel: channel as u32,
276            snapshot_seqs: Some(api::v1::SnapshotSequences {
277                snapshot_seqs: snapshot_seqs.read().unwrap().clone(),
278                sst_min_sequences: sst_min_sequences.read().unwrap().clone(),
279            }),
280            explain,
281        }
282    }
283}
284
285impl From<&QueryContext> for api::v1::QueryContext {
286    fn from(ctx: &QueryContext) -> Self {
287        ctx.clone().into()
288    }
289}
290
291impl QueryContext {
292    /// Forks this context with an independent snapshot of mutable session data.
293    ///
294    /// Unlike [`Clone`], changes to the schema, user, timezone, and other fields
295    /// held in mutable session data do not affect this context.
296    pub fn fork(&self) -> Self {
297        let mut fork = self.clone();
298        fork.mutable_session_data = Arc::new(RwLock::new(
299            self.mutable_session_data.read().unwrap().clone(),
300        ));
301        fork
302    }
303
304    /// Forks a context for internal writes after admitting the complete request.
305    /// Call only after write admission succeeds for the current database.
306    pub fn with_write_admission(&self) -> Self {
307        let mut ctx = self.fork();
308        ctx.admitted_write = Some((ctx.current_catalog().to_string(), ctx.current_schema()));
309        ctx
310    }
311
312    /// Returns zero for writes covered by this request's database admission.
313    /// Usage accounting and the original protocol channel remain unchanged.
314    pub fn write_rows_to_admit(&self, catalog: &str, schema: &str, rows: u64) -> u64 {
315        if self
316            .admitted_write
317            .as_ref()
318            .is_some_and(|(c, s)| c == catalog && s == schema)
319        {
320            0
321        } else {
322            rows
323        }
324    }
325
326    pub fn arc() -> QueryContextRef {
327        Arc::new(
328            QueryContextBuilder::default()
329                .set_extension(
330                    REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
331                    generate_remote_query_id(),
332                )
333                .build(),
334        )
335    }
336
337    /// Create a new  datafusion's ConfigOptions instance based on the current QueryContext.
338    pub fn create_config_options(&self) -> ConfigOptions {
339        let mut config = ConfigOptions::default();
340        config.execution.time_zone = Some(self.timezone().to_string());
341        config
342    }
343
344    pub fn with(catalog: &str, schema: &str) -> QueryContext {
345        QueryContextBuilder::default()
346            .current_catalog(catalog.to_string())
347            .current_schema(schema.to_string())
348            .set_extension(
349                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
350                generate_remote_query_id(),
351            )
352            .build()
353    }
354
355    pub fn with_channel(catalog: &str, schema: &str, channel: Channel) -> QueryContext {
356        QueryContextBuilder::default()
357            .current_catalog(catalog.to_string())
358            .current_schema(schema.to_string())
359            .channel(channel)
360            .set_extension(
361                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
362                generate_remote_query_id(),
363            )
364            .build()
365    }
366
367    pub fn with_db_name(db_name: Option<&str>) -> QueryContext {
368        let (catalog, schema) = db_name
369            .map(|db| {
370                let (catalog, schema) = parse_catalog_and_schema_from_db_string(db);
371                (catalog, schema)
372            })
373            .unwrap_or_else(|| {
374                (
375                    DEFAULT_CATALOG_NAME.to_string(),
376                    DEFAULT_SCHEMA_NAME.to_string(),
377                )
378            });
379        QueryContextBuilder::default()
380            .current_catalog(catalog)
381            .current_schema(schema.clone())
382            .set_extension(
383                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
384                generate_remote_query_id(),
385            )
386            .build()
387    }
388
389    pub fn current_schema(&self) -> String {
390        self.mutable_session_data.read().unwrap().schema.clone()
391    }
392
393    pub fn set_current_schema(&self, new_schema: &str) {
394        self.mutable_session_data.write().unwrap().schema = new_schema.to_string();
395    }
396
397    pub fn current_catalog(&self) -> &str {
398        &self.current_catalog
399    }
400
401    pub fn set_current_catalog(&mut self, new_catalog: &str) {
402        self.current_catalog = new_catalog.to_string();
403    }
404
405    pub fn sql_dialect(&self) -> &(dyn Dialect + Send + Sync) {
406        &*self.sql_dialect
407    }
408
409    pub fn get_db_string(&self) -> String {
410        let catalog = self.current_catalog();
411        let schema = self.current_schema();
412        build_db_string(catalog, &schema)
413    }
414
415    pub fn timezone(&self) -> Timezone {
416        self.mutable_session_data.read().unwrap().timezone.clone()
417    }
418
419    pub fn set_timezone(&self, timezone: Timezone) {
420        self.mutable_session_data.write().unwrap().timezone = timezone;
421    }
422
423    /// Returns whether ordinary inserts in this request should skip WAL.
424    pub fn skip_wal(&self) -> bool {
425        self.mutable_session_data.read().unwrap().skip_wal
426    }
427
428    pub fn set_skip_wal(&self, skip_wal: bool) {
429        self.mutable_session_data.write().unwrap().skip_wal = skip_wal;
430    }
431
432    pub fn read_preference(&self) -> ReadPreference {
433        self.mutable_session_data.read().unwrap().read_preference
434    }
435
436    pub fn set_read_preference(&self, read_preference: ReadPreference) {
437        self.mutable_session_data.write().unwrap().read_preference = read_preference;
438    }
439
440    pub fn current_user(&self) -> UserInfoRef {
441        self.mutable_session_data.read().unwrap().user_info.clone()
442    }
443
444    pub fn set_current_user(&self, user: UserInfoRef) {
445        self.mutable_session_data.write().unwrap().user_info = user;
446    }
447
448    pub fn set_extension<S1: Into<String>, S2: Into<String>>(&mut self, key: S1, value: S2) {
449        self.extensions.insert(key.into(), value.into());
450    }
451
452    #[cfg(feature = "enterprise")]
453    pub fn remove_extension(&mut self, key: &str) {
454        self.extensions.remove(key);
455    }
456
457    pub fn extension<S: AsRef<str>>(&self, key: S) -> Option<&str> {
458        self.extensions.get(key.as_ref()).map(|v| v.as_str())
459    }
460
461    pub fn remote_query_id(&self) -> Option<&str> {
462        self.extension(REMOTE_QUERY_ID_EXTENSION_KEY)
463    }
464
465    pub fn remote_query_id_value(&self) -> Option<QueryId> {
466        self.remote_query_id()
467            .and_then(|query_id| query_id.parse().ok())
468    }
469
470    pub fn enable_live_analyze_metrics(&mut self) {
471        if let Some(remote_query_id) = self.remote_query_id().map(str::to_string) {
472            self.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, remote_query_id);
473        }
474    }
475
476    pub fn live_analyze_metrics_enabled(&self) -> bool {
477        self.remote_query_id()
478            .zip(self.extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY))
479            .is_some_and(|(remote_query_id, value)| value == remote_query_id)
480    }
481
482    pub fn extensions(&self) -> HashMap<String, String> {
483        self.extensions.clone()
484    }
485
486    /// Returns normalized query options from the request over session defaults.
487    pub fn effective_query_options(
488        &self,
489    ) -> Result<HashMap<String, String>, crate::session_config::Error> {
490        let mut request_options = HashMap::new();
491        for (key, value) in &self.extensions {
492            if crate::hints::is_reserved_extension_key(key) {
493                continue;
494            }
495            // Older Frontends and Flownodes forward `query_parallelism=0` (their
496            // "use the server default" encoding) untouched inside the wire
497            // query context, so accept it as unset instead of failing the
498            // query on the Datanode side. The canonical
499            // `query.parallelism = 0` stays rejected.
500            if key.eq_ignore_ascii_case("query_parallelism") && value == "0" {
501                continue;
502            }
503            // Wire decoding preserves legacy frontend nonzero values; public entry points enforce the limit.
504            if let Some((canonical, normalized)) =
505                crate::query_options::parse_query_option(key, value, u16::MAX)?
506            {
507                if request_options
508                    .get(&canonical)
509                    .is_some_and(|previous| previous != &normalized)
510                {
511                    return crate::session_config::InvalidConfigValueSnafu {
512                        name: key,
513                        value,
514                        hint: "Conflicting aliases for the same query option",
515                    }
516                    .fail();
517                }
518                request_options.insert(canonical, normalized);
519            }
520        }
521        let mut effective = if self.query_options_snapshotted {
522            HashMap::new()
523        } else {
524            self.configuration_parameter.query_options()
525        };
526        effective.extend(request_options);
527        Ok(effective)
528    }
529
530    /// Creates a per-query immutable option snapshot while retaining shared mutable session data.
531    pub fn query_option_snapshot(&self) -> Result<Self, crate::session_config::Error> {
532        let mut snapshot = self.clone();
533        let options = snapshot.effective_query_options()?;
534        snapshot
535            .extensions
536            .retain(|key, _| crate::query_options::canonical_query_option_name(key).is_none());
537        // Older Datanodes only understand the legacy `query_parallelism` hint key, so
538        // mirror the canonical parallelism override for the wire format while keeping
539        // the canonical key for in-process consumers.
540        let legacy_parallelism = options
541            .get("query.parallelism")
542            .map(|parallelism| ("query_parallelism".to_string(), parallelism.clone()));
543        snapshot.extensions.extend(options);
544        if let Some((key, value)) = legacy_parallelism {
545            snapshot.extensions.entry(key).or_insert(value);
546        }
547        snapshot.query_options_snapshotted = true;
548        Ok(snapshot)
549    }
550
551    /// Default to double quote and fallback to back quote
552    pub fn quote_style(&self) -> char {
553        if self.sql_dialect().is_delimited_identifier_start('"') {
554            '"'
555        } else if self.sql_dialect().is_delimited_identifier_start('\'') {
556            '\''
557        } else {
558            '`'
559        }
560    }
561
562    pub fn configuration_parameter(&self) -> &ConfigurationVariables {
563        &self.configuration_parameter
564    }
565
566    /// Whether the local HTTP entry point selected logical-table batching.
567    pub fn logical_batching_enabled(&self) -> bool {
568        self.logical_batching_enabled
569    }
570
571    /// Sets local logical-table batching selection without adding a wire-visible extension.
572    pub fn set_logical_batching_enabled(&mut self, enabled: bool) {
573        self.logical_batching_enabled = enabled;
574    }
575
576    /// Whether the local protocol entry point selected ordinary-table batching.
577    pub fn batching_enabled(&self) -> bool {
578        self.batching_enabled
579    }
580
581    /// Sets local write batching selection without adding a wire-visible extension.
582    pub fn set_batching_enabled(&mut self, enabled: bool) {
583        self.batching_enabled = enabled;
584    }
585
586    pub fn channel(&self) -> Channel {
587        self.channel
588    }
589
590    pub fn set_channel(&mut self, channel: Channel) {
591        self.channel = channel;
592    }
593
594    pub fn warning(&self) -> Option<String> {
595        self.mutable_query_context_data
596            .read()
597            .unwrap()
598            .warning
599            .clone()
600    }
601
602    pub fn set_warning(&self, msg: String) {
603        self.mutable_query_context_data.write().unwrap().warning = Some(msg);
604    }
605
606    pub fn explain_format(&self) -> Option<String> {
607        self.mutable_query_context_data
608            .read()
609            .unwrap()
610            .explain_format
611            .clone()
612    }
613
614    pub fn set_explain_format(&self, format: String) {
615        self.mutable_query_context_data
616            .write()
617            .unwrap()
618            .explain_format = Some(format);
619    }
620
621    pub fn explain_verbose(&self) -> bool {
622        self.mutable_query_context_data
623            .read()
624            .unwrap()
625            .explain_options
626            .map(|opts| opts.verbose)
627            .unwrap_or(false)
628    }
629
630    pub fn set_explain_verbose(&self, verbose: bool) {
631        self.mutable_query_context_data
632            .write()
633            .unwrap()
634            .explain_options
635            .get_or_insert_default()
636            .verbose = verbose;
637    }
638
639    pub fn query_timeout(&self) -> Option<Duration> {
640        self.mutable_session_data.read().unwrap().query_timeout
641    }
642
643    pub fn query_timeout_as_millis(&self) -> u128 {
644        let timeout = self.mutable_session_data.read().unwrap().query_timeout;
645        if let Some(t) = timeout {
646            return t.as_millis();
647        }
648        0
649    }
650
651    pub fn set_query_timeout(&self, timeout: Duration) {
652        self.mutable_session_data.write().unwrap().query_timeout = Some(timeout);
653    }
654
655    pub fn insert_cursor(&self, name: String, rb: RecordBatchStreamCursor) {
656        let mut guard = self.mutable_session_data.write().unwrap();
657        guard.cursors.insert(name, Arc::new(rb));
658
659        let cursor_count = guard.cursors.len();
660        if cursor_count > CURSOR_COUNT_WARNING_LIMIT {
661            warn!("Current connection has {} open cursors", cursor_count);
662        }
663    }
664
665    pub fn remove_cursor(&self, name: &str) {
666        let mut guard = self.mutable_session_data.write().unwrap();
667        guard.cursors.remove(name);
668    }
669
670    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
671        let guard = self.mutable_session_data.read().unwrap();
672        let rb = guard.cursors.get(name);
673        rb.cloned()
674    }
675
676    pub fn snapshots(&self) -> HashMap<u64, u64> {
677        self.snapshot_seqs.read().unwrap().clone()
678    }
679
680    pub fn sst_min_sequences(&self) -> HashMap<u64, u64> {
681        self.sst_min_sequences.read().unwrap().clone()
682    }
683
684    pub fn get_snapshot(&self, region_id: u64) -> Option<u64> {
685        self.snapshot_seqs.read().unwrap().get(&region_id).cloned()
686    }
687
688    pub fn set_snapshot(&self, region_id: u64, sequence: u64) {
689        self.snapshot_seqs
690            .write()
691            .unwrap()
692            .insert(region_id, sequence);
693    }
694
695    /// Returns `true` if the session can cast strings to numbers in MySQL style.
696    pub fn auto_string_to_numeric(&self) -> bool {
697        matches!(self.channel, Channel::Mysql)
698    }
699
700    /// Finds the minimal sequence of SST files to scan of a Region.
701    pub fn sst_min_sequence(&self, region_id: u64) -> Option<u64> {
702        self.sst_min_sequences
703            .read()
704            .unwrap()
705            .get(&region_id)
706            .copied()
707    }
708
709    pub fn process_id(&self) -> u32 {
710        self.process_id
711    }
712
713    /// Get client information
714    pub fn conn_info(&self) -> &ConnInfo {
715        &self.conn_info
716    }
717
718    pub fn protocol_ctx(&self) -> &ProtocolCtx {
719        &self.protocol_ctx
720    }
721
722    pub fn set_protocol_ctx(&mut self, protocol_ctx: ProtocolCtx) {
723        self.protocol_ctx = protocol_ctx;
724    }
725}
726
727impl QueryContextBuilder {
728    pub fn build(self) -> QueryContext {
729        let channel = self.channel.unwrap_or_default();
730        let mut extensions = self.extensions.unwrap_or_default();
731        extensions
732            .entry(REMOTE_QUERY_ID_EXTENSION_KEY.to_string())
733            .or_insert_with(generate_remote_query_id);
734        QueryContext {
735            current_catalog: self
736                .current_catalog
737                .unwrap_or_else(|| DEFAULT_CATALOG_NAME.to_string()),
738            snapshot_seqs: self.snapshot_seqs.unwrap_or_default(),
739            sst_min_sequences: self.sst_min_sequences.unwrap_or_default(),
740            mutable_session_data: self.mutable_session_data.unwrap_or_default(),
741            mutable_query_context_data: self.mutable_query_context_data.unwrap_or_default(),
742            sql_dialect: self
743                .sql_dialect
744                .unwrap_or_else(|| Arc::new(GreptimeDbDialect {})),
745            extensions,
746            query_options_snapshotted: false,
747            configuration_parameter: self
748                .configuration_parameter
749                .unwrap_or_else(|| Arc::new(ConfigurationVariables::default())),
750            channel,
751            batching_enabled: self.batching_enabled.unwrap_or_default(),
752            admitted_write: None,
753            logical_batching_enabled: self.logical_batching_enabled.unwrap_or_default(),
754            process_id: self.process_id.unwrap_or_default(),
755            conn_info: self.conn_info.unwrap_or_default(),
756            protocol_ctx: self.protocol_ctx.unwrap_or_default(),
757        }
758    }
759
760    pub fn set_extension(mut self, key: String, value: String) -> Self {
761        self.extensions
762            .get_or_insert_with(HashMap::new)
763            .insert(key, value);
764        self
765    }
766}
767
768#[derive(Debug, Clone, Default)]
769pub struct ConnInfo {
770    pub client_addr: Option<SocketAddr>,
771    pub channel: Channel,
772}
773
774impl Display for ConnInfo {
775    fn fmt(&self, f: &mut Formatter) -> std::fmt::Result {
776        write!(
777            f,
778            "{}[{}]",
779            self.channel,
780            self.client_addr
781                .map(|addr| addr.to_string())
782                .as_deref()
783                .unwrap_or("unknown client addr")
784        )
785    }
786}
787
788impl ConnInfo {
789    pub fn new(client_addr: Option<SocketAddr>, channel: Channel) -> Self {
790        Self {
791            client_addr,
792            channel,
793        }
794    }
795}
796
797/// Returns the SQL dialect for the given query channel.
798pub fn dialect_for_channel(channel: Channel) -> Arc<dyn Dialect + Send + Sync> {
799    match channel {
800        Channel::Mysql => Arc::new(MySqlDialect {}),
801        Channel::Postgres => Arc::new(PostgreSqlDialect {}),
802        _ => Arc::new(GenericDialect {}),
803    }
804}
805
806#[derive(Default, Debug)]
807pub struct ConfigurationVariables {
808    postgres_bytea_output: ArcSwap<PGByteaOutputValue>,
809    pg_datestyle_format: ArcSwap<(PGDateTimeStyle, PGDateOrder)>,
810    pg_intervalstyle_format: ArcSwap<PGIntervalStyle>,
811    query_options: RwLock<HashMap<String, String>>,
812}
813
814impl Clone for ConfigurationVariables {
815    fn clone(&self) -> Self {
816        Self {
817            postgres_bytea_output: ArcSwap::new(self.postgres_bytea_output.load().clone()),
818            pg_datestyle_format: ArcSwap::new(self.pg_datestyle_format.load().clone()),
819            pg_intervalstyle_format: ArcSwap::new(self.pg_intervalstyle_format.load().clone()),
820            query_options: RwLock::new(self.query_options.read().unwrap().clone()),
821        }
822    }
823}
824
825impl ConfigurationVariables {
826    pub fn new() -> Self {
827        Self::default()
828    }
829
830    pub fn set_postgres_bytea_output(&self, value: PGByteaOutputValue) {
831        let _ = self.postgres_bytea_output.swap(Arc::new(value));
832    }
833
834    pub fn postgres_bytea_output(&self) -> Arc<PGByteaOutputValue> {
835        self.postgres_bytea_output.load().clone()
836    }
837
838    pub fn pg_datetime_style(&self) -> Arc<(PGDateTimeStyle, PGDateOrder)> {
839        self.pg_datestyle_format.load().clone()
840    }
841
842    pub fn set_pg_datetime_style(&self, style: PGDateTimeStyle, order: PGDateOrder) {
843        self.pg_datestyle_format.swap(Arc::new((style, order)));
844    }
845
846    pub fn pg_intervalstyle_format(&self) -> Arc<PGIntervalStyle> {
847        self.pg_intervalstyle_format.load().clone()
848    }
849
850    pub fn set_pg_intervalstyle_format(&self, value: PGIntervalStyle) {
851        self.pg_intervalstyle_format.swap(Arc::new(value));
852    }
853
854    pub fn allow_query_fallback(&self) -> bool {
855        self.query_options().get("query.allow_query_fallback") == Some(&"true".to_string())
856    }
857
858    pub fn set_allow_query_fallback(&self, allow: bool) {
859        self.query_options
860            .write()
861            .unwrap()
862            .insert("query.allow_query_fallback".to_string(), allow.to_string());
863    }
864
865    /// Validates and stores one canonicalized per-session query option.
866    pub fn set_query_option(
867        &self,
868        key: &str,
869        value: &str,
870    ) -> Result<(), crate::session_config::Error> {
871        let (canonical, normalized) = crate::query_options::parse_query_option(
872            key,
873            value,
874            crate::query_options::MAX_PUBLIC_PARALLELISM,
875        )?
876        .ok_or_else(|| crate::session_config::Error::InvalidConfigValue {
877            name: key.to_string(),
878            value: value.to_string(),
879            hint: "Not a supported query option".to_string(),
880            location: snafu::Location::new(file!(), line!(), column!()),
881        })?;
882        self.query_options
883            .write()
884            .unwrap()
885            .insert(canonical, normalized);
886        Ok(())
887    }
888
889    /// Returns the session's canonical query option overrides.
890    pub fn query_options(&self) -> HashMap<String, String> {
891        self.query_options.read().unwrap().clone()
892    }
893}
894
895#[cfg(test)]
896mod test {
897    use std::collections::HashMap;
898
899    use common_catalog::consts::DEFAULT_CATALOG_NAME;
900
901    use crate::Session;
902    use crate::context::{Channel, *};
903
904    #[test]
905    fn test_query_parallelism_wire_accepts_legacy_value() {
906        let mut context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
907        context.set_extension("query_parallelism", "2048");
908
909        assert_eq!(
910            context
911                .effective_query_options()
912                .unwrap()
913                .get("query.parallelism"),
914            Some(&"2048".to_string())
915        );
916    }
917
918    #[test]
919    fn test_set_query_parallelism_public_limit() {
920        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
921        assert!(
922            context
923                .configuration_parameter()
924                .set_query_option("query.parallelism", "1025")
925                .is_err()
926        );
927        assert!(
928            context
929                .configuration_parameter()
930                .set_query_option("query.parallelism", "1024")
931                .is_ok()
932        );
933    }
934
935    #[test]
936    fn test_session() {
937        let session = Session::new(
938            Some("127.0.0.1:9000".parse().unwrap()),
939            Channel::Mysql,
940            Default::default(),
941            100,
942        );
943        // test user_info
944        assert_eq!(session.user_info().username(), "greptime");
945
946        // test channel
947        assert_eq!(session.conn_info().channel, Channel::Mysql);
948        let client_addr = session.conn_info().client_addr.as_ref().unwrap();
949        assert_eq!(client_addr.ip().to_string(), "127.0.0.1");
950        assert_eq!(client_addr.port(), 9000);
951
952        assert_eq!("mysql[127.0.0.1:9000]", session.conn_info().to_string());
953        assert_eq!(100, session.process_id());
954
955        let query_ctx = session.new_query_context();
956        assert!(query_ctx.remote_query_id().is_some());
957    }
958
959    #[test]
960    fn test_context_db_string() {
961        let context = QueryContext::with("a0b1c2d3", "test");
962        assert_eq!("a0b1c2d3-test", context.get_db_string());
963
964        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "test");
965        assert_eq!("test", context.get_db_string());
966    }
967
968    #[test]
969    fn test_fork_has_independent_mutable_session_data() {
970        let context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
971        let fork = context.fork();
972
973        fork.set_current_schema("private");
974
975        assert_eq!(context.current_schema(), "public");
976        assert_eq!(fork.current_schema(), "private");
977    }
978
979    #[test]
980    fn test_skip_wal_default_builder_and_fork() {
981        let default_context = QueryContext::with(DEFAULT_CATALOG_NAME, "public");
982        assert!(!default_context.skip_wal());
983        assert!(!QueryContextBuilder::default().build().skip_wal());
984        let context = QueryContextBuilder::default().skip_wal(true).build();
985        assert!(context.skip_wal());
986        let fork = context.fork();
987        assert!(fork.skip_wal());
988        fork.set_skip_wal(false);
989        assert!(context.skip_wal());
990        assert!(!fork.skip_wal());
991        context.set_skip_wal(false);
992        fork.set_skip_wal(true);
993        assert!(!context.skip_wal());
994        assert!(fork.skip_wal());
995    }
996
997    #[test]
998    fn test_write_admission_is_local_and_database_scoped() {
999        let ctx = QueryContext::with_channel("greptime", "public", Channel::Otlp);
1000        let admitted = ctx.with_write_admission();
1001        assert_eq!(ctx.write_rows_to_admit("greptime", "public", 100), 100);
1002        assert_eq!(admitted.write_rows_to_admit("greptime", "public", 100), 0);
1003        assert_eq!(admitted.write_rows_to_admit("greptime", "other", 100), 100);
1004        assert_eq!(admitted.write_rows_to_admit("other", "public", 100), 100);
1005        assert_eq!(admitted.channel(), Channel::Otlp);
1006        assert_eq!(
1007            admitted
1008                .fork()
1009                .write_rows_to_admit("greptime", "public", 100),
1010            0
1011        );
1012        let wire: api::v1::QueryContext = admitted.into();
1013        let restored = QueryContext::try_from(wire).unwrap();
1014        assert_eq!(restored.write_rows_to_admit("greptime", "public", 100), 100);
1015    }
1016
1017    #[test]
1018    fn test_skip_wal_is_not_serialized_in_query_context() {
1019        let context = QueryContextBuilder::default().skip_wal(true).build();
1020        let api_context: api::v1::QueryContext = context.into();
1021        assert!(
1022            !api_context
1023                .extensions
1024                .contains_key(crate::hints::INSERT_SKIP_WAL_HINT)
1025        );
1026        let restored: QueryContext = api_context.try_into().unwrap();
1027        assert!(!restored.skip_wal());
1028    }
1029
1030    #[test]
1031    fn test_api_query_context_roundtrip_with_sequences() {
1032        let api_ctx = api::v1::QueryContext {
1033            current_catalog: "c1".to_string(),
1034            current_schema: "s1".to_string(),
1035            timezone: "UTC".to_string(),
1036            extensions: HashMap::from([("flow.return_region_seq".to_string(), "true".to_string())]),
1037            channel: Channel::Internal as u32,
1038            snapshot_seqs: Some(api::v1::SnapshotSequences {
1039                snapshot_seqs: HashMap::from([(1, 100)]),
1040                sst_min_sequences: HashMap::from([(1, 90)]),
1041            }),
1042            explain: None,
1043        };
1044
1045        let session_ctx: QueryContext = api_ctx.clone().try_into().unwrap();
1046        let roundtrip_api: api::v1::QueryContext = session_ctx.into();
1047
1048        assert_eq!(roundtrip_api.current_catalog, api_ctx.current_catalog);
1049        assert_eq!(roundtrip_api.current_schema, api_ctx.current_schema);
1050        assert_eq!(roundtrip_api.timezone, api_ctx.timezone);
1051        assert_eq!(
1052            roundtrip_api.extensions.get("flow.return_region_seq"),
1053            Some(&"true".to_string())
1054        );
1055        assert!(
1056            roundtrip_api
1057                .extensions
1058                .contains_key(REMOTE_QUERY_ID_EXTENSION_KEY)
1059        );
1060        assert_eq!(roundtrip_api.channel, api_ctx.channel);
1061        assert_eq!(roundtrip_api.snapshot_seqs, api_ctx.snapshot_seqs);
1062    }
1063
1064    #[test]
1065    fn test_query_context_remote_query_id_round_trip() {
1066        let query_id = "0195f4fd-c503-7c54-8b8f-7dfb8f6f9c4a";
1067        let ctx = QueryContextBuilder::default()
1068            .current_catalog(DEFAULT_CATALOG_NAME.to_string())
1069            .current_schema("public".to_string())
1070            .set_extension(
1071                REMOTE_QUERY_ID_EXTENSION_KEY.to_string(),
1072                query_id.to_string(),
1073            )
1074            .build();
1075
1076        assert_eq!(ctx.remote_query_id(), Some(query_id));
1077        assert_eq!(ctx.remote_query_id_value().unwrap().to_string(), query_id);
1078
1079        let proto: api::v1::QueryContext = (&ctx).into();
1080        let restored = QueryContext::try_from(proto).unwrap();
1081        assert_eq!(restored.remote_query_id(), Some(query_id));
1082        assert_eq!(
1083            restored.remote_query_id_value().unwrap().to_string(),
1084            query_id
1085        );
1086    }
1087
1088    #[test]
1089    fn test_query_context_read_preference_wire_values() {
1090        let wire = |value: Option<&str>| api::v1::QueryContext {
1091            current_catalog: DEFAULT_CATALOG_NAME.to_string(),
1092            current_schema: DEFAULT_SCHEMA_NAME.to_string(),
1093            timezone: "UTC".to_string(),
1094            extensions: value
1095                .map(|value| {
1096                    HashMap::from([(READ_PREFERENCE_EXTENSION_KEY.to_string(), value.to_string())])
1097                })
1098                .unwrap_or_default(),
1099            channel: Channel::Internal as u32,
1100            snapshot_seqs: None,
1101            explain: None,
1102        };
1103
1104        let header = RegionRequestHeader::default();
1105        let header_context = QueryContext::try_from(&header).unwrap();
1106        assert_eq!(header_context.read_preference(), ReadPreference::Leader);
1107
1108        for value in [None, Some("leader"), Some("LEADER")] {
1109            let ctx = QueryContext::try_from(wire(value)).unwrap();
1110            assert_eq!(ctx.read_preference(), ReadPreference::Leader);
1111            assert!(!ctx.extensions().contains_key(READ_PREFERENCE_EXTENSION_KEY));
1112        }
1113        for value in ["", "unknown"] {
1114            let error = QueryContext::try_from(wire(Some(value))).unwrap_err();
1115            assert!(matches!(
1116                error,
1117                crate::session_config::Error::InvalidConfigValue {
1118                    name,
1119                    value: actual,
1120                    ..
1121                } if name == "read_preference" && actual == value
1122            ));
1123        }
1124    }
1125
1126    #[test]
1127    fn test_legacy_query_context_wire_defaults_to_leader() {
1128        use prost::Message;
1129
1130        let legacy = api::v1::QueryContext::decode(
1131            include_bytes!(
1132                "../../../tests/compatibility/wire/read_preference/legacy_query_context.pb"
1133            )
1134            .as_slice(),
1135        )
1136        .unwrap();
1137        assert_eq!(legacy.current_catalog, "c1");
1138        assert_eq!(legacy.current_schema, "s1");
1139        assert_eq!(legacy.timezone, "UTC");
1140        assert_eq!(legacy.channel, Channel::Internal as u32);
1141        assert_eq!(
1142            legacy
1143                .extensions
1144                .get("flow.return_region_seq")
1145                .map(String::as_str),
1146            Some("true")
1147        );
1148        assert!(
1149            !legacy
1150                .extensions
1151                .contains_key(READ_PREFERENCE_EXTENSION_KEY)
1152        );
1153
1154        let restored = QueryContext::try_from(legacy).unwrap();
1155        assert_eq!(restored.read_preference(), ReadPreference::Leader);
1156        assert_eq!(restored.current_catalog(), "c1");
1157        assert_eq!(restored.current_schema(), "s1");
1158        assert_eq!(restored.timezone().to_string(), "UTC");
1159        assert_eq!(restored.channel(), Channel::Internal);
1160        assert_eq!(restored.extension("flow.return_region_seq"), Some("true"));
1161
1162        let outbound: api::v1::QueryContext = restored.into();
1163        let roundtrip = api::v1::QueryContext::decode(outbound.encode_to_vec().as_slice()).unwrap();
1164        assert_eq!(roundtrip.current_catalog, "c1");
1165        assert_eq!(roundtrip.current_schema, "s1");
1166        assert_eq!(roundtrip.timezone, "UTC");
1167        assert_eq!(roundtrip.channel, Channel::Internal as u32);
1168        assert_eq!(
1169            roundtrip
1170                .extensions
1171                .get("flow.return_region_seq")
1172                .map(String::as_str),
1173            Some("true")
1174        );
1175        assert!(
1176            !roundtrip
1177                .extensions
1178                .contains_key(READ_PREFERENCE_EXTENSION_KEY)
1179        );
1180    }
1181
1182    #[test]
1183    fn test_outbound_read_preference_clears_stale_extension() {
1184        let context = QueryContextBuilder::default()
1185            .set_extension(
1186                READ_PREFERENCE_EXTENSION_KEY.to_string(),
1187                "stale".to_string(),
1188            )
1189            .set_extension("unrelated".to_string(), "preserved".to_string())
1190            .build();
1191        let wire: api::v1::QueryContext = context.into();
1192        assert!(!wire.extensions.contains_key(READ_PREFERENCE_EXTENSION_KEY));
1193        assert_eq!(
1194            wire.extensions.get("unrelated").map(String::as_str),
1195            Some("preserved")
1196        );
1197    }
1198
1199    #[test]
1200    fn test_live_analyze_metrics_requires_matching_remote_query_id() {
1201        let mut ctx = QueryContext::arc().as_ref().clone();
1202        assert!(!ctx.live_analyze_metrics_enabled());
1203
1204        ctx.enable_live_analyze_metrics();
1205        assert!(ctx.live_analyze_metrics_enabled());
1206
1207        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "true");
1208        assert!(!ctx.live_analyze_metrics_enabled());
1209
1210        ctx.set_extension(LIVE_ANALYZE_METRICS_EXTENSION_KEY, "another-query-id");
1211        assert!(!ctx.live_analyze_metrics_enabled());
1212    }
1213    #[test]
1214    fn test_batching_selection_is_local_only() {
1215        let mut ctx = QueryContextBuilder::default().build();
1216        assert!(!ctx.batching_enabled());
1217        assert!(!ctx.logical_batching_enabled());
1218        ctx.set_logical_batching_enabled(true);
1219        assert!(ctx.clone().logical_batching_enabled());
1220        assert!(ctx.fork().logical_batching_enabled());
1221        ctx.set_batching_enabled(true);
1222        assert!(ctx.clone().batching_enabled());
1223        assert!(ctx.fork().batching_enabled());
1224        let wire: api::v1::QueryContext = ctx.into();
1225        let restored = QueryContext::try_from(wire).unwrap();
1226        assert!(!restored.batching_enabled());
1227        assert!(!restored.logical_batching_enabled());
1228        let ctx = QueryContextBuilder::default()
1229            .set_extension("batching_enabled".to_string(), "true".to_string())
1230            .build();
1231        assert!(!ctx.batching_enabled());
1232    }
1233}
session/context.rs

session/
context.rs
