aherman
/
proxysql
mirror of https://github.com/sysown/proxysql
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Improved TAP test

Browse Source
pull/5133/head
Rahim Kanji 7 months ago
parent 9e9fdc23c2
commit 7a30447eaa
1 changed files with 169 additions and 157 deletions
Show Stats Download Patch File Download Diff File

326
test/tap/tests/pgsql-extended_query_protocol_query_rules_test-t.cpp
Unescape View File

@ -55,42 +55,44 @@ PGConnPtr createNewConnection(ConnType conn_type, const std::string& options = "
}
// Helper to clear rules
void clear_rules() {
PGConnPtr admin = createNewConnection(ADMIN);
PQclear(PQexec(admin.get(), "DELETE FROM pgsql_query_rules"));
PQclear(PQexec(admin.get(), "LOAD PGSQL QUERY RULES TO RUNTIME"));
void clear_rules(PGconn* admin) {
PGresult* res = PQexec(admin, "DELETE FROM pgsql_query_rules");
PQclear(res);
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure deletion is processed
res = PQexec(admin, "LOAD PGSQL QUERY RULES TO RUNTIME");
PQclear(res);
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure deletion is processed
}
// Helper to insert rule and load
void insert_rule(const std::string& sql) {
PGConnPtr admin = createNewConnection(ADMIN);
PGresult* res = PQexec(admin.get(), sql.c_str());
void insert_rule(PGconn* admin, const std::string& sql) {
PGresult* res = PQexec(admin, sql.c_str());
if (PQresultStatus(res) != PGRES_COMMAND_OK) {
fprintf(stderr, "Insert rule failed: %s\n", PQerrorMessage(admin.get()));
fprintf(stderr, "Insert rule failed: %s\n", PQerrorMessage(admin));
}
PQclear(res);
res = PQexec(admin.get(), "LOAD PGSQL QUERY RULES TO RUNTIME");
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure deletion is processed
res = PQexec(admin, "LOAD PGSQL QUERY RULES TO RUNTIME");
PQclear(res);
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure deletion is processed
}
void delete_rule(int rule_id) {
PGConnPtr admin = createNewConnection(ADMIN);
void delete_rule(PGconn* admin, int rule_id) {
std::string query = "DELETE FROM pgsql_query_rules WHERE rule_id=" + std::to_string(rule_id);
PGresult* res = PQexec(admin.get(), query.c_str());
PGresult* res = PQexec(admin, query.c_str());
if (PQresultStatus(res) != PGRES_COMMAND_OK) {
fprintf(stderr, "Delete rule failed: %s\n", PQerrorMessage(admin.get()));
fprintf(stderr, "Delete rule failed: %s\n", PQerrorMessage(admin));
}
PQclear(res);
res = PQexec(admin.get(), "LOAD PGSQL QUERY RULES TO RUNTIME");
res = PQexec(admin, "LOAD PGSQL QUERY RULES TO RUNTIME");
PQclear(res);
}
// Helper to get hits for rule_id
int get_hits(int rule_id) {
std::this_thread::sleep_for(std::chrono::milliseconds(3000)); // Ensure stats are updated
PGConnPtr admin = createNewConnection(ADMIN);
int get_hits(PGconn* admin_conn, int rule_id) {
std::this_thread::sleep_for(std::chrono::milliseconds(2000)); // Ensure hits are updated
std::string query = "SELECT hits FROM stats_pgsql_query_rules WHERE rule_id=" + std::to_string(rule_id);
PGresult* res = PQexec(admin.get(), query.c_str());
PGresult* res = PQexec(admin_conn, query.c_str());
if (PQresultStatus(res) != PGRES_TUPLES_OK || PQntuples(res) == 0) {
PQclear(res);
return -1; // Error
@ -101,13 +103,13 @@ int get_hits(int rule_id) {
}
// Function to perform the full flow and check hits after each step
void test_rule_flow(int rule_id, const std::string& test_name, const std::string& query, const std::string& expected_exec_result, bool expect_prepare_fail = false, bool expect_exec_fail = false, int expected_hits_after_prepare = 1, int expected_hits_after_desc = 1, int expected_hits_after_exec = 1, int expected_hits_after_close = 1) {
void test_rule_flow(PGconn* admin, int rule_id, const std::string& test_name, const std::string& query, const std::string& expected_exec_result, bool expect_prepare_fail = false, bool expect_exec_fail = false, int expected_hits_after_prepare = 1, int expected_hits_after_desc = 1, int expected_hits_after_exec = 1, int expected_hits_after_close = 1) {
PGConnPtr conn = createNewConnection(BACKEND);
PGresult* res = PQprepare(conn.get(), "stmt", query.c_str(), 0, nullptr);
if (expect_prepare_fail) {
ok(PQresultStatus(res) != PGRES_COMMAND_OK, (test_name + ": Prepare failed as expected").c_str());
PQclear(res);
int hits = get_hits(rule_id);
int hits = get_hits(admin, rule_id);
ok(hits == expected_hits_after_prepare, (test_name + ": Hits after failed prepare").c_str());
// Skip further
@ -115,14 +117,13 @@ void test_rule_flow(int rule_id, const std::string& test_name, const std::string
} else {
ok(PQresultStatus(res) == PGRES_COMMAND_OK, (test_name + ": Prepare succeeded: %s").c_str(), PQerrorMessage(conn.get()));
PQclear(res);
int hits_after_prepare = get_hits(rule_id);
int hits_after_prepare = get_hits(admin, rule_id);
ok(hits_after_prepare == expected_hits_after_prepare, (test_name + ": '%d/%d' Hits after prepare").c_str(), hits_after_prepare, expected_hits_after_prepare);
// Describe
res = PQdescribePrepared(conn.get(), "stmt");
ok(PQresultStatus(res) == PGRES_COMMAND_OK, (test_name + ": Describe succeeded: %s").c_str(), PQerrorMessage(conn.get()));
PQclear(res);
int hits_after_desc = get_hits(rule_id);
int hits_after_desc = get_hits(admin, rule_id);
ok(hits_after_desc == expected_hits_after_desc, (test_name + ": '%d/%d' Hits after describe").c_str(), hits_after_desc, expected_hits_after_desc);
// Execute
@ -135,14 +136,14 @@ void test_rule_flow(int rule_id, const std::string& test_name, const std::string
(PQgetisnull(res, 0, 0) ? "(null)" : PQgetvalue(res, 0, 0)), expected_exec_result.c_str());
}
PQclear(res);
int hits_after_exec = get_hits(rule_id);
int hits_after_exec = get_hits(admin, rule_id);
ok(hits_after_exec == expected_hits_after_exec, (test_name + ": '%d/%d' Hits after execute").c_str(), hits_after_exec, expected_hits_after_exec);
// Close
res = PQexec(conn.get(), "DEALLOCATE stmt");
ok(PQresultStatus(res) == PGRES_COMMAND_OK, (test_name + ": Close succeeded: %s").c_str(), PQerrorMessage(conn.get()));
PQclear(res);
int hits_after_close = get_hits(rule_id);
int hits_after_close = get_hits(admin, rule_id);
ok(hits_after_close == expected_hits_after_close, (test_name + ": '%d/%d' Hits after close").c_str(), hits_after_close, expected_hits_after_close);
}
}
@ -212,13 +213,7 @@ void consume_results(PGconn* conn) {
}
}
void test_query_processor() {
PGConnPtr admin_conn = createNewConnection(ADMIN);
if (!admin_conn || PQstatus(admin_conn.get()) != CONNECTION_OK) {
BAIL_OUT("Error: failed to connect to the database in file %s, line %d", __FILE__, __LINE__);
return;
}
void test_query_processor(PGconn* admin_conn) {
PGConnPtr backend_conn = createNewConnection(ConnType::BACKEND, "", false);
if (!backend_conn || PQstatus(backend_conn.get()) != CONNECTION_OK) {
@ -232,145 +227,148 @@ void test_query_processor() {
try {
// Test 1: Parse and Sync
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (1,1,'^SELECT 1$',1)");
int initial = get_hits(1);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (1,1,'^SELECT 1$',1)");
int initial = get_hits(admin_conn, 1);
PQsendPrepare(backend_conn.get(), "", "SELECT 1", 0, NULL);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
int after = get_hits(1);
int after = get_hits(admin_conn, 1);
ok((after - initial == 1), "Parse and Sync applies rule once (hits: %d/1)", after - initial);
// Test 2: BIND EXECUTE SYNC (after Parse)
clear_rules();
clear_rules(admin_conn);
PQsendPrepare(backend_conn.get(), "stmt2", "SELECT 1", 0, NULL);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (2,1,'^SELECT 1$',1);");
initial = get_hits(2);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (2,1,'^SELECT 1$',1);");
initial = get_hits(admin_conn, 2);
PQsendQueryPrepared(backend_conn.get(), "stmt2", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after = get_hits(2);
ok(after - initial == 1, "BIND EXECUTE SYNC applies rule once (hits: %d/1)", after - initial);
after = get_hits(admin_conn, 2);
ok(after - initial == 2, "BIND EXECUTE SYNC applies rule twice (hits: %d/2)", after - initial);
// Test 3: BIND DESCRIBE EXECUTE SYNC (after Prepare)
clear_rules();
// Test 3: DESCRIBE BIND DESCRIBE EXECUTE SYNC (after Prepare)
clear_rules(admin_conn);
PQsendPrepare(backend_conn.get(), "stmt3", "SELECT 1", 0, NULL);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (3,1,'^SELECT 1$',1);");
initial = get_hits(3);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (3,1,'^SELECT 1$',1);");
initial = get_hits(admin_conn, 3);
PQsendDescribePrepared(backend_conn.get(), "stmt3");
PQsendQueryPrepared(backend_conn.get(), "stmt3", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after = get_hits(3);
ok(after - initial == 1, "BIND DESCRIBE EXECUTE SYNC applies rule once (hits: %d/1)", after - initial);
after = get_hits(admin_conn, 3);
ok(after - initial == 2, "DESCRIBE BIND DESCRIBE EXECUTE SYNC applies rule twice (hits: %d/2)", after - initial);
// Test 4: PREPARE BIND EXECUTE PREPARE BIND EXECUTE SYNC with same query
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (4,1,'^SELECT 1$',1);");
initial = get_hits(4);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (4,1,'^SELECT 1$',1);");
initial = get_hits(admin_conn, 4);
PQsendPrepare(backend_conn.get(), "stmt4a", "SELECT 1", 0, NULL);
PQsendQueryPrepared(backend_conn.get(), "stmt4a", 0, NULL, NULL, NULL, 0);
PQsendPrepare(backend_conn.get(), "stmt4b", "SELECT 1", 0, NULL);
PQsendQueryPrepared(backend_conn.get(), "stmt4b", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after = get_hits(4);
ok(after - initial == 1, "Multiple PREPARE BIND EXECUTE SYNC with same query applies rule once (hits: %d/1)", after - initial);
after = get_hits(admin_conn, 4);
ok(after - initial == 4, "Multiple PREPARE BIND EXECUTE SYNC with same query (hits: %d/4)", after - initial);
// Test 5: PREPARE BIND EXECUTE PREPARE BIND EXECUTE SYNC with different queries to confirm same HG
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (5,1,'^SELECT 1$',0,1);");
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (6,1,'^SELECT 2$',200,1);");
int initial1 = get_hits(5);
int initial2 = get_hits(6);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (5,1,'^SELECT 1$',0,1);");
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (6,1,'^SELECT 2$',200,1);");
int initial1 = get_hits(admin_conn, 5);
int initial2 = get_hits(admin_conn, 6);
std::string query1 = "SELECT 1";
std::string query2 = "SELECT 2";
PQsendPrepare(backend_conn.get(), "stmt5a", query1.c_str(), 0, NULL);
PQsendDescribePrepared(backend_conn.get(), "stmt5a");
PQsendQueryPrepared(backend_conn.get(), "stmt5a", 0, NULL, NULL, NULL, 0);
PQsendPrepare(backend_conn.get(), "stmt5b", query2.c_str(), 0, NULL);
PQsendDescribePrepared(backend_conn.get(), "stmt5b");
PQsendQueryPrepared(backend_conn.get(), "stmt5b", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
int after1 = get_hits(5);
int after2 = get_hits(6);
bool hits_once = (after1 - initial1 == 1) && (after2 - initial2 == 0);
ok(hits_once, "Multiple message, but only first packet should run query processor. First message (%d/1) Second Message (%d/0)",
int after1 = get_hits(admin_conn, 5);
int after2 = get_hits(admin_conn, 6);
bool hits_once = (after1 - initial1 == 3) && (after2 - initial2 == 3);
ok(hits_once, "Multiple message. First message (%d/3) Second Message (%d/3)",
(after1 - initial1), (after2 - initial2));
// Test 6: Prepare Sync, then BIND EXECUTE SYNC to confirm separate pipelines apply rules separately
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (7,1,'^SELECT 1$',1);");
initial = get_hits(7);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (7,1,'^SELECT 1$',1);");
initial = get_hits(admin_conn, 7);
PQsendPrepare(backend_conn.get(), "stmt6", "SELECT 1", 0, NULL);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after = get_hits(7);
after = get_hits(admin_conn, 7);
int initial_for_bind = after;
PQsendQueryPrepared(backend_conn.get(), "stmt6", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
int after_bind = get_hits(7);
ok((after - initial == 1) && (after_bind - initial_for_bind == 1), "Separate pipelines apply rules separately");
int after_bind = get_hits(admin_conn, 7);
ok((after - initial == 1) && (after_bind - initial_for_bind == 2),
"Prepare Sync then BIND EXECUTE SYNC applies rule (hits: %d/1 and %d/2)", after - initial, after_bind - initial_for_bind);
// Test 7: BIND DESCRIBE EXECUTE SYNC with different query setup
clear_rules();
clear_rules(admin_conn);
PQsendPrepare(backend_conn.get(), "stmt7", "SELECT 1", 0, NULL);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (8,1,'^SELECT 1$',0,1);");
initial = get_hits(8);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (8,1,'^SELECT 1$',0,1);");
initial = get_hits(admin_conn, 8);
PQsendQueryPrepared(backend_conn.get(), "stmt7", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after = get_hits(8);
ok(after - initial == 1, "BIND DESCRIBE EXECUTE SYNC applies rule once (hits: %d/1)", after - initial);
after = get_hits(admin_conn, 8);
ok(after - initial == 2, "BIND DESCRIBE EXECUTE SYNC applies rule twice (hits: %d/2)", after - initial);
// Test 8: Complex mix: PREPARE BIND EXECUTE SYNC then BIND EXECUTE PREPARE BIND EXECUTE SYNC
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (9,1,'^SELECT 1$',0,1);");
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (10,1,'^SELECT 2$',200,1);");
initial1 = get_hits(9);
initial2 = get_hits(10);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (9,1,'^SELECT 1$',0,1);");
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (10,1,'^SELECT 2$',200,1);");
initial1 = get_hits(admin_conn, 9);
initial2 = get_hits(admin_conn, 10);
PQsendPrepare(backend_conn.get(), "stmt8a", "SELECT 1", 0, NULL);
PQsendQueryPrepared(backend_conn.get(), "stmt8a", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
//std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
consume_results(backend_conn.get());
PQsendQueryPrepared(backend_conn.get(), "stmt8a", 0, NULL, NULL, NULL, 0);
PQsendPrepare(backend_conn.get(), "stmt8c", "SELECT 2", 0, NULL);
PQsendQueryPrepared(backend_conn.get(), "stmt8c", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
consume_results(backend_conn.get());
after1 = get_hits(9);
after2 = get_hits(10);
bool hits_correct = (after1 - initial1 == 2) && (after2 - initial2 == 0); // First pipeline hits rule1 on prepare, second pipeline hits rule1 on bind (same query as first)
ok(hits_correct, "Complex mix: hits rule1 twice (%d/2) and rule2 not at all (%d/0)", after1 - initial1, after2 - initial2);
after1 = get_hits(admin_conn, 9);
after2 = get_hits(admin_conn, 10);
bool hits_correct = (after1 - initial1 == 4) && (after2 - initial2 == 2); // First pipeline hits rule1 on prepare, second pipeline hits rule1 on bind (same query as first)
ok(hits_correct, "Complex mix: hits rule1 hits (%d/4) and rule2 hits (%d/2)", after1 - initial1, after2 - initial2);
// Test 9: PARSE EXECUTE SYNC with SET statement
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (11,1,'^SELECT 1$',1);");
initial = get_hits(11);
clear_rules(admin_conn);
insert_rule(admin_conn, "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (11,1,'^SELECT 1$',1);");
initial = get_hits(admin_conn, 11);
PQsendPrepare(backend_conn.get(), "stmt9a", "SELECT 1", 0, NULL);
PQsendPrepare(backend_conn.get(), "stmt9b", "SET client_min_messages = 'error'", 0, NULL);
PQpipelineSync(backend_conn.get());
std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
//std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
consume_results(backend_conn.get());
PQsendQueryPrepared(backend_conn.get(), "stmt9a", 0, NULL, NULL, NULL, 0);
PQsendQueryPrepared(backend_conn.get(), "stmt9b", 0, NULL, NULL, NULL, 0);
PQpipelineSync(backend_conn.get());
std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
//std::this_thread::sleep_for(std::chrono::milliseconds(1000)); // Ensure pipeline is ready
consume_results(backend_conn.get());
consume_results(backend_conn.get());
int exited = PQexitPipelineMode(backend_conn.get());
ok(exited == 1, "Exited pipeline mode successfully %s", PQerrorMessage(backend_conn.get()));
PQsetnonblocking(backend_conn.get(), 0);
after = get_hits(11);
ok(after - initial == 2, "PARSE EXECUTE SYNC with SET statement applies rule once (hits: %d/2)", after - initial);
after = get_hits(admin_conn, 11);
ok(after - initial == 3, "PARSE EXECUTE SYNC with SET statement applies rule once (hits: %d/3)", after - initial);
PGresult* res = PQexec(backend_conn.get(), "SHOW client_min_messages");
const char* encoding = PQgetvalue(res, 0, 0);
ok(PQresultStatus(res) == PGRES_TUPLES_OK && strcmp(encoding, "error") == 0, "client_min_messages is 'error': %s", encoding);
@ -388,48 +386,62 @@ int main(int argc, char** argv) {
plan(101);
PGConnPtr admin_conn = createNewConnection(ADMIN);
if (!admin_conn || PQstatus(admin_conn.get()) != CONNECTION_OK) {
BAIL_OUT("Error: failed to connect to the database in file %s, line %d", __FILE__, __LINE__);
return exit_status();
}
{
PGresult* res = PQexec(admin_conn.get(), "SET pgsql-poll_timeout=1000");
PQclear(res);
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure setting is processed
res = PQexec(admin_conn.get(), "LOAD PGSQL VARIABLES TO RUNTIME");
PQclear(res);
std::this_thread::sleep_for(std::chrono::milliseconds(100)); // Ensure setting is processed
}
// Test 1: Query rewrite
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (1,1,'^SELECT 1$','SELECT 2',1)");
test_rule_flow(1, "Rewrite", "SELECT 1", "2");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (1,1,'^SELECT 1$','SELECT 2',1)");
test_rule_flow(admin_conn.get(), 1, "Rewrite", "SELECT 1", "2");
// Test 2: Query error (block with error_msg)
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, error_msg, apply) VALUES (2,1,'^SELECT 1$','blocked',1)");
test_rule_flow(2, "Error Block", "SELECT 1", "", true, false, 1, 1, 1, 1);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, error_msg, apply) VALUES (2,1,'^SELECT 1$','blocked',1)");
test_rule_flow(admin_conn.get(), 2, "Error Block", "SELECT 1", "", true, false, 1, 1, 1, 1);
// Test 3: Query pass (no change)
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (3,1,'^SELECT 1$',1)");
test_rule_flow(3, "Pass", "SELECT 1", "1", false, false, 1, 2, 3, 3);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (3,1,'^SELECT 1$',1)");
test_rule_flow(admin_conn.get(), 3, "Pass", "SELECT 1", "1", false, false, 1, 2, 4, 4);
// Test 4: Query OK_msg
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, OK_msg, apply) VALUES (4,1,'^SELECT 1$','passed',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, OK_msg, apply) VALUES (4,1,'^SELECT 1$','passed',1)");
PGConnPtr conn = createNewConnection(BACKEND);
PGresult* res = PQprepare(conn.get(), "", "SELECT 1", 0, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
ok(true, "OK_msg: Prepare succeeded");
PQclear(res);
res = PQexecPrepared(conn.get(), "", 0, nullptr, nullptr, nullptr, 0);
//ok(PQresultStatus(res) == PGRES_COMMAND_OK, "OK_msg: Execution passed (Command Completion)");
ok(PQresultStatus(res) == PGRES_TUPLES_OK, "OK_msg: Execution passed");
ok(PQresultStatus(res) == PGRES_COMMAND_OK, "OK_msg: Execution passed (Command Completion)");
//ok(PQresultStatus(res) == PGRES_TUPLES_OK, "OK_msg: Execution passed");
PQclear(res);
int hits = get_hits(4);
ok(hits == 2, "OK_msg: '%d/2' Hits", hits);
int hits = get_hits(admin_conn.get(), 4);
ok(hits == 3, "OK_msg: '%d/3' Hits", hits);
} else {
ok(false, "OK_msg: Prepare failed");
}
// Test 5: Query routing
clear_rules();
clear_rules(admin_conn.get());
// Assume HG 0 default, HG1 exists
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (5,1,'^SELECT 10$',200,1)");
test_rule_flow(5, "Routing", "SELECT 10", "", true,true,1);
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, destination_hostgroup, apply) VALUES (5,1,'^SELECT 10$',200,1)");
test_rule_flow(admin_conn.get(), 5, "Routing", "SELECT 10", "", true,true,1);
// Edge case 6: Unnamed statement rewrite
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (6,1,'^SELECT 1$','SELECT 3',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (6,1,'^SELECT 1$','SELECT 3',1)");
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "", "SELECT 1", 0, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
@ -439,15 +451,15 @@ int main(int argc, char** argv) {
ok(PQresultStatus(res) == PGRES_TUPLES_OK && std::string(PQgetvalue(res, 0, 0)) == "3",
"Unnamed Rewrite: Execute rewritten");
PQclear(res);
int hits = get_hits(6);
int hits = get_hits(admin_conn.get(), 6);
ok(hits == 1, "Unnamed Rewrite: '%d/1' Hits", hits);
} else {
ok(false, "Unnamed Rewrite: Prepare failed");
}
// Edge case 7: Rewrite with parameters
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (7,1,'^SELECT \\$1 AS val$','SELECT 4 AS val',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (7,1,'^SELECT \\$1 AS val$','SELECT 4 AS val',1)");
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "stmt", "SELECT $1 AS val", 1, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
@ -458,7 +470,7 @@ int main(int argc, char** argv) {
ok(PQresultStatus(res) != PGRES_TUPLES_OK,
"Param Rewrite: Execute rewritten, param failed: %s", PQerrorMessage(conn.get()));
PQclear(res);
int hits = get_hits(7);
int hits = get_hits(admin_conn.get(), 7);
ok(hits == 1, "Param Rewrite: '%d/1' Hits", hits);
res = PQexec(conn.get(), "DEALLOCATE stmt");
PQclear(res);
@ -467,41 +479,41 @@ int main(int argc, char** argv) {
}
// Edge case 8: Rewrite to invalid query
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (8,1,'^SELECT 1$','SELECT INVALID',1)");
test_rule_flow(8, "Invalid Rewrite", "SELECT 1", "", true);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (8,1,'^SELECT 1$','SELECT INVALID',1)");
test_rule_flow(admin_conn.get(), 8, "Invalid Rewrite", "SELECT 1", "", true);
// Edge case 9: Case insensitive match
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, re_modifiers, replace_pattern, apply) VALUES (9,1,'^select 1$','CASELESS','SELECT 5',1)");
test_rule_flow(9, "Case Insensitive", "SELECT 1", "5");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, re_modifiers, replace_pattern, apply) VALUES (9,1,'^select 1$','CASELESS','SELECT 5',1)");
test_rule_flow(admin_conn.get(), 9, "Case Insensitive", "SELECT 1", "5");
// Edge case 10: Negate match
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, negate_match_pattern, apply, flagOUT) VALUES (10,1,'^SELECT 1$',1, 0, 200)");
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply, flagIN) VALUES (11,1,'^SELECT 2$','SELECT 600',1, 200)");
test_rule_flow(10, "Negate", "SELECT 2", "600", false, false, 1, 2, 3, 4);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, negate_match_pattern, apply, flagOUT) VALUES (10,1,'^SELECT 1$',1, 0, 200)");
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply, flagIN) VALUES (11,1,'^SELECT 2$','SELECT 600',1, 200)");
test_rule_flow(admin_conn.get(), 10, "Negate", "SELECT 2", "600", false, false, 1, 2, 4, 5);
// Additional test 11: Multi-execute
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (11,1,'^SELECT 1$',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (11,1,'^SELECT 1$',1)");
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "stmt", "SELECT 1", 0, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
ok(true, "Multi-Exec: Prepare succeeded");
PQclear(res);
int hits_after_prepare = get_hits(11);
int hits_after_prepare = get_hits(admin_conn.get(), 11);
ok(hits_after_prepare == 1, "Multi-Exec: '%d/1' Hits after prepare", hits_after_prepare);
res = PQexecPrepared(conn.get(), "stmt", 0, nullptr, nullptr, nullptr, 0);
ok(PQresultStatus(res) == PGRES_TUPLES_OK, "Multi-Exec: First execute succeeded");
PQclear(res);
int hits_after_first = get_hits(11);
ok(hits_after_first == 2, "Multi-Exec: Hits still '%d/2' after first exec", hits_after_first);
int hits_after_first = get_hits(admin_conn.get(), 11);
ok(hits_after_first == 3, "Multi-Exec: Hits '%d/3' after first exec", hits_after_first);
res = PQexecPrepared(conn.get(), "stmt", 0, nullptr, nullptr, nullptr, 0);
ok(PQresultStatus(res) == PGRES_TUPLES_OK, "Multi-Exec: Second execute succeeded");
PQclear(res);
int hits_after_second = get_hits(11);
ok(hits_after_second == 3, "Multi-Exec: Hits still '%d/3' after second exec", hits_after_second);
int hits_after_second = get_hits(admin_conn.get(), 11);
ok(hits_after_second == 5, "Multi-Exec: Hits '%d/5' after second exec", hits_after_second);
res = PQexec(conn.get(), "DEALLOCATE stmt");
PQclear(res);
} else {
@ -509,35 +521,35 @@ int main(int argc, char** argv) {
}
// Additional test 12: Rule with username filter matching
clear_rules();
clear_rules(admin_conn.get());
std::string username = cl.pgsql_username;
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, username, match_pattern, replace_pattern, apply) VALUES (12,1,'" + username + "','^SELECT 1$','SELECT 7',1)");
test_rule_flow(12, "Username Match", "SELECT 1", "7");
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, username, match_pattern, replace_pattern, apply) VALUES (12,1,'" + username + "','^SELECT 1$','SELECT 7',1)");
test_rule_flow(admin_conn.get(), 12, "Username Match", "SELECT 1", "7");
// Additional test 13: Rule with username filter not matching
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, username, match_pattern, replace_pattern, apply) VALUES (13,1,'wronguser','^SELECT 1$','SELECT 8',1)");
test_rule_flow(13, "Username No Match", "SELECT 1", "1" /* original */, false, false, 0 /* no hit */, 0, 0, 0);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, username, match_pattern, replace_pattern, apply) VALUES (13,1,'wronguser','^SELECT 1$','SELECT 8',1)");
test_rule_flow(admin_conn.get(), 13, "Username No Match", "SELECT 1", "1" /* original */, false, false, 0 /* no hit */, 0, 0, 0);
// Additional test 14: Chain of rules - first sets flagOUT, second rewrites
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, flagOUT, apply) VALUES (14,1,'^SELECT 1$',100,0)");
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, flagIN, match_pattern, replace_pattern, apply) VALUES (15,1,100,'^SELECT 1$','SELECT 9',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, flagOUT, apply) VALUES (14,1,'^SELECT 1$',100,0)");
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, flagIN, match_pattern, replace_pattern, apply) VALUES (15,1,100,'^SELECT 1$','SELECT 9',1)");
// Check hits for both
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "stmt", "SELECT 1", 0, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
ok(true, "Chain: Prepare succeeded");
PQclear(res);
int hits14 = get_hits(14);
int hits15 = get_hits(15);
int hits14 = get_hits(admin_conn.get(), 14);
int hits15 = get_hits(admin_conn.get(), 15);
ok(hits14 == 1 && hits15 == 1, "Chain: Both rules hit after prepare 14:'%d/1' 15:'%d/1'", hits14, hits15);
res = PQexecPrepared(conn.get(), "stmt", 0, nullptr, nullptr, nullptr, 0);
ok(PQresultStatus(res) == PGRES_TUPLES_OK && std::string(PQgetvalue(res, 0, 0)) == "9",
"Chain: Execute rewritten by second rule");
PQclear(res);
// Hits should not increase on exec
ok(get_hits(14) == 1 && get_hits(15) == 1, "Chain: Hits unchanged after exec");
ok(get_hits(admin_conn.get(), 14) == 1 && get_hits(admin_conn.get(), 15) == 1, "Chain: Hits unchanged after exec");
res = PQexec(conn.get(), "DEALLOCATE stmt");
PQclear(res);
} else {
@ -545,19 +557,19 @@ int main(int argc, char** argv) {
}
// Additional test 17: Timeout rule on long query
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, timeout, apply) VALUES (17,1,'^SELECT pg_sleep\\(10\\)$',1000,1)"); // 1s timeout for 2s sleep
test_rule_flow(17, "Timeout", "SELECT pg_sleep(10)", "", false, true, 1, 2, 3, 3);
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, timeout, apply) VALUES (17,1,'^SELECT pg_sleep\\(10\\)$',1000,1)"); // 1s timeout for 2s sleep
test_rule_flow(admin_conn.get(), 17, "Timeout", "SELECT pg_sleep(10)", "", false, true, 1, 2, 4, 4);
// Additional test 18: Parameter bind with different values
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (18,1,'^SELECT \\$1 AS val$',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, apply) VALUES (18,1,'^SELECT \\$1 AS val$',1)");
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "stmt", "SELECT $1 AS val", 1, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
ok(true, "Param Bind: Prepare succeeded");
PQclear(res);
int hits_after_prepare = get_hits(18);
int hits_after_prepare = get_hits(admin_conn.get(), 18);
ok(hits_after_prepare == 1, "Param Bind: '%d/1' Hits after prepare", hits_after_prepare);
const char* param1 = "10";
@ -566,8 +578,8 @@ int main(int argc, char** argv) {
"Param Bind: First bind exec succeeded");
PQclear(res);
int hits_after_first = get_hits(18);
ok(hits_after_first == 2, "Param Bind: Hits still '%d/2' after first", hits_after_first);
int hits_after_first = get_hits(admin_conn.get(), 18);
ok(hits_after_first == 3, "Param Bind: Hits still '%d/3' after first", hits_after_first);
const char* param2 = "20";
res = PQexecPrepared(conn.get(), "stmt", 1, &param2, nullptr, nullptr, 0);
@ -575,8 +587,8 @@ int main(int argc, char** argv) {
"Param Bind: Second bind exec succeeded");
PQclear(res);
int hits_after_second = get_hits(18);
ok(hits_after_second == 3, "Param Bind: Hits still '%d/3' after second", hits_after_second);
int hits_after_second = get_hits(admin_conn.get(), 18);
ok(hits_after_second == 5, "Param Bind: Hits still '%d/5' after second", hits_after_second);
res = PQexec(conn.get(), "DEALLOCATE stmt");
PQclear(res);
@ -586,8 +598,8 @@ int main(int argc, char** argv) {
}
// Additional test 20: Error on bind after rewrite to no param
clear_rules();
insert_rule("INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (19,1,'^SELECT \\$1$','SELECT 1',1)");
clear_rules(admin_conn.get());
insert_rule(admin_conn.get(), "INSERT INTO pgsql_query_rules (rule_id, active, match_pattern, replace_pattern, apply) VALUES (19,1,'^SELECT \\$1$','SELECT 1',1)");
conn = createNewConnection(BACKEND);
res = PQprepare(conn.get(), "stmt", "SELECT $1", 1, nullptr);
if (PQresultStatus(res) == PGRES_COMMAND_OK) {
@ -598,13 +610,13 @@ int main(int argc, char** argv) {
ok(PQresultStatus(res) != PGRES_TUPLES_OK, "Rewrite No Param: Exec failed due to param mismatch");
PQclear(res);
int hits = get_hits(19);
int hits = get_hits(admin_conn.get(), 19);
ok(hits == 1, "Rewrite No Param: '%d/1' Hits", hits);
} else {
ok(false, "Rewrite No Param: Prepare failed");
}
test_query_processor();
test_query_processor(admin_conn.get());
return exit_status();
}

Write Preview
Loading…
Cancel
Save