Repository: TobikoData/sqlmesh
Pull request #1882 (2023-12-21 22:13:59)

Fix: Properly Quote Set Catalog


Patch:
diff --git a/sqlmesh/core/engine_adapter/duckdb.py b/sqlmesh/core/engine_adapter/duckdb.py
index 9d381bcc..c5b1875a 100644
--- a/sqlmesh/core/engine_adapter/duckdb.py
+++ b/sqlmesh/core/engine_adapter/duckdb.py
@@ -23,7 +23,7 @@ class DuckDBEngineAdapter(LogicalMergeMixin, GetCurrentCatalogFromFunctionMixin)
 
     def set_current_catalog(self, catalog: str) -> None:
         """Sets the catalog name of the current connection."""
-        self.execute(f"USE {catalog}")
+        self.execute(exp.Use(this=exp.to_identifier(catalog)))
 
     def _df_to_source_queries(
         self,
diff --git a/sqlmesh/core/engine_adapter/trino.py b/sqlmesh/core/engine_adapter/trino.py
index cecbec54..68b05404 100644
--- a/sqlmesh/core/engine_adapter/trino.py
+++ b/sqlmesh/core/engine_adapter/trino.py
@@ -6,7 +6,7 @@ import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
 
-from sqlmesh.core.dialect import to_schema
+from sqlmesh.core.dialect import schema_, to_schema
 from sqlmesh.core.engine_adapter.base import (
     CatalogSupport,
     InsertOverwriteStrategy,
@@ -50,7 +50,7 @@ class TrinoEngineAdapter(
 
     def set_current_catalog(self, catalog: str) -> None:
         """Sets the catalog name of the current connection."""
-        self.execute(f"USE {catalog}.information_schema")
+        self.execute(exp.Use(this=schema_(catalog, "information_schema")))
 
     def _insert_overwrite_by_condition(
         self,


Test patch:
diff --git a/tests/core/engine_adapter/test_duckdb.py b/tests/core/engine_adapter/test_duckdb.py
index edad0d8e..6f941549 100644
--- a/tests/core/engine_adapter/test_duckdb.py
+++ b/tests/core/engine_adapter/test_duckdb.py
@@ -1,9 +1,12 @@
+import typing as t
+
 import pandas as pd
 import pytest
 from sqlglot import expressions as exp
 from sqlglot import parse_one
 
 from sqlmesh.core.engine_adapter import DuckDBEngineAdapter, EngineAdapter
+from tests.core.engine_adapter import to_sql_calls
 
 
 @pytest.fixture
@@ -61,3 +64,12 @@ def test_replace_query_pandas(adapter: EngineAdapter, duck_conn):
         "test_table", df, {"a": exp.DataType.build("long"), "b": exp.DataType.build("long")}
     )
     pd.testing.assert_frame_equal(adapter.fetchdf("SELECT * FROM test_table"), df)
+
+
+def test_set_current_catalog(make_mocked_engine_adapter: t.Callable, duck_conn):
+    adapter = make_mocked_engine_adapter(DuckDBEngineAdapter)
+    adapter.set_current_catalog("test_catalog")
+
+    assert to_sql_calls(adapter) == [
+        'USE "test_catalog"',
+    ]
diff --git a/tests/core/engine_adapter/test_trino.py b/tests/core/engine_adapter/test_trino.py
new file mode 100644
index 00000000..76a82169
--- /dev/null
+++ b/tests/core/engine_adapter/test_trino.py
@@ -0,0 +1,13 @@
+import typing as t
+
+from sqlmesh.core.engine_adapter import TrinoEngineAdapter
+from tests.core.engine_adapter import to_sql_calls
+
+
+def test_set_current_catalog(make_mocked_engine_adapter: t.Callable, duck_conn):
+    adapter = make_mocked_engine_adapter(TrinoEngineAdapter)
+    adapter.set_current_catalog("test_catalog")
+
+    assert to_sql_calls(adapter) == [
+        'USE "information_schema"."test_catalog"',
+    ]