Skip to content

Commit

Permalink
chore: short_sql from 10KB to 30KB (#15575)
Browse files Browse the repository at this point in the history
  • Loading branch information
BohuTANG authored May 19, 2024
1 parent 6a89aa1 commit ac12ad1
Show file tree
Hide file tree
Showing 2 changed files with 12 additions and 9 deletions.
9 changes: 6 additions & 3 deletions src/query/service/src/sessions/query_ctx_shared.rs
Original file line number Diff line number Diff line change
Expand Up @@ -582,14 +582,17 @@ impl Drop for QueryContextShared {

pub fn short_sql(sql: String) -> String {
use unicode_segmentation::UnicodeSegmentation;
const MAX_LENGTH: usize = 30 * 1024; // 30KB

let query = sql.trim_start();
if query.as_bytes().len() > 10240 && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT") {
// keep first 10KB (10,240 bytes)
if query.as_bytes().len() > MAX_LENGTH && query.as_bytes()[..6].eq_ignore_ascii_case(b"INSERT")
{
// keep first 30KB
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 {
if bytes_taken + grapheme_bytes.len() <= MAX_LENGTH {
result.extend_from_slice(grapheme_bytes);
bytes_taken += grapheme_bytes.len();
} else {
Expand Down
12 changes: 6 additions & 6 deletions src/query/service/tests/it/sessions/query_ctx.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,19 +69,19 @@ async fn test_get_storage_accessor_fs() -> Result<()> {

#[test]
fn test_short_sql() {
// Test case 1: SQL query shorter than 10KB
// Test case 1: SQL query shorter than 30KB
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"
// Test case 2: SQL query longer than 30KB and starts with "INSERT"
let long_sql = "INSERT INTO users (id, name, email) VALUES ".to_string()
+ &"(1, 'John Doe', '[email protected]'), ".repeat(1000);
let expected_result = long_sql.as_bytes()[..10240].to_vec();
+ &"(1, 'John Doe', '[email protected]'), ".repeat(1500); // Adjusted for 30KB
let expected_result = long_sql.as_bytes()[..30 * 1024].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);
// Test case 3: SQL query longer than 30KB but does not start with "INSERT"
let long_sql = "SELECT * FROM users WHERE ".to_string() + &"id = 1 OR ".repeat(1500); // Adjusted for 30KB
assert_eq!(short_sql(long_sql.clone()), long_sql);

// Test case 4: Empty SQL query
Expand Down

0 comments on commit ac12ad1

Please sign in to comment.