Skip to content

Commit

Permalink
Merge pull request #4605 from sysown/v2.x_4598
Browse files Browse the repository at this point in the history
Better parsing of USE queries sent with COM_QUERY #4598
  • Loading branch information
renecannao authored Aug 13, 2024
2 parents 01d1fc9 + c049c92 commit 050907e
Show file tree
Hide file tree
Showing 4 changed files with 181 additions and 24 deletions.
6 changes: 6 additions & 0 deletions include/set_parser.h
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,12 @@ class SetParser {
// First implemenation of the parser for TRANSACTION ISOLATION LEVEL and TRANSACTION READ/WRITE
std::map<std::string, std::vector<std::string>> parse2();
std::string parse_character_set();
std::string parse_USE_query();
std::string remove_comments(const std::string& q);
#ifdef DEBUG
// built-in testing
void test_parse_USE_query();
#endif // DEBUG
~SetParser();
};

Expand Down
41 changes: 19 additions & 22 deletions lib/MySQL_Session.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -6138,30 +6138,27 @@ void MySQL_Session::handler___status_WAITING_CLIENT_DATA___STATE_SLEEP___MYSQL_C
if (session_type == PROXYSQL_SESSION_MYSQL) {
__sync_fetch_and_add(&MyHGM->status.frontend_use_db, 1);
string nq=string((char *)pkt->ptr+sizeof(mysql_hdr)+1,pkt->size-sizeof(mysql_hdr)-1);
RE2::GlobalReplace(&nq,(char *)"(?U)/\\*.*\\*/",(char *)" ");
char *sn_tmp = (char *)nq.c_str();
while (sn_tmp < ( nq.c_str() + nq.length() - 4 ) && *sn_tmp == ' ')
sn_tmp++;
//char *schemaname=strdup(nq.c_str()+4);
char *schemaname=strdup(sn_tmp+3);
char *schemanameptr=trim_spaces_and_quotes_in_place(schemaname);
// handle cases like "USE `schemaname`
if(schemanameptr[0]=='`' && schemanameptr[strlen(schemanameptr)-1]=='`') {
schemanameptr[strlen(schemanameptr)-1]='\0';
schemanameptr++;
}
client_myds->myconn->userinfo->set_schemaname(schemanameptr,strlen(schemanameptr));
free(schemaname);
if (mirror==false) {
SetParser parser(nq);
string schemaname = parser.parse_USE_query();
if (schemaname != "") {
client_myds->myconn->userinfo->set_schemaname((char *)schemaname.c_str(),schemaname.length());
if (mirror==false) {
RequestEnd(NULL);
}
l_free(pkt->size,pkt->ptr);
client_myds->setDSS_STATE_QUERY_SENT_NET();
unsigned int nTrx=NumActiveTransactions();
uint16_t setStatus = (nTrx ? SERVER_STATUS_IN_TRANS : 0 );
if (autocommit) setStatus |= SERVER_STATUS_AUTOCOMMIT;
client_myds->myprot.generate_pkt_OK(true,NULL,NULL,1,0,0,setStatus,0,NULL);
GloMyLogger->log_audit_entry(PROXYSQL_MYSQL_INITDB, this, NULL);
} else {
l_free(pkt->size,pkt->ptr);
client_myds->setDSS_STATE_QUERY_SENT_NET();
std::string msg = "Unable to parse: " + nq;
client_myds->myprot.generate_pkt_ERR(true,NULL,NULL,client_myds->pkt_sid+1,1148,(char *)"42000", msg.c_str());
RequestEnd(NULL);
}
l_free(pkt->size,pkt->ptr);
client_myds->setDSS_STATE_QUERY_SENT_NET();
unsigned int nTrx=NumActiveTransactions();
uint16_t setStatus = (nTrx ? SERVER_STATUS_IN_TRANS : 0 );
if (autocommit) setStatus |= SERVER_STATUS_AUTOCOMMIT;
client_myds->myprot.generate_pkt_OK(true,NULL,NULL,1,0,0,setStatus,0,NULL);
GloMyLogger->log_audit_entry(PROXYSQL_MYSQL_INITDB, this, NULL);
client_myds->DSS=STATE_SLEEP;
} else {
l_free(pkt->size,pkt->ptr);
Expand Down
151 changes: 149 additions & 2 deletions lib/set_parser.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -3,9 +3,11 @@
#include <string>
#include <vector>
#include <map>
#ifdef PARSERDEBUG
#include <cassert>
#include <utility> // for std::pair
//#ifdef PARSERDEBUG
#include <iostream>
#endif
//#endif

using namespace std;

Expand Down Expand Up @@ -506,3 +508,148 @@ std::string SetParser::parse_character_set() {
return value4;
}

std::string SetParser::parse_USE_query() {
#ifdef DEBUG
proxy_debug(PROXY_DEBUG_MYSQL_QUERY_PROCESSOR, 4, "Parsing query %s\n", query.c_str());
#endif // DEBUG

re2::RE2::Options opt2(RE2::Quiet);
opt2.set_case_sensitive(false);
opt2.set_longest_match(false);

std::string dbname = remove_comments(query);
re2::RE2 re0("^\\s*", opt2);
re2::RE2::Replace(&dbname, re0, "");
if (dbname.size() >= 4) {
if (
strncasecmp(dbname.c_str(), "USE ",4) == 0
||
strncasecmp(dbname.c_str(), "USE`",4) == 0
) {
re2::RE2 re1("^USE\\s*", opt2);
re2::RE2::Replace(&dbname, re1, "");
re2::RE2 re2("\\s*$", opt2);
re2::RE2::Replace(&dbname, re2, "");
if (dbname[0] == '`') {
if (dbname.length() > 2) {
if (dbname[dbname.length()-1] == '`') {
// Remove the first character
dbname.erase(0, 1);
// Remove the last character
dbname.erase(dbname.length() - 1);
return dbname;
}
}
} else {
return dbname;
}
}
} else {
return "";
}

return "";
}


std::string SetParser::remove_comments(const std::string& q) {
std::string result = "";
bool in_multiline_comment = false;

for (size_t i = 0; i < query.size(); ++i) {
char current_char = query[i];

// Check for multiline comment start
if (current_char == '/' && i + 1 < query.size() && query[i + 1] == '*') {
in_multiline_comment = true;
i++; // Skip the '*'
continue;
}

// Check for multiline comment end
if (in_multiline_comment && current_char == '*' && i + 1 < query.size() && query[i + 1] == '/') {
in_multiline_comment = false;
i++; // Skip the '/'
continue;
}

// Skip characters inside multiline comment
if (in_multiline_comment) {
continue;
}

// Check for single-line comments
if (current_char == '#' || (current_char == '-' && i + 1 < query.size() && query[i + 1] == '-')) {
// Skip until the end of the line
while (i < query.size() && query[i] != '\n') {
i++;
}
continue;
}

// Append the character to the result if it's not a comment
result += current_char;
}

return result;
}


#ifdef DEBUG
void SetParser::test_parse_USE_query() {

// Define vector of pairs (query, expected dbname)
std::vector<std::pair<std::string, std::string>> testCases = {
{"USE my_database", "my_database"}, // Basic Case
{"USE my_database", "my_database"}, // Basic Case
{"USE my_database ", "my_database"}, // Basic Case
{"/* comment */USE dbname /* comment */", "dbname"}, // With Comments
{"/* comment */ USE dbname", "dbname"}, // With Comments
{"USE dbname /* comment */", "dbname"}, // With Comments
{"/* comment */USE `dbname` /* comment */", "dbname"}, // With backtick
{"/* comment */USE `dbname`/* comment */", "dbname"}, // With backtick
{"/* comment */USE`dbname` /* comment */", "dbname"}, // With backtick
{"/* comment */USE `dbname`/* comment */", "dbname"}, // With backtick
{"/* comment\nmultiline comment */USE dbname /* comment */", "dbname"}, // Multiline Comment
{"/* comment */USE dbname # comment", "dbname"}, // Hash Comment
{"/* comment */USE dbname -- comment", "dbname"}, // Double Dash Comment
{"/* comment */USE dbname # comment", "dbname"}, // Hash Comment
{"/* comment */USE dbname -- comment", "dbname"}, // Double Dash Comment
{"USE dbname # comment", "dbname"}, // Hash Comment
{"USE dbname -- comment", "dbname"}, // Double Dash Comment
{"SELECT * FROM my_table", ""}, // No match
{"/*+ placeholder_comment */ USE test_use_comment", "test_use_comment"},

{"USE /*+ placeholder_comment */ `test_use_comment-a1`", "test_use_comment-a1"},
{"USE /*+ placeholder_comment */ `test_use_comment_1`", "test_use_comment_1"},
{"USE/*+ placeholder_comment */ `test_use_comment_2`", "test_use_comment_2"},
{"USE /*+ placeholder_comment */`test_use_comment_3`", "test_use_comment_3"},
{"USE /*+ placeholder_comment */ test_use_comment_4", "test_use_comment_4"},
{"USE/*+ placeholder_comment */ test_use_comment_5", "test_use_comment_5"},
{"USE /*+ placeholder_comment */test_use_comment_6", "test_use_comment_6"},
{"USE /*+ placeholder_comment */ `test_use_comment-1`", "test_use_comment-1"},
{"use my_database", "my_database"},
{"/* comment */ use dbname -- comment", "dbname"},
{"/* comment\nmultiline comment */USE dbname /* comment\nmultiline comment */", "dbname"}, // Multiline Comment

{"USE/*+ placeholder_comment */ `test_use_comment-2`", "test_use_comment-2"},
{"USE /*+ placeholder_comment */`test_use_comment-3`", "test_use_comment-3"},
{"/*+ placeholder_comment */USE `test_use_comment-4`", "test_use_comment-4"},
{"USE/*+ placeholder_comment */`test_use_comment-5`", "test_use_comment-5"},
{"/* comment */USE`test_use_comment-6`", "test_use_comment-6"},
{"USE`test_use_comment-7`", "test_use_comment-7"},
};

// Run tests for each pair
for (const auto& p : testCases) {
set_query(p.first);
std::string dbname = parse_USE_query();
if (dbname != p.second) {
// we call parse_USE_query() again just to make it easier to create a breakpoint
std::string s = parse_USE_query();
assert(s == p.second);
}
}

}
#endif // DEBUG
7 changes: 7 additions & 0 deletions src/main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1976,6 +1976,13 @@ int main(int argc, const char * argv[]) {
// std::cerr << "Main init phase0 completed in ";
#endif
}
#ifdef DEBUG
{
// Automated testing
SetParser parser("");
parser.test_parse_USE_query();
}
#endif // DEBUG
{
cpu_timer t;
ProxySQL_Main_process_global_variables(argc, argv);
Expand Down

0 comments on commit 050907e

Please sign in to comment.