More robust MySQL catalog usage detection (#366)

aws · Jan 14, 2025 · aa04779 · aa04779
1 parent f0dfc1c
commit aa04779
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/mysql/lib/dialect/mysql_database_dialect.ts b/mysql/lib/dialect/mysql_database_dialect.ts
@@ -195,8 +195,9 @@ export class MySQLDatabaseDialect implements DatabaseDialect {
   }
 
   doesStatementSetCatalog(statement: string): string | undefined {
-    if (statement.includes("use")) {
-      return statement.split(" ")[1];
+    const catalogRegexp = /^use\s+(\w+)/i;
+    if (catalogRegexp.test(statement)) {
+      return statement.split(catalogRegexp)[1];
     }
 
     return undefined;

diff --git a/tests/unit/sql_method_utils.test.ts b/tests/unit/sql_method_utils.test.ts
@@ -127,7 +127,9 @@ describe("test sql method utils", () => {
     [[" /* COMMENT */ select /* COMMENT */ 1 "], undefined],
     [[" use dbName "], "dbname"],
     [[" use/* COMMENT use dbName3*/ dbName "], "dbname"],
-    [[" use dbName1 ", " use dbName2 "], "dbname2"]
+    [[" use dbName1 ", " use dbName2 "], "dbname2"],
+    [[" select * from user", " select /* use dbName */ * from user "], undefined],
+    [[" use dbName; select 1"], "dbname"]
   ])("test catalog", (sql: string[], expectedResult: string | undefined) => {
     expect(SqlMethodUtils.doesSetCatalog(sql, new MySQLDatabaseDialect())).toBe(expectedResult);
   });