chore(query): add new session setting: query_tag (#17430)

set session query_tag=query-tag1;

all query in this session will be recored into system table query_log.

Author: TCeason

src/query/service/src/interpreters/common/query_log.rs

@@ -143,6 +143,11 @@ impl InterpreterQueryLog {
                 .expect("write to string must succeed");
         }
 
+        let query_tag = if let Ok(tag) = current_session.get_settings().get_query_tag() {
+            tag
+        } else {
+            "".to_string()
+        };
         session_settings.push_str("scope: SESSION");
 
         // Error
@@ -213,6 +218,7 @@ impl InterpreterQueryLog {
             exception_text,
             stack_trace,
             server_version: DATABEND_COMMIT_VERSION.to_string(),
+            query_tag,
             session_settings,
             extra: "".to_string(),
             has_profiles: false,
@@ -309,6 +315,12 @@ impl InterpreterQueryLog {
                 .expect("write to string must succeed");
         }
 
+        // Session
+        let query_tag = if let Ok(tag) = current_session.get_settings().get_query_tag() {
+            tag
+        } else {
+            "".to_string()
+        };
         session_settings.push_str("scope: SESSION");
 
         // Error
@@ -380,6 +392,7 @@ impl InterpreterQueryLog {
             exception_text,
             stack_trace,
             server_version: DATABEND_COMMIT_VERSION.to_string(),
+            query_tag,
             session_settings,
             extra: "".to_string(),
             has_profiles,

src/query/service/tests/it/storages/testdata/columns_table.txt

@@ -368,6 +368,7 @@ DB.Table: 'system'.'columns', Table: columns-table_id:1, ver:0, Engine: SystemCo
 | 'query_parameterized_hash'        | 'system'             | 'query_log'              | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'query_queued_duration_ms'        | 'system'             | 'query_log'              | 'Int64'               | 'BIGINT'            | ''       | ''       | 'NO'     | ''       |
 | 'query_start_time'                | 'system'             | 'query_log'              | 'Timestamp'           | 'TIMESTAMP'         | ''       | ''       | 'NO'     | ''       |
+| 'query_tag'                       | 'system'             | 'query_log'              | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'query_text'                      | 'system'             | 'query_log'              | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'range'                           | 'system'             | 'settings'               | 'String'              | 'VARCHAR'           | ''       | ''       | 'NO'     | ''       |
 | 'referenced_column_name'          | 'information_schema' | 'key_column_usage'       | 'NULL'                | 'NULL'              | ''       | ''       | 'NO'     | ''       |

src/query/settings/src/settings_default.rs

@@ -296,6 +296,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::String(vec!["PostgreSQL".into(), "MySQL".into(), "Experimental".into(), "Hive".into(), "Prql".into()])),
                 }),
+                ("query_tag", DefaultSettingValue {
+                    value: UserSettingValue::String("".to_owned()),
+                    desc: "Sets the query tag for this session.",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Session,
+                    range: None,
+                }),
                 ("enable_dphyp", DefaultSettingValue {
                     value: UserSettingValue::UInt64(1),
                     desc: "Enables dphyp join order algorithm.",

src/query/settings/src/settings_getter_setter.rs

@@ -412,6 +412,10 @@ impl Settings {
         self.try_get_string("sandbox_tenant")
     }
 
+    pub fn get_query_tag(&self) -> Result<String> {
+        self.try_get_string("query_tag")
+    }
+
     pub fn get_hide_options_in_show_create_table(&self) -> Result<bool> {
         Ok(self.try_get_u64("hide_options_in_show_create_table")? != 0)
     }

src/query/storages/system/src/query_log_table.rs

@@ -164,7 +164,8 @@ pub struct QueryLogElement {
     // Server.
     pub server_version: String,
 
-    // Session settings
+    // Session
+    pub query_tag: String,
     #[serde(skip_serializing)]
     pub session_settings: String,
 
@@ -311,7 +312,8 @@ impl SystemLogElement for QueryLogElement {
             TableField::new("stack_trace", TableDataType::String),
             // Server.
             TableField::new("server_version", TableDataType::String),
-            // Session settings
+            // Session
+            TableField::new("query_tag", TableDataType::String),
             TableField::new("session_settings", TableDataType::String),
             // Extra.
             TableField::new("extra", TableDataType::String),
@@ -548,7 +550,11 @@ impl SystemLogElement for QueryLogElement {
             .next()
             .unwrap()
             .push(Scalar::String(self.server_version.clone()).as_ref());
-        // Session settings
+        // Session
+        columns
+            .next()
+            .unwrap()
+            .push(Scalar::String(self.query_tag.clone()).as_ref());
         columns
             .next()
             .unwrap()

tests/sqllogictests/suites/base/06_show/06_0003_show_settings.test

@@ -1,3 +1,6 @@
+statement error 1001
+set global query_tag = 'x';
+
 statement ok
 unset max_threads;
 

tests/sqllogictests/suites/mode/cluster/explain_v2.test

@@ -428,20 +428,20 @@ query T
 EXPLAIN SELECT a.cluster_node, b.query_node FROM (SELECT name as cluster_node FROM system.clusters) AS a LEFT JOIN (SELECT DISTINCT node_id as query_node FROM system.query_log) AS b ON a.cluster_node = b.query_node
 ----
 Exchange
-├── output columns: [clusters.name (#0), query_node (#62)]
+├── output columns: [clusters.name (#0), query_node (#63)]
 ├── exchange type: Merge
 └── HashJoin
-    ├── output columns: [clusters.name (#0), query_node (#62)]
+    ├── output columns: [clusters.name (#0), query_node (#63)]
     ├── join type: LEFT OUTER
-    ├── build keys: [b.query_node (#62)]
+    ├── build keys: [b.query_node (#63)]
     ├── probe keys: [CAST(a.cluster_node (#0) AS String NULL)]
     ├── filters: []
     ├── estimated rows: 0.00
     ├── Exchange(Build)
-    │   ├── output columns: [query_node (#62)]
+    │   ├── output columns: [query_node (#63)]
     │   ├── exchange type: Broadcast
     │   └── EvalScalar
-    │       ├── output columns: [query_node (#62)]
+    │       ├── output columns: [query_node (#63)]
     │       ├── expressions: [CAST(b.query_node (#10) AS String NULL)]
     │       ├── estimated rows: 0.00
     │       └── AggregateFinal

tests/suites/0_stateless/20+_others/20_0024_session_query_tag.result

@@ -0,0 +1,4 @@
+1
+2
+true
+true

tests/suites/0_stateless/20+_others/20_0024_session_query_tag.sh

@@ -0,0 +1,12 @@
+#!/usr/bin/env bash
+CURDIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)
+. "$CURDIR"/../../../shell_env.sh
+
+$BENDSQL_CLIENT_CONNECT --query="set query_tag = 'test-query_tag'; select 1; create database if not exists c;"
+$BENDSQL_CLIENT_CONNECT --query="set query_tag = 'test-query_tag-2'; select 2; create database if not exists b;"
+
+# count % 5 = 0 because (set command + select start/finished + create start/finished) must be a positive multiple of 5
+$BENDSQL_CLIENT_CONNECT --query="select count(query_text) % 5 = 0 from system.query_log group by query_tag having query_tag='test-query_tag'"
+$BENDSQL_CLIENT_CONNECT --query="select count(query_text) % 5 = 0 from system.query_log group by query_tag having query_tag='test-query_tag-2'"
+
+$BENDSQL_CLIENT_CONNECT --query="drop database if exists b;drop database if exists c;"