From 31596ea2d30a2ce6eaf4f2de3f4b51069fca4425 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 13 Oct 2023 15:43:35 -0700 Subject: [PATCH 1/6] Fixed `run_tests.sh` to properly handle the log level and also delete toxiproxy (if one exists). This helps with re-running the tests, even if it fails in the middle (making the tests execution idempotent). --- .circleci/run_tests.sh | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.circleci/run_tests.sh b/.circleci/run_tests.sh index 4ba497c3..3e076c06 100644 --- a/.circleci/run_tests.sh +++ b/.circleci/run_tests.sh @@ -10,7 +10,7 @@ set -o xtrace # for inspection. function start_pgcat() { kill -s SIGINT $(pgrep pgcat) || true - RUST_LOG=${1} ./target/debug/pgcat .circleci/pgcat.toml & + ./target/debug/pgcat .circleci/pgcat.toml --log-level ${1} & sleep 1 } @@ -29,6 +29,9 @@ PGPASSWORD=sharding_user pgbench -h 127.0.0.1 -U sharding_user shard2 -i LOG_LEVEL=error toxiproxy-server & sleep 1 +# Remove toxiproxy if it exists. +toxiproxy-cli delete postgres_replica || true + # Create a database at port 5433, forward it to Postgres toxiproxy-cli create -l 127.0.0.1:5433 -u 127.0.0.1:5432 postgres_replica From 5c7fcdcefb2ef634b6201cf0424485b15a704b27 Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Fri, 13 Oct 2023 15:44:09 -0700 Subject: [PATCH 2/6] Avoid generating warning on queries that are known not to be supported by `sqlparser`. --- src/client.rs | 50 ++++++++++++++++++++++++++++++--------------- src/query_router.rs | 21 +++++++++++++++++-- 2 files changed, 53 insertions(+), 18 deletions(-) diff --git a/src/client.rs b/src/client.rs index 98a0669c..4989f8e3 100644 --- a/src/client.rs +++ b/src/client.rs @@ -890,7 +890,7 @@ where 'Q' => { if query_router.query_parser_enabled() { - match query_router.parse(&message) { + initial_parsed_ast = Some(match query_router.parse(&message) { Ok(ast) => { let plugin_result = query_router.execute_plugins(&ast).await; @@ -910,15 +910,20 @@ where let _ = query_router.infer(&ast); - initial_parsed_ast = Some(ast); + ast } Err(error) => { - warn!( - "Query parsing error: {} (client: {})", - error, client_identifier - ); + if error != Error::UnsupportedStatement { + warn!( + "Query parsing error: {} (client: {})", + error, client_identifier + ); + } + // Returning an empty vector shows the failed parse attempt and + // avoids another round of parse attempt below. + vec![] } - } + }) } } @@ -940,10 +945,12 @@ where let _ = query_router.infer(&ast); } Err(error) => { - warn!( - "Query parsing error: {} (client: {})", - error, client_identifier - ); + if error != Error::UnsupportedStatement { + warn!( + "Query parsing error: {} (client: {})", + error, client_identifier + ); + } } }; } @@ -1273,14 +1280,25 @@ where if query_router.query_parser_enabled() { // We don't want to parse again if we already parsed it as the initial message let ast = match initial_parsed_ast { - Some(_) => Some(initial_parsed_ast.take().unwrap()), + Some(_) => { + let parsed_ast = initial_parsed_ast.take().unwrap(); + // if 'parsed_ast' is empty, it means that there was a failed + // attempt to parse the query as a custom command, earlier above. + if parsed_ast.is_empty() { + None + } else { + Some(parsed_ast) + } + } None => match query_router.parse(&message) { Ok(ast) => Some(ast), Err(error) => { - warn!( - "Query parsing error: {} (client: {})", - error, client_identifier - ); + if error != Error::UnsupportedStatement { + warn!( + "Query parsing error: {} (client: {})", + error, client_identifier + ); + } None } }, diff --git a/src/query_router.rs b/src/query_router.rs index 8b451dd3..343c5c93 100644 --- a/src/query_router.rs +++ b/src/query_router.rs @@ -338,6 +338,8 @@ impl QueryRouter { Some((command, value)) } + const UNSUPPORTED_STATEMENTS_FOR_PARSING: [&'static str; 3] = ["COPY", "TRUNCATE", "VACUUM"]; + pub fn parse(&self, message: &BytesMut) -> Result, Error> { let mut message_cursor = Cursor::new(message); @@ -380,8 +382,22 @@ impl QueryRouter { match Parser::parse_sql(&PostgreSqlDialect {}, &query) { Ok(ast) => Ok(ast), Err(err) => { - debug!("{}: {}", err, query); - Err(Error::QueryRouterParserError(err.to_string())) + let qry_upper = query.to_ascii_uppercase(); + + // Check for unsupported statements to avoid producing a warning. + // Note 1: this is not a complete list of unsupported statements. + // Note 2: we do not check for unsupported statements before going through the + // parser, as the plugin system might be able to handle them, once sqlparser + // is able to correctly parse these (rather valid) queries. + if Self::UNSUPPORTED_STATEMENTS_FOR_PARSING + .iter() + .any(|s| qry_upper.starts_with(s)) + { + Err(Error::UnsupportedStatement) + } else { + debug!("{}: {}", err, query); + Err(Error::QueryRouterParserError(err.to_string())) + } } } } @@ -1864,6 +1880,7 @@ mod test { let query = simple_query("SELECT * FROM pg_database"); let ast = qr.parse(&query).unwrap(); + println!(">>> {:?}", ast); let res = qr.execute_plugins(&ast).await; From a442b31604f7095e43739bfdd8ff7e795781e00c Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 15 Oct 2023 21:27:12 -0700 Subject: [PATCH 3/6] Speed up running tests during development, by adding a `no-clean` argument to avoid doing a clean build. --- tests/docker/run.sh | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/docker/run.sh b/tests/docker/run.sh index ae30c978..5f437a89 100644 --- a/tests/docker/run.sh +++ b/tests/docker/run.sh @@ -1,6 +1,17 @@ #!/bin/bash -rm -rf /app/target/ || true +set -e + +CLEAN_BUILD=true + +if [ $1 = "no-clean" ]; then + echo "INFO: clean build is NOT going to be performed." + CLEAN_BUILD=false +fi + +if $CLEAN_BUILD ; then + rm -rf /app/target/ || true +fi rm /app/*.profraw || true rm /app/pgcat.profdata || true rm -rf /app/cov || true @@ -12,7 +23,9 @@ export RUSTFLAGS="-Zprofile -Ccodegen-units=1 -Copt-level=0 -Clink-dead-code -Co export RUSTDOCFLAGS="-Cpanic=abort" cd /app/ -cargo clean +if $CLEAN_BUILD ; then + cargo clean +fi cargo build cargo test --tests From 9fba9a096c6aafa5c99b6dee5d89935008c8617c Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 15 Oct 2023 21:28:33 -0700 Subject: [PATCH 4/6] `END` is not supported by `sqlparser`. --- tests/pgbench/simple.sql | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/pgbench/simple.sql b/tests/pgbench/simple.sql index ad5e6139..8a01cb4a 100644 --- a/tests/pgbench/simple.sql +++ b/tests/pgbench/simple.sql @@ -26,7 +26,7 @@ UPDATE pgbench_branches SET bbalance = bbalance + :delta WHERE bid = :bid; INSERT INTO pgbench_history (tid, bid, aid, delta, mtime) VALUES (:tid, :bid, :aid, :delta, CURRENT_TIMESTAMP); -END; +COMMIT; SET SHARDING KEY TO :aid; From eaf1ea9102cfcc71ca075c8b69ab0cf4c6dc291e Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Sun, 15 Oct 2023 21:33:26 -0700 Subject: [PATCH 5/6] Nitz. --- src/query_router.rs | 1 - 1 file changed, 1 deletion(-) diff --git a/src/query_router.rs b/src/query_router.rs index 343c5c93..25031f1d 100644 --- a/src/query_router.rs +++ b/src/query_router.rs @@ -1880,7 +1880,6 @@ mod test { let query = simple_query("SELECT * FROM pg_database"); let ast = qr.parse(&query).unwrap(); - println!(">>> {:?}", ast); let res = qr.execute_plugins(&ast).await; From 2af5169c9ea5a021c54e75f8ffe594c72711574c Mon Sep 17 00:00:00 2001 From: Mohammad Dashti Date: Mon, 16 Oct 2023 06:34:19 +0000 Subject: [PATCH 6/6] Added `SET` to unsupported queries + Fixed an issue with running tests on non-clean builds. --- src/query_router.rs | 3 ++- tests/docker/run.sh | 1 + 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/src/query_router.rs b/src/query_router.rs index 25031f1d..4b6873dc 100644 --- a/src/query_router.rs +++ b/src/query_router.rs @@ -338,7 +338,8 @@ impl QueryRouter { Some((command, value)) } - const UNSUPPORTED_STATEMENTS_FOR_PARSING: [&'static str; 3] = ["COPY", "TRUNCATE", "VACUUM"]; + const UNSUPPORTED_STATEMENTS_FOR_PARSING: [&'static str; 4] = + ["COPY", "SET", "TRUNCATE", "VACUUM"]; pub fn parse(&self, message: &BytesMut) -> Result, Error> { let mut message_cursor = Cursor::new(message); diff --git a/tests/docker/run.sh b/tests/docker/run.sh index 5f437a89..bfffc570 100644 --- a/tests/docker/run.sh +++ b/tests/docker/run.sh @@ -7,6 +7,7 @@ CLEAN_BUILD=true if [ $1 = "no-clean" ]; then echo "INFO: clean build is NOT going to be performed." CLEAN_BUILD=false + find /app/target/debug/deps -name *.gcda -exec rm {} \; fi if $CLEAN_BUILD ; then