diff --git a/src/common/base/src/base/string.rs b/src/common/base/src/base/string.rs index e240e2630d13..0b462d0d4c93 100644 --- a/src/common/base/src/base/string.rs +++ b/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 { diff --git a/src/common/base/tests/it/string.rs b/src/common/base/tests/it/string.rs index 4b49c57f2673..5cdf951f7b4d 100644 --- a/src/common/base/tests/it/string.rs +++ b/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', 'john@example.com');" .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 + ); } diff --git a/src/query/service/src/interpreters/interpreter.rs b/src/query/service/src/interpreters/interpreter.rs index fa1c9e311377..c74892dc32a0 100644 --- a/src/query/service/src/interpreters/interpreter.rs +++ b/src/query/service/src/interpreters/interpreter.rs @@ -207,7 +207,10 @@ pub async fn interpreter_plan_sql(ctx: Arc, 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 { diff --git a/src/query/service/src/servers/http/clickhouse_handler.rs b/src/query/service/src/servers/http/clickhouse_handler.rs index 8020159f1d14..22f41136144f 100644 --- a/src/query/service/src/servers/http/clickhouse_handler.rs +++ b/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() }; diff --git a/src/query/service/src/servers/http/v1/query/http_query.rs b/src/query/service/src/servers/http/v1/query/http_query.rs index 8641e665392f..f714a72b9ba0 100644 --- a/src/query/service/src/servers/http/v1/query/http_query.rs +++ b/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) diff --git a/src/query/service/src/sessions/query_ctx_shared.rs b/src/query/service/src/sessions/query_ctx_shared.rs index 55c1f44bdd7a..cb073e69e772 100644 --- a/src/query/service/src/sessions/query_ctx_shared.rs +++ b/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), + )); } { diff --git a/src/query/settings/src/settings_default.rs b/src/query/settings/src/settings_default.rs index bef9c5ae9dbc..a1a6ddb005f1 100644 --- a/src/query/settings/src/settings_default.rs +++ b/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 { diff --git a/src/query/settings/src/settings_getter_setter.rs b/src/query/settings/src/settings_getter_setter.rs index 521fa92c11d4..e54806e1550f 100644 --- a/src/query/settings/src/settings_getter_setter.rs +++ b/src/query/settings/src/settings_getter_setter.rs @@ -748,4 +748,12 @@ impl Settings { pub fn get_max_spill_io_requests(&self) -> Result { self.try_get_u64("max_spill_io_requests") } + + pub fn get_short_sql_max_length(&self) -> Result { + 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) + } } diff --git a/src/query/settings/tests/it/setting.rs b/src/query/settings/tests/it/setting.rs index b318034a0d79..f99634773add 100644 --- a/src/query/settings/tests/it/setting.rs +++ b/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)]