chore: change the short_sql from 64bytes to 10KB (#15554)

Author: BohuTANG

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

@@ -583,18 +583,21 @@ impl Drop for QueryContextShared {
 pub fn short_sql(sql: String) -> String {
     use unicode_segmentation::UnicodeSegmentation;
     let query = sql.trim_start();
-    if query.as_bytes().len() >= 64 && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT") {
-        // keep first 64 graphemes
-        String::from_utf8(
-            query
-                .graphemes(true)
-                .take(64)
-                .flat_map(|g| g.as_bytes().iter())
-                .copied() // copied converts &u8 into u8
-                .chain(b"...".iter().copied())
-                .collect::<Vec<u8>>(),
-        )
-        .unwrap() // by construction, this cannot panic as we extracted unicode grapheme
+    if query.as_bytes().len() > 10240 && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT") {
+        // keep first 10KB (10,240 bytes)
+        let mut result = Vec::new();
+        let mut bytes_taken = 0;
+        for grapheme in query.graphemes(true) {
+            let grapheme_bytes = grapheme.as_bytes();
+            if bytes_taken + grapheme_bytes.len() <= 10240 {
+                result.extend_from_slice(grapheme_bytes);
+                bytes_taken += grapheme_bytes.len();
+            } else {
+                break;
+            }
+        }
+        result.extend_from_slice(b"...");
+        String::from_utf8(result).unwrap() // by construction, this cannot panic as we extracted unicode graphemes
     } else {
         sql
     }

src/query/service/tests/it/sessions/query_ctx.rs

@@ -17,6 +17,7 @@ use databend_common_exception::Result;
 use databend_common_meta_app::storage::StorageFsConfig;
 use databend_common_meta_app::storage::StorageParams;
 use databend_common_meta_app::storage::StorageS3Config;
+use databend_query::sessions::short_sql;
 use databend_query::sessions::TableContext;
 use databend_query::test_kits::ConfigBuilder;
 use databend_query::test_kits::TestFixture;
@@ -65,3 +66,29 @@ async fn test_get_storage_accessor_fs() -> Result<()> {
 
     Ok(())
 }
+
+#[test]
+fn test_short_sql() {
+    // Test case 1: SQL query shorter than 10KB
+    let sql1 = "SELECT * FROM users WHERE id = 1;".to_string();
+    assert_eq!(short_sql(sql1.clone()), sql1);
+
+    // Test case 2: SQL query longer than 10KB and starts with "INSERT"
+    let long_sql = "INSERT INTO users (id, name, email) VALUES ".to_string()
+        + &"(1, 'John Doe', 'john@example.com'), ".repeat(1000);
+    let expected_result = long_sql.as_bytes()[..10240].to_vec();
+    let expected_result = String::from_utf8(expected_result).unwrap() + "...";
+    assert_eq!(short_sql(long_sql), expected_result);
+
+    // Test case 3: SQL query longer than 10KB but does not start with "INSERT"
+    let long_sql = "SELECT * FROM users WHERE ".to_string() + &"id = 1 OR ".repeat(1000);
+    assert_eq!(short_sql(long_sql.clone()), long_sql);
+
+    // Test case 4: Empty SQL query
+    let empty_sql = "".to_string();
+    assert_eq!(short_sql(empty_sql.clone()), empty_sql);
+
+    // Test case 5: SQL query with leading whitespace
+    let sql_with_whitespace = "   SELECT * FROM users;".to_string();
+    assert_eq!(short_sql(sql_with_whitespace.clone()), sql_with_whitespace);
+}