Skip to content

Commit 50a7f9b

Browse files
authored
chore: change the short_sql from 64bytes to 10KB (#15554)
1 parent c4d8b00 commit 50a7f9b

File tree

2 files changed

+42
-12
lines changed

2 files changed

+42
-12
lines changed

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

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -583,18 +583,21 @@ impl Drop for QueryContextShared {
583583
pub fn short_sql(sql: String) -> String {
584584
use unicode_segmentation::UnicodeSegmentation;
585585
let query = sql.trim_start();
586-
if query.as_bytes().len() >= 64 && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT") {
587-
// keep first 64 graphemes
588-
String::from_utf8(
589-
query
590-
.graphemes(true)
591-
.take(64)
592-
.flat_map(|g| g.as_bytes().iter())
593-
.copied() // copied converts &u8 into u8
594-
.chain(b"...".iter().copied())
595-
.collect::<Vec<u8>>(),
596-
)
597-
.unwrap() // by construction, this cannot panic as we extracted unicode grapheme
586+
if query.as_bytes().len() > 10240 && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT") {
587+
// keep first 10KB (10,240 bytes)
588+
let mut result = Vec::new();
589+
let mut bytes_taken = 0;
590+
for grapheme in query.graphemes(true) {
591+
let grapheme_bytes = grapheme.as_bytes();
592+
if bytes_taken + grapheme_bytes.len() <= 10240 {
593+
result.extend_from_slice(grapheme_bytes);
594+
bytes_taken += grapheme_bytes.len();
595+
} else {
596+
break;
597+
}
598+
}
599+
result.extend_from_slice(b"...");
600+
String::from_utf8(result).unwrap() // by construction, this cannot panic as we extracted unicode graphemes
598601
} else {
599602
sql
600603
}

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

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ use databend_common_exception::Result;
1717
use databend_common_meta_app::storage::StorageFsConfig;
1818
use databend_common_meta_app::storage::StorageParams;
1919
use databend_common_meta_app::storage::StorageS3Config;
20+
use databend_query::sessions::short_sql;
2021
use databend_query::sessions::TableContext;
2122
use databend_query::test_kits::ConfigBuilder;
2223
use databend_query::test_kits::TestFixture;
@@ -65,3 +66,29 @@ async fn test_get_storage_accessor_fs() -> Result<()> {
6566

6667
Ok(())
6768
}
69+
70+
#[test]
71+
fn test_short_sql() {
72+
// Test case 1: SQL query shorter than 10KB
73+
let sql1 = "SELECT * FROM users WHERE id = 1;".to_string();
74+
assert_eq!(short_sql(sql1.clone()), sql1);
75+
76+
// Test case 2: SQL query longer than 10KB and starts with "INSERT"
77+
let long_sql = "INSERT INTO users (id, name, email) VALUES ".to_string()
78+
+ &"(1, 'John Doe', 'john@example.com'), ".repeat(1000);
79+
let expected_result = long_sql.as_bytes()[..10240].to_vec();
80+
let expected_result = String::from_utf8(expected_result).unwrap() + "...";
81+
assert_eq!(short_sql(long_sql), expected_result);
82+
83+
// Test case 3: SQL query longer than 10KB but does not start with "INSERT"
84+
let long_sql = "SELECT * FROM users WHERE ".to_string() + &"id = 1 OR ".repeat(1000);
85+
assert_eq!(short_sql(long_sql.clone()), long_sql);
86+
87+
// Test case 4: Empty SQL query
88+
let empty_sql = "".to_string();
89+
assert_eq!(short_sql(empty_sql.clone()), empty_sql);
90+
91+
// Test case 5: SQL query with leading whitespace
92+
let sql_with_whitespace = " SELECT * FROM users;".to_string();
93+
assert_eq!(short_sql(sql_with_whitespace.clone()), sql_with_whitespace);
94+
}

0 commit comments

Comments
 (0)