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

session/
lib.rs
