Fix: Include catalog name in schema creation (#1249)

* include catalog schema creation

* update tests

* style

Author: eakmanrq

sqlmesh/core/engine_adapter/base.py

@@ -457,13 +457,19 @@ def create_view(
         )
 
     def create_schema(
-        self, schema_name: str, ignore_if_exists: bool = True, warn_on_error: bool = True
+        self,
+        schema_name: str,
+        catalog_name: t.Optional[str] = None,
+        ignore_if_exists: bool = True,
+        warn_on_error: bool = True,
     ) -> None:
         """Create a schema from a name or qualified table name."""
+        if catalog_name:
+            schema_name = f"{catalog_name}.{schema_name}"
         try:
             self.execute(
                 exp.Create(
-                    this=exp.to_identifier(schema_name.split(".")[0]),
+                    this=exp.to_table(schema_name),
                     kind="SCHEMA",
                     exists=ignore_if_exists,
                 )

sqlmesh/core/engine_adapter/bigquery.py

@@ -103,14 +103,21 @@ def _is_session_active(self) -> bool:
         return self._session_id is not None
 
     def create_schema(
-        self, schema_name: str, ignore_if_exists: bool = True, warn_on_error: bool = True
+        self,
+        schema_name: str,
+        catalog_name: t.Optional[str] = None,
+        ignore_if_exists: bool = True,
+        warn_on_error: bool = True,
     ) -> None:
         """Create a schema from a name or qualified table name."""
         from google.api_core.exceptions import Conflict
 
         try:
             super().create_schema(
-                schema_name, ignore_if_exists=ignore_if_exists, warn_on_error=False
+                schema_name,
+                catalog_name=catalog_name,
+                ignore_if_exists=ignore_if_exists,
+                warn_on_error=False,
             )
         except Exception as e:
             is_already_exists_error = isinstance(e, Conflict) and "Already Exists:" in str(e)

sqlmesh/core/snapshot/evaluator.py

@@ -700,7 +700,7 @@ def promote(
     ) -> None:
         schema = view_name.schema_for_environment(environment_naming_info=environment_naming_info)
         if schema is not None:
-            self.adapter.create_schema(schema)
+            self.adapter.create_schema(schema, catalog_name=view_name.catalog)
 
         target_name = view_name.for_environment(environment_naming_info)
         logger.info("Updating view '%s' to point at table '%s'", target_name, table_name)
@@ -736,7 +736,8 @@ def create(
         name: str,
         **render_kwargs: t.Any,
     ) -> None:
-        self.adapter.create_schema(exp.to_table(name).db)
+        table = exp.to_table(name)
+        self.adapter.create_schema(table.db, catalog_name=table.catalog)
 
         logger.info("Creating table '%s'", name)
         if model.annotated:
@@ -913,7 +914,8 @@ def create(
         name: str,
         **render_kwargs: t.Any,
     ) -> None:
-        self.adapter.create_schema(exp.to_table(name).db)
+        table = exp.to_table(name)
+        self.adapter.create_schema(table.db, catalog_name=table.catalog)
 
         logger.info("Creating view '%s'", name)
         self.adapter.create_view(

sqlmesh/core/test/definition.py

@@ -65,17 +65,18 @@ def __init__(
 
     def setUp(self) -> None:
         """Load all input tables"""
-        for table, rows in self.body.get("inputs", {}).items():
+        for table_name, rows in self.body.get("inputs", {}).items():
             df = pd.DataFrame.from_records(rows)  # noqa
             columns_to_types: dict[str, exp.DataType] = {}
 
             for i, v in rows[0].items():
                 # convert ruamel into python
                 v = v.real if hasattr(v, "real") else v
                 columns_to_types[i] = parse_one(type(v).__name__, into=exp.DataType)
-
-            self.engine_adapter.create_schema(table)
-            self.engine_adapter.create_view(_test_fixture_name(table), df, columns_to_types)
+            table = exp.to_table(table_name)
+            if table.db:
+                self.engine_adapter.create_schema(table.db, catalog_name=table.catalog)
+            self.engine_adapter.create_view(_test_fixture_name(table_name), df, columns_to_types)
 
     def tearDown(self) -> None:
         """Drop all input tables"""

tests/core/engine_adapter/test_base.py

@@ -46,11 +46,13 @@ def test_create_schema(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
     adapter.create_schema("test_schema")
     adapter.create_schema("test_schema", ignore_if_exists=False)
+    adapter.create_schema("test_schema", catalog_name="test_catalog")
 
     adapter.cursor.execute.assert_has_calls(
         [
             call('CREATE SCHEMA IF NOT EXISTS "test_schema"'),
             call('CREATE SCHEMA "test_schema"'),
+            call('CREATE SCHEMA IF NOT EXISTS "test_catalog"."test_schema"'),
         ]
     )
 

tests/core/engine_adapter/test_duckdb.py

@@ -21,7 +21,6 @@ def test_create_view(adapter: EngineAdapter, duck_conn):
 
 
 def test_create_schema(adapter: EngineAdapter, duck_conn):
-    adapter.create_schema("test_schema")
     adapter.create_schema("test_schema")
     assert duck_conn.execute(
         "SELECT 1 FROM information_schema.schemata WHERE schema_name = 'test_schema'"

tests/core/test_snapshot_evaluator.py

@@ -137,7 +137,7 @@ def x(evaluator, y=None) -> None:
 
     adapter_mock.create_schema.assert_has_calls(
         [
-            call("sqlmesh__test_schema"),
+            call("sqlmesh__test_schema", catalog_name=""),
         ]
     )
 
@@ -189,7 +189,7 @@ def test_promote(mocker: MockerFixture, adapter_mock, make_snapshot):
 
     evaluator.promote([snapshot], EnvironmentNamingInfo(name="test_env"))
 
-    adapter_mock.create_schema.assert_called_once_with("test_schema__test_env")
+    adapter_mock.create_schema.assert_called_once_with("test_schema__test_env", catalog_name=None)
     adapter_mock.create_view.assert_called_once_with(
         "test_schema__test_env.test_model",
         parse_one(
@@ -359,7 +359,7 @@ def test_promote_model_info(mocker: MockerFixture):
         EnvironmentNamingInfo(name="test_env"),
     )
 
-    adapter_mock.create_schema.assert_called_once_with("test_schema__test_env")
+    adapter_mock.create_schema.assert_called_once_with("test_schema__test_env", catalog_name=None)
     adapter_mock.create_view.assert_called_once_with(
         "test_schema__test_env.test_model",
         parse_one("SELECT * FROM physical_schema.test_schema__test_model__1"),