feat: sql dialect for different protocols (#1631)

* feat: add SqlDialect to query context

* feat: use session in postgrel handlers

* chore: refactor sql dialect

* feat: use different dialects for different sql protocols

* feat: adds GreptimeDbDialect

* refactor: replace GenericDialect with GreptimeDbDialect

* feat: save user info to session

* fix: compile error

* fix: test

src/session/Cargo.toml

@@ -9,3 +9,4 @@ arc-swap = "1.5"
 common-catalog = { path = "../common/catalog" }
 common-telemetry = { path = "../common/telemetry" }
 common-time = { path = "../common/time" }
+sql = { path = "../sql" }

src/session/src/context.rs

@@ -21,6 +21,7 @@ use common_catalog::build_db_string;
 use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 use common_telemetry::debug;
 use common_time::TimeZone;
+use sql::dialect::{Dialect, GreptimeDbDialect, MySqlDialect, PostgreSqlDialect};
 
 pub type QueryContextRef = Arc<QueryContext>;
 pub type ConnInfoRef = Arc<ConnInfo>;
@@ -30,6 +31,7 @@ pub struct QueryContext {
     current_catalog: ArcSwap<String>,
     current_schema: ArcSwap<String>,
     time_zone: ArcSwap<Option<TimeZone>>,
+    sql_dialect: Box<dyn Dialect + Send + Sync>,
 }
 
 impl Default for QueryContext {
@@ -59,25 +61,42 @@ impl QueryContext {
             current_catalog: ArcSwap::new(Arc::new(DEFAULT_CATALOG_NAME.to_string())),
             current_schema: ArcSwap::new(Arc::new(DEFAULT_SCHEMA_NAME.to_string())),
             time_zone: ArcSwap::new(Arc::new(None)),
+            sql_dialect: Box::new(GreptimeDbDialect {}),
         }
     }
 
     pub fn with(catalog: &str, schema: &str) -> Self {
+        Self::with_sql_dialect(catalog, schema, Box::new(GreptimeDbDialect {}))
+    }
+
+    pub fn with_sql_dialect(
+        catalog: &str,
+        schema: &str,
+        sql_dialect: Box<dyn Dialect + Send + Sync>,
+    ) -> Self {
         Self {
             current_catalog: ArcSwap::new(Arc::new(catalog.to_string())),
             current_schema: ArcSwap::new(Arc::new(schema.to_string())),
             time_zone: ArcSwap::new(Arc::new(None)),
+            sql_dialect,
         }
     }
 
+    #[inline]
     pub fn current_schema(&self) -> String {
         self.current_schema.load().as_ref().clone()
     }
 
+    #[inline]
     pub fn current_catalog(&self) -> String {
         self.current_catalog.load().as_ref().clone()
     }
 
+    #[inline]
+    pub fn sql_dialect(&self) -> &(dyn Dialect + Send + Sync) {
+        &*self.sql_dialect
+    }
+
     pub fn set_current_schema(&self, schema: &str) {
         let last = self.current_schema.swap(Arc::new(schema.to_string()));
         if schema != last.as_str() {
@@ -142,15 +161,30 @@ impl UserInfo {
     }
 }
 
+#[derive(Debug)]
 pub struct ConnInfo {
-    pub client_host: SocketAddr,
+    pub client_addr: Option<SocketAddr>,
     pub channel: Channel,
 }
 
+impl std::fmt::Display for ConnInfo {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        write!(
+            f,
+            "{}[{}]",
+            self.channel,
+            self.client_addr
+                .map(|addr| addr.to_string())
+                .as_deref()
+                .unwrap_or("unknown client addr")
+        )
+    }
+}
+
 impl ConnInfo {
-    pub fn new(client_host: SocketAddr, channel: Channel) -> Self {
+    pub fn new(client_addr: Option<SocketAddr>, channel: Channel) -> Self {
         Self {
-            client_host,
+            client_addr,
             channel,
         }
     }
@@ -158,13 +192,26 @@ impl ConnInfo {
 
 #[derive(Debug, PartialEq)]
 pub enum Channel {
-    Grpc,
-    Http,
     Mysql,
     Postgres,
-    Opentsdb,
-    Influxdb,
-    Prometheus,
+}
+
+impl Channel {
+    pub fn dialect(&self) -> Box<dyn Dialect + Send + Sync> {
+        match self {
+            Channel::Mysql => Box::new(MySqlDialect {}),
+            Channel::Postgres => Box::new(PostgreSqlDialect {}),
+        }
+    }
+}
+
+impl std::fmt::Display for Channel {
+    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
+        match self {
+            Channel::Mysql => write!(f, "mysql"),
+            Channel::Postgres => write!(f, "postgres"),
+        }
+    }
 }
 
 #[cfg(test)]
@@ -175,7 +222,7 @@ mod test {
 
     #[test]
     fn test_session() {
-        let session = Session::new("127.0.0.1:9000".parse().unwrap(), Channel::Mysql);
+        let session = Session::new(Some("127.0.0.1:9000".parse().unwrap()), Channel::Mysql);
         // test user_info
         assert_eq!(session.user_info().username(), "greptime");
         session.set_user_info(UserInfo::new("root"));
@@ -183,11 +230,11 @@ mod test {
 
         // test channel
         assert_eq!(session.conn_info().channel, Channel::Mysql);
-        assert_eq!(
-            session.conn_info().client_host.ip().to_string(),
-            "127.0.0.1"
-        );
-        assert_eq!(session.conn_info().client_host.port(), 9000);
+        let client_addr = session.conn_info().client_addr.as_ref().unwrap();
+        assert_eq!(client_addr.ip().to_string(), "127.0.0.1");
+        assert_eq!(client_addr.port(), 9000);
+
+        assert_eq!("mysql[127.0.0.1:9000]", session.conn_info().to_string());
     }
 
     #[test]

src/session/src/lib.rs

@@ -18,33 +18,54 @@ use std::net::SocketAddr;
 use std::sync::Arc;
 
 use arc_swap::ArcSwap;
+use common_catalog::consts::{DEFAULT_CATALOG_NAME, DEFAULT_SCHEMA_NAME};
 
-use crate::context::{Channel, ConnInfo, ConnInfoRef, QueryContext, QueryContextRef, UserInfo};
+use crate::context::{Channel, ConnInfo, QueryContext, QueryContextRef, UserInfo};
 
+/// Session for persistent connection such as MySQL, PostgreSQL etc.
+#[derive(Debug)]
 pub struct Session {
     query_ctx: QueryContextRef,
     user_info: ArcSwap<UserInfo>,
-    conn_info: ConnInfoRef,
+    conn_info: ConnInfo,
 }
 
+pub type SessionRef = Arc<Session>;
+
 impl Session {
-    pub fn new(addr: SocketAddr, channel: Channel) -> Self {
+    pub fn new(addr: Option<SocketAddr>, channel: Channel) -> Self {
         Session {
-            query_ctx: Arc::new(QueryContext::new()),
+            query_ctx: Arc::new(QueryContext::with_sql_dialect(
+                DEFAULT_CATALOG_NAME,
+                DEFAULT_SCHEMA_NAME,
+                channel.dialect(),
+            )),
             user_info: ArcSwap::new(Arc::new(UserInfo::default())),
-            conn_info: Arc::new(ConnInfo::new(addr, channel)),
+            conn_info: ConnInfo::new(addr, channel),
         }
     }
 
+    #[inline]
     pub fn context(&self) -> QueryContextRef {
         self.query_ctx.clone()
     }
-    pub fn conn_info(&self) -> ConnInfoRef {
-        self.conn_info.clone()
+
+    #[inline]
+    pub fn conn_info(&self) -> &ConnInfo {
+        &self.conn_info
     }
+
+    #[inline]
+    pub fn mut_conn_info(&mut self) -> &mut ConnInfo {
+        &mut self.conn_info
+    }
+
+    #[inline]
     pub fn user_info(&self) -> Arc<UserInfo> {
         self.user_info.load().clone()
     }
+
+    #[inline]
     pub fn set_user_info(&self, user_info: UserInfo) {
         self.user_info.store(Arc::new(user_info));
     }