Repository: TobikoData/sqlmesh
Pull request #2251 (2024-03-12 04:15:22)

fix: warn instead of error single catalog multiple catalogs
Prior to this PR if we detected a single catalog engine (like Postgres) referred to another catalog then it would error. This is likely excessive since the single catalog would error on it's own anyways and it can potentially raise errors when they are needed. Therefore this PR switches this to log a warning instead. 


Patch:
diff --git a/sqlmesh/core/engine_adapter/shared.py b/sqlmesh/core/engine_adapter/shared.py
index 1c98d40f..4c299e8e 100644
--- a/sqlmesh/core/engine_adapter/shared.py
+++ b/sqlmesh/core/engine_adapter/shared.py
@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import functools
 import inspect
+import logging
 import types
 import typing as t
 from enum import Enum
@@ -18,6 +19,9 @@ if t.TYPE_CHECKING:
     from sqlmesh.core.engine_adapter.base import EngineAdapter
 
 
+logger = logging.getLogger(__name__)
+
+
 class DataObjectType(str, Enum):
     UNKNOWN = "unknown"
     TABLE = "table"
@@ -290,8 +294,8 @@ def set_catalog(override_mapping: t.Optional[t.Dict[str, CatalogSupport]] = None
             container[key] = expression  # type: ignore
             if catalog_support.is_single_catalog_only:
                 if catalog_name != engine_adapter._default_catalog:
-                    raise UnsupportedCatalogOperationError(
-                        f"{engine_adapter.dialect} requires that all catalog operations be against a single catalog: {engine_adapter._default_catalog}"
+                    logger.warning(
+                        f"{engine_adapter.dialect} requires that all catalog operations be against a single catalog: {engine_adapter._default_catalog}. Ignoring catalog: {catalog_name}"
                     )
                 return func(*list_args, **kwargs)
             # Set the catalog name on the engine adapter if needed
diff --git a/sqlmesh/core/model/definition.py b/sqlmesh/core/model/definition.py
index 90bc570c..c24b8700 100644
--- a/sqlmesh/core/model/definition.py
+++ b/sqlmesh/core/model/definition.py
@@ -422,13 +422,14 @@ class _Model(ModelMeta, frozen=True):
         Return:
             The mocked out ctas query.
         """
-        # the query is expanded so it's been copied, it's safe to mutate.
-        query = self.render_query_or_raise(**render_kwarg)
+        query = self.render_query_or_raise(**render_kwarg).copy()
 
         for select in query.find_all(exp.Select):
             if select.args.get("from"):
                 select.where(exp.false(), copy=False)
-                if not isinstance(select.parent, exp.Union):
+                if not isinstance(select.parent, exp.Union) or (
+                    select.parent.parent is None and select.arg_key == "expression"
+                ):
                     select.limit(0, copy=False)
 
         if self.managed_columns:


Test patch:
diff --git a/tests/core/engine_adapter/test_integration.py b/tests/core/engine_adapter/test_integration.py
index ce474168..69e511a7 100644
--- a/tests/core/engine_adapter/test_integration.py
+++ b/tests/core/engine_adapter/test_integration.py
@@ -19,7 +19,6 @@ from sqlmesh.core.dialect import normalize_model_name
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.utils import random_id
 from sqlmesh.utils.date import now, to_date, to_ds, to_time_column, yesterday
-from sqlmesh.utils.errors import UnsupportedCatalogOperationError
 from sqlmesh.utils.pydantic import PydanticModel
 from tests.conftest import SushiDataValidator
 from tests.utils.pandas import compare_dataframes
@@ -577,7 +576,7 @@ def test_catalog_operations(ctx: TestContext):
     assert ctx.engine_adapter.get_current_catalog() == current_catalog
 
 
-def test_drop_schema_catalog(ctx: TestContext):
+def test_drop_schema_catalog(ctx: TestContext, caplog):
     def drop_schema_and_validate(schema_name: str):
         ctx.engine_adapter.drop_schema(schema_name, cascade=True)
         results = ctx.get_metadata_results(schema_name)
@@ -643,11 +642,8 @@ def test_drop_schema_catalog(ctx: TestContext):
 
     schema = ctx.schema("drop_schema_catalog_test", catalog_name)
     if ctx.engine_adapter.CATALOG_SUPPORT.is_single_catalog_only:
-        with pytest.raises(
-            UnsupportedCatalogOperationError,
-            match=".*requires that all catalog operations be against a single catalog.*",
-        ):
-            drop_schema_and_validate(schema)
+        drop_schema_and_validate(schema)
+        assert "requires that all catalog operations be against a single catalog" in caplog.text
         return
     drop_schema_and_validate(schema)
     create_objects_and_validate(schema)
diff --git a/tests/core/engine_adapter/test_postgres.py b/tests/core/engine_adapter/test_postgres.py
index 746db0aa..a5718c87 100644
--- a/tests/core/engine_adapter/test_postgres.py
+++ b/tests/core/engine_adapter/test_postgres.py
@@ -5,7 +5,6 @@ from pytest_mock import MockFixture
 from sqlglot.helper import ensure_list
 
 from sqlmesh.core.engine_adapter import PostgresEngineAdapter
-from sqlmesh.utils.errors import UnsupportedCatalogOperationError
 from tests.core.engine_adapter import to_sql_calls
 
 pytestmark = [pytest.mark.engine, pytest.mark.postgres]
@@ -52,10 +51,12 @@ def test_drop_schema(kwargs, expected, make_mocked_engine_adapter: t.Callable):
     assert to_sql_calls(adapter) == ensure_list(expected)
 
 
-def test_drop_schema_with_catalog(make_mocked_engine_adapter: t.Callable, mocker: MockFixture):
+def test_drop_schema_with_catalog(
+    make_mocked_engine_adapter: t.Callable, mocker: MockFixture, caplog
+):
     adapter = make_mocked_engine_adapter(PostgresEngineAdapter)
 
     adapter.get_current_catalog = mocker.MagicMock(return_value="other_catalog")
 
-    with pytest.raises(UnsupportedCatalogOperationError):
-        adapter.drop_schema("test_catalog.test_schema")
+    adapter.drop_schema("test_catalog.test_schema")
+    assert "requires that all catalog operations be against a single catalog" in caplog.text
diff --git a/tests/core/test_model.py b/tests/core/test_model.py
index d72afc72..9be879a0 100644
--- a/tests/core/test_model.py
+++ b/tests/core/test_model.py
@@ -1767,6 +1767,30 @@ def test_model_ctas_query():
         == 'SELECT 1 AS "a" FROM "b" AS "b" WHERE FALSE LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        SELECT 1 AS a FROM t UNION ALL SELECT 2 AS a FROM t UNION ALL SELECT 3 AS a FROM t UNION ALL SELECT 4 AS a FROM t
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 3 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 4 AS "a" FROM "t" AS "t" WHERE FALSE LIMIT 0'
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        SELECT 1 AS a FROM t UNION ALL SELECT 2 AS a FROM t
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})