fix: Fix schema wrapped in square brackets

dodopizza · Sep 7, 2022 · 3f71712 · 3f71712
1 parent 0a3d77b
commit 3f71712
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py
@@ -164,7 +164,7 @@ def _convert_schema_in_statement(query: str) -> str:
             - MyTable                     -> MyTable
         """
 
-        pattern = r"^([a-zA-Z0-9]+\b|\"[a-zA-Z0-9 \-_.]+\")?\.?([a-zA-Z0-9]+\b|\"[a-zA-Z0-9 \-_.]+\")"
+        pattern = r"^\[?([a-zA-Z0-9]+\b|\"[a-zA-Z0-9 \-_.]+\")?\]?\.?\[?([a-zA-Z0-9]+\b|\"[a-zA-Z0-9 \-_.]+\")\]?"
         match = re.search(pattern, query)
 
         if not match or not match.group(1):

diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py
@@ -185,6 +185,8 @@ def test_schema_from_metadata(table_name: str, schema_name: str, expected_table_
         ('"schema name"."table name"', 'database("schema name")."table name"'),
         ('"table.name"', '"table.name"'),
         ("MyTable", "MyTable"),
+        ('["schema"].["table"]', 'database("schema")."table"'),
+        ('["table"]', '["table"]'),
     ],
 )
 def test_schema_from_query(query_table_name: str, expected_table_name: str):