session/
lib.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
15pub mod context;
16pub mod hints;
17pub mod protocol_ctx;
18pub mod query_id;
19pub mod query_options;
20pub mod session_config;
21pub mod table_name;
22
23use std::collections::{HashMap, VecDeque};
24use std::net::SocketAddr;
25use std::sync::{Arc, RwLock};
26use std::time::Duration;
27
28use auth::UserInfoRef;
29use common_catalog::build_db_string;
30use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
31use common_recordbatch::cursor::RecordBatchStreamCursor;
32pub use common_session::ReadPreference;
33use common_time::Timezone;
34use common_time::timezone::get_timezone;
35use context::{ConfigurationVariables, QueryContextBuilder};
36use derive_more::Debug;
37
38use crate::context::{Channel, ConnInfo, QueryContextRef, dialect_for_channel};
39
40/// Maximum number of warnings to store per session (similar to MySQL's max_error_count)
41const MAX_WARNINGS: usize = 64;
42
43/// Session for persistent connection such as MySQL, PostgreSQL etc.
44#[derive(Debug)]
45pub struct Session {
46    catalog: RwLock<String>,
47    mutable_inner: Arc<RwLock<MutableInner>>,
48    conn_info: ConnInfo,
49    configuration_variables: Arc<ConfigurationVariables>,
50    // the process id to use when killing the query
51    process_id: u32,
52}
53
54pub type SessionRef = Arc<Session>;
55
56/// A container for mutable items in query context
57#[derive(Debug, Clone)]
58pub(crate) struct MutableInner {
59    schema: String,
60    user_info: UserInfoRef,
61    timezone: Timezone,
62    query_timeout: Option<Duration>,
63    read_preference: ReadPreference,
64    /// Latest procedure ID recorded during statement execution, before
65    /// `post_execute` hooks and protocol transmission.
66    last_procedure_id: Option<String>,
67    /// Request-level WAL policy for ordinary inserts.
68    skip_wal: bool,
69    #[debug(skip)]
70    pub(crate) cursors: HashMap<String, Arc<RecordBatchStreamCursor>>,
71    /// Warning messages for MySQL SHOW WARNINGS support
72    warnings: VecDeque<String>,
73}
74
75impl Default for MutableInner {
76    fn default() -> Self {
77        Self {
78            schema: DEFAULT_SCHEMA_NAME.into(),
79            user_info: auth::userinfo_by_name(None),
80            timezone: get_timezone(None).clone(),
81            query_timeout: None,
82            read_preference: ReadPreference::Leader,
83            last_procedure_id: None,
84            skip_wal: false,
85            cursors: HashMap::with_capacity(0),
86            warnings: VecDeque::new(),
87        }
88    }
89}
90
91impl Session {
92    pub fn new(
93        addr: Option<SocketAddr>,
94        channel: Channel,
95        configuration_variables: ConfigurationVariables,
96        process_id: u32,
97    ) -> Self {
98        Session {
99            catalog: RwLock::new(DEFAULT_CATALOG_NAME.into()),
100            conn_info: ConnInfo::new(addr, channel),
101            configuration_variables: Arc::new(configuration_variables),
102            mutable_inner: Arc::new(RwLock::new(MutableInner::default())),
103            process_id,
104        }
105    }
106
107    pub fn new_query_context(&self) -> QueryContextRef {
108        QueryContextBuilder::default()
109            // catalog is not allowed for update in query context so we use
110            // string here
111            .current_catalog(self.catalog.read().unwrap().clone())
112            .mutable_session_data(self.mutable_inner.clone())
113            .sql_dialect(dialect_for_channel(self.conn_info.channel))
114            .configuration_parameter(self.configuration_variables.clone())
115            .channel(self.conn_info.channel)
116            .process_id(self.process_id)
117            .conn_info(self.conn_info.clone())
118            .build()
119            .into()
120    }
121
122    /// Cursors are shared across query contexts created from this session.
123    pub fn get_cursor(&self, name: &str) -> Option<Arc<RecordBatchStreamCursor>> {
124        let guard = self.mutable_inner.read().unwrap();
125        guard.cursors.get(name).cloned()
126    }
127
128    pub fn conn_info(&self) -> &ConnInfo {
129        &self.conn_info
130    }
131
132    pub fn timezone(&self) -> Timezone {
133        self.mutable_inner.read().unwrap().timezone.clone()
134    }
135
136    pub fn read_preference(&self) -> ReadPreference {
137        self.mutable_inner.read().unwrap().read_preference
138    }
139
140    pub fn set_timezone(&self, tz: Timezone) {
141        let mut inner = self.mutable_inner.write().unwrap();
142        inner.timezone = tz;
143    }
144
145    pub fn set_read_preference(&self, read_preference: ReadPreference) {
146        self.mutable_inner.write().unwrap().read_preference = read_preference;
147    }
148
149    pub fn user_info(&self) -> UserInfoRef {
150        self.mutable_inner.read().unwrap().user_info.clone()
151    }
152
153    pub fn set_user_info(&self, user_info: UserInfoRef) {
154        self.mutable_inner.write().unwrap().user_info = user_info;
155    }
156
157    pub fn set_catalog(&self, catalog: String) {
158        *self.catalog.write().unwrap() = catalog;
159    }
160
161    pub fn catalog(&self) -> String {
162        self.catalog.read().unwrap().clone()
163    }
164
165    pub fn schema(&self) -> String {
166        self.mutable_inner.read().unwrap().schema.clone()
167    }
168
169    pub fn set_schema(&self, schema: String) {
170        self.mutable_inner.write().unwrap().schema = schema;
171    }
172
173    pub fn get_db_string(&self) -> String {
174        build_db_string(&self.catalog(), &self.schema())
175    }
176
177    pub fn process_id(&self) -> u32 {
178        self.process_id
179    }
180
181    pub fn warnings_count(&self) -> usize {
182        self.mutable_inner.read().unwrap().warnings.len()
183    }
184
185    pub fn warnings(&self) -> Vec<String> {
186        self.mutable_inner
187            .read()
188            .unwrap()
189            .warnings
190            .iter()
191            .cloned()
192            .collect()
193    }
194
195    /// Add a warning message. If the limit is reached, discard the oldest warning.
196    pub fn add_warning(&self, warning: String) {
197        let mut inner = self.mutable_inner.write().unwrap();
198        if inner.warnings.len() >= MAX_WARNINGS {
199            inner.warnings.pop_front();
200        }
201        inner.warnings.push_back(warning);
202    }
203
204    pub fn clear_warnings(&self) {
205        let mut inner = self.mutable_inner.write().unwrap();
206        if inner.warnings.is_empty() {
207            return;
208        }
209        inner.warnings.clear();
210    }
211}
session/lib.rs

session/
lib.rs
