feat: short sql length setting (databendlabs#16502)

feature: short sql length setting

src/common/base/src/base/string.rs

@@ -193,8 +193,7 @@ pub fn mask_connection_info(sql: &str) -> String {
 /// Maximum length of the SQL query to be displayed or log.
 /// If the query exceeds this length and starts with keywords,
 /// it will be truncated and appended with the remaining length.
-pub fn short_sql(sql: String) -> String {
-    const MAX_LENGTH: usize = 128;
+pub fn short_sql(sql: String, max_length: u64) -> String {
     let keywords = ["INSERT"];
 
     fn starts_with_any(query: &str, keywords: &[&str]) -> bool {
@@ -209,10 +208,10 @@ pub fn short_sql(sql: String) -> String {
     // of multiple Unicode code points.
     // This ensures that we handle complex characters like emojis or
     // accented characters properly.
-    if query.graphemes(true).count() > MAX_LENGTH && starts_with_any(query, &keywords) {
-        let truncated: String = query.graphemes(true).take(MAX_LENGTH).collect();
+    if query.graphemes(true).count() > max_length as usize && starts_with_any(query, &keywords) {
+        let truncated: String = query.graphemes(true).take(max_length as usize).collect();
         let original_length = query.graphemes(true).count();
-        let remaining_length = original_length.saturating_sub(MAX_LENGTH);
+        let remaining_length = original_length.saturating_sub(max_length as usize);
         // Append the remaining length indicator
         truncated + &format!("...[{} more characters]", remaining_length)
     } else {

src/common/base/tests/it/string.rs

@@ -111,9 +111,10 @@ fn test_mask_connection_info() {
 
 #[test]
 fn test_short_sql() {
+    let max_length: u64 = 128;
     // Test case 1: SQL query shorter than 128 characters
     let sql1 = "SELECT * FROM users WHERE id = 1;".to_string();
-    assert_eq!(short_sql(sql1.clone()), sql1);
+    assert_eq!(short_sql(sql1.clone(), max_length), sql1);
 
     // Test case 2: SQL query longer than 128 characters and starts with "INSERT"
     let long_sql_insert = "INSERT INTO users (id, name, email) VALUES ".to_string()
@@ -123,23 +124,32 @@ fn test_short_sql() {
         let truncated: String = long_sql_insert.graphemes(true).take(128).collect();
         truncated + &format!("...[{} more characters]", expected_length_insert)
     };
-    assert_eq!(short_sql(long_sql_insert.clone()), expected_result_insert);
+    assert_eq!(
+        short_sql(long_sql_insert.clone(), max_length),
+        expected_result_insert
+    );
 
     // Test case 3: SQL query longer than 128 characters but does not start with "INSERT"
     let long_sql_update =
         "UPDATE users SET name = 'John' WHERE id = 1;".to_string() + &"id = 1 OR ".repeat(20); // Make sure this creates a string longer than 128 characters
-    assert_eq!(short_sql(long_sql_update.clone()), long_sql_update);
+    assert_eq!(
+        short_sql(long_sql_update.clone(), max_length),
+        long_sql_update
+    );
 
     // Test case 4: Empty SQL query
     let empty_sql = "".to_string();
-    assert_eq!(short_sql(empty_sql.clone()), empty_sql);
+    assert_eq!(short_sql(empty_sql.clone(), max_length), empty_sql);
 
     // Test case 5: SQL query with leading whitespace
     let sql_with_whitespace =
         "   INSERT INTO users (id, name, email) VALUES (1, 'John Doe', '[email protected]');"
             .to_string();
     let trimmed_sql = sql_with_whitespace.trim_start().to_string();
-    assert_eq!(short_sql(sql_with_whitespace.clone()), trimmed_sql);
+    assert_eq!(
+        short_sql(sql_with_whitespace.clone(), max_length),
+        trimmed_sql
+    );
 
     // Test case 6: SQL query with multiple emojis to test truncation at an emoji point
     let emoji_sql = "INSERT INTO users (id, name) VALUES (1, 'John Doe 😊😊😊😊😊😊😊😊😊😊');"
@@ -150,5 +160,8 @@ fn test_short_sql() {
         let remaining_length = emoji_sql.graphemes(true).count().saturating_sub(128);
         truncated + &format!("...[{} more characters]", remaining_length)
     };
-    assert_eq!(short_sql(emoji_sql.clone()), expected_emoji_result);
+    assert_eq!(
+        short_sql(emoji_sql.clone(), max_length),
+        expected_emoji_result
+    );
 }

src/query/service/src/interpreters/interpreter.rs

@@ -207,7 +207,10 @@ pub async fn interpreter_plan_sql(ctx: Arc<QueryContext>, sql: &str) -> Result<(
         Arc::new(ServiceQueryExecutor::new(ctx.clone())),
     );
     let result = planner.plan_sql(sql).await;
-    let short_sql = short_sql(sql.to_string());
+    let short_sql = short_sql(
+        sql.to_string(),
+        ctx.get_settings().get_short_sql_max_length()?,
+    );
     let mut stmt = if let Ok((_, extras)) = &result {
         Some(extras.statement.clone())
     } else {

src/query/service/src/servers/http/clickhouse_handler.rs

@@ -330,7 +330,16 @@ pub async fn clickhouse_handler_post(
         // other parts of the request already logged in middleware
         let len = sql.len();
         let msg = if len > n {
-            format!("{}...(omit {} bytes)", short_sql(sql.clone()), len - n)
+            format!(
+                "{}...(omit {} bytes)",
+                short_sql(
+                    sql.clone(),
+                    ctx.get_settings()
+                        .get_short_sql_max_length()
+                        .unwrap_or(1000)
+                ),
+                len - n
+            )
         } else {
             sql.to_string()
         };

src/query/service/src/servers/http/v1/query/http_query.rs

@@ -130,7 +130,7 @@ impl Debug for HttpQueryRequest {
         f.debug_struct("HttpQueryRequest")
             .field("session_id", &self.session_id)
             .field("session", &self.session)
-            .field("sql", &short_sql(self.sql.clone()))
+            .field("sql", &short_sql(self.sql.clone(), 1000))
             .field("pagination", &self.pagination)
             .field("string_fields", &self.string_fields)
             .field("stage_attachment", &self.stage_attachment)

src/query/service/src/sessions/query_ctx_shared.rs

@@ -487,7 +487,12 @@ impl QueryContextShared {
     pub fn attach_query_str(&self, kind: QueryKind, query: String) {
         {
             let mut running_query = self.running_query.write();
-            *running_query = Some(short_sql(query));
+            *running_query = Some(short_sql(
+                query,
+                self.get_settings()
+                    .get_short_sql_max_length()
+                    .unwrap_or(1000),
+            ));
         }
 
         {

src/query/settings/src/settings_default.rs

@@ -909,6 +909,12 @@ impl DefaultSettings {
                     mode: SettingMode::Both,
                     range: Some(SettingRange::Numeric(0..=u64::MAX)),
                 }),
+                ("short_sql_max_length", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(128),
+                    desc: "Sets the maximum length for truncating SQL queries in short_sql function.",
+                    mode: SettingMode::Both,
+                    range: Some(SettingRange::Numeric(1..=1024*1024)),
+                }),
             ]);
 
             Ok(Arc::new(DefaultSettings {

src/query/settings/src/settings_getter_setter.rs

@@ -748,4 +748,12 @@ impl Settings {
     pub fn get_max_spill_io_requests(&self) -> Result<u64> {
         self.try_get_u64("max_spill_io_requests")
     }
+
+    pub fn get_short_sql_max_length(&self) -> Result<u64> {
+        self.try_get_u64("short_sql_max_length")
+    }
+
+    pub fn set_short_sql_max_length(&self, val: u64) -> Result<()> {
+        self.try_set_u64("short_sql_max_length", val)
+    }
 }

src/query/settings/tests/it/setting.rs

@@ -98,6 +98,25 @@ async fn test_set_settings() {
         let expect = "WrongValueForVariable. Code: 2803, Text = Value xx is not within the allowed values [\"None\", \"LZ4\", \"ZSTD\"].";
         assert_eq!(expect, format!("{}", result.unwrap_err()));
     }
+
+    // Number Range
+    {
+        // Ok
+        settings
+            .set_setting("short_sql_max_length".to_string(), "1000".to_string())
+            .unwrap();
+
+        // Range 1024*1024
+        let result =
+            settings.set_setting("short_sql_max_length".to_string(), "1048577".to_string());
+        let expect = "WrongValueForVariable. Code: 2803, Text = Value 1048577 is not within the range [1, 1048576].";
+        assert_eq!(expect, format!("{}", result.unwrap_err()));
+
+        // Range 1
+        let result = settings.set_setting("short_sql_max_length".to_string(), "0".to_string());
+        let expect = "WrongValueForVariable. Code: 2803, Text = Value 0 is not within the range [1, 1048576].";
+        assert_eq!(expect, format!("{}", result.unwrap_err()));
+    }
 }
 
 #[tokio::test(flavor = "multi_thread", worker_threads = 1)]