Add janitor test across all adapters and fix drop cascade in BigQuery

Author: erindru

sqlmesh/core/engine_adapter/bigquery.py

@@ -1260,6 +1260,21 @@ def _native_df_to_pandas_df(
 
         return super()._native_df_to_pandas_df(query_or_df)
 
+    def _drop_object(
+        self,
+        name: TableName | SchemaName,
+        exists: bool = True,
+        kind: str = "TABLE",
+        **drop_args: t.Any,
+    ) -> None:
+        if kind.upper() == "TABLE" and "cascade" in drop_args:
+            # BigQuery doesnt support DROP CASCADE for tables
+            # ref: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#drop_table_statement
+            # so set it to False here so SQLGlot doesnt output a CASCADE argument
+            drop_args["cascade"] = False
+
+        super()._drop_object(name=name, exists=exists, kind=kind, **drop_args)
+
     @property
     def _query_data(self) -> t.Any:
         return self._connection_pool.get_attribute("query_data")

tests/core/engine_adapter/integration/__init__.py

@@ -27,7 +27,7 @@
 from _pytest.mark.structures import ParameterSet
 
 if t.TYPE_CHECKING:
-    from sqlmesh.core._typing import TableName
+    from sqlmesh.core._typing import TableName, SchemaName
     from sqlmesh.core.engine_adapter._typing import Query
 
 TEST_SCHEMA = "test_schema"
@@ -222,6 +222,13 @@ def df_type(self) -> t.Optional[str]:
             return self._test_type.split("-", maxsplit=1)[1]
         return None
 
+    @property
+    def engine_type(self) -> str:
+        if self.mark.startswith("gcp_postgres"):
+            return "gcp_postgres"
+
+        return self.mark.split("_")[0]
+
     @property
     def columns_to_types(self):
         if self._columns_to_types is None:
@@ -307,7 +314,7 @@ def default_table_format(self) -> t.Optional[str]:
     def add_test_suffix(self, value: str) -> str:
         return f"{value}_{self.test_id}"
 
-    def get_metadata_results(self, schema: t.Optional[str] = None) -> MetadataResults:
+    def get_metadata_results(self, schema: t.Optional[SchemaName] = None) -> MetadataResults:
         schema = schema if schema else self.schema(TEST_SCHEMA)
         return MetadataResults.from_data_objects(self.engine_adapter.get_data_objects(schema))
 

tests/core/engine_adapter/integration/test_integration.py

@@ -21,6 +21,7 @@
 from sqlmesh.core.config import load_config_from_paths
 from sqlmesh.core.config.connection import ConnectionConfig
 import sqlmesh.core.dialect as d
+from sqlmesh.core.environment import EnvironmentSuffixTarget
 from sqlmesh.core.dialect import select_from_values
 from sqlmesh.core.model import Model, load_sql_based_model
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
@@ -2333,11 +2334,7 @@ def _normalize_snowflake(name: str, prefix_regex: str = "(sqlmesh__)(.*)"):
             k: [_normalize_snowflake(name) for name in v] for k, v in object_names.items()
         }
 
-    if ctx.mark.startswith("gcp_postgres"):
-        engine_type = "gcp_postgres"
-    else:
-        engine_type = ctx.mark.split("_")[0]
-    init_example_project(tmp_path, engine_type, schema_name=schema_name)
+    init_example_project(tmp_path, ctx.engine_type, schema_name=schema_name)
 
     config = load_config_from_paths(
         Config,
@@ -3557,3 +3554,137 @@ def test_identifier_length_limit(ctx: TestContext):
         match=re.escape(match),
     ):
         adapter.create_table(long_table_name, {"col": exp.DataType.build("int")})
+
+
+@pytest.mark.parametrize(
+    "environment_suffix_target",
+    [
+        EnvironmentSuffixTarget.TABLE,
+        EnvironmentSuffixTarget.SCHEMA,
+        EnvironmentSuffixTarget.CATALOG,
+    ],
+)
+def test_janitor(
+    ctx: TestContext, tmp_path: pathlib.Path, environment_suffix_target: EnvironmentSuffixTarget
+):
+    if (
+        environment_suffix_target == EnvironmentSuffixTarget.CATALOG
+        and not ctx.engine_adapter.SUPPORTS_CREATE_DROP_CATALOG
+    ):
+        pytest.skip("Engine does not support catalog-based virtual environments")
+
+    schema = ctx.schema()  # catalog.schema
+    parsed_schema = d.to_schema(schema)
+
+    init_example_project(tmp_path, ctx.engine_type, schema_name=parsed_schema.db)
+
+    def _set_config(_gateway: str, config: Config) -> None:
+        config.environment_suffix_target = environment_suffix_target
+        config.model_defaults.dialect = ctx.dialect
+
+    sqlmesh = ctx.create_context(path=tmp_path, config_mutator=_set_config)
+
+    sqlmesh.plan(auto_apply=True)
+
+    # create a new model in dev
+    (tmp_path / "models" / "new_model.sql").write_text(f"""
+        MODEL (
+            name {schema}.new_model,
+            kind FULL
+        );
+
+        select * from {schema}.full_model
+    """)
+    sqlmesh.load()
+
+    result = sqlmesh.plan(environment="dev", auto_apply=True)
+    assert result.context_diff.is_new_environment
+    assert len(result.context_diff.new_snapshots) == 1
+    new_model = list(result.context_diff.new_snapshots.values())[0]
+    assert "new_model" in new_model.name.lower()
+
+    # check physical objects
+    snapshot_table_name = exp.to_table(new_model.table_name(), dialect=ctx.dialect)
+    snapshot_schema = snapshot_table_name.db
+
+    prod_schema = normalize_identifiers(d.to_schema(schema), dialect=ctx.dialect)
+    dev_env_schema = prod_schema.copy()
+    if environment_suffix_target == EnvironmentSuffixTarget.CATALOG:
+        dev_env_schema.set("catalog", exp.to_identifier(f"{prod_schema.catalog}__dev"))
+    else:
+        dev_env_schema.set("db", exp.to_identifier(f"{prod_schema.db}__dev"))
+    normalize_identifiers(dev_env_schema, dialect=ctx.dialect)
+
+    md = ctx.get_metadata_results(prod_schema)
+    if environment_suffix_target == EnvironmentSuffixTarget.TABLE:
+        assert sorted([v.lower() for v in md.views]) == [
+            "full_model",
+            "incremental_model",
+            "new_model__dev",
+            "seed_model",
+        ]
+    else:
+        assert sorted([v.lower() for v in md.views]) == [
+            "full_model",
+            "incremental_model",
+            "seed_model",
+        ]
+    assert not md.tables
+    assert not md.managed_tables
+
+    if environment_suffix_target != EnvironmentSuffixTarget.TABLE:
+        # note: this is "catalog__dev.schema" for EnvironmentSuffixTarget.CATALOG and "catalog.schema__dev" for EnvironmentSuffixTarget.SCHEMA
+        md = ctx.get_metadata_results(dev_env_schema)
+        assert [v.lower() for v in md.views] == ["new_model"]
+        assert not md.tables
+        assert not md.managed_tables
+
+    md = ctx.get_metadata_results(snapshot_schema)
+    assert not md.views
+    assert not md.managed_tables
+    assert sorted(t.split("__")[1].lower() for t in md.tables) == [
+        "full_model",
+        "incremental_model",
+        "new_model",
+        "seed_model",
+    ]
+
+    # invalidate dev and run the janitor to clean it up
+    sqlmesh.invalidate_environment("dev")
+    assert sqlmesh.run_janitor(
+        ignore_ttl=True
+    )  # ignore_ttl to delete the new_model snapshot even though it hasnt expired yet
+
+    # there should be no dev environment or dev tables / schemas
+    md = ctx.get_metadata_results(prod_schema)
+    assert sorted([v.lower() for v in md.views]) == [
+        "full_model",
+        "incremental_model",
+        "seed_model",
+    ]
+    assert not md.tables
+    assert not md.managed_tables
+
+    if environment_suffix_target != EnvironmentSuffixTarget.TABLE:
+        if environment_suffix_target == EnvironmentSuffixTarget.SCHEMA:
+            md = ctx.get_metadata_results(dev_env_schema)
+        else:
+            try:
+                md = ctx.get_metadata_results(dev_env_schema)
+            except Exception as e:
+                # Most engines will raise an error when @set_catalog tries to set a catalog that doesnt exist
+                # in this case, we just swallow the error. We know this call already worked before in the earlier checks
+                md = MetadataResults()
+
+        assert not md.views
+        assert not md.tables
+        assert not md.managed_tables
+
+    md = ctx.get_metadata_results(snapshot_schema)
+    assert not md.views
+    assert not md.managed_tables
+    assert sorted(t.split("__")[1].lower() for t in md.tables) == [
+        "full_model",
+        "incremental_model",
+        "seed_model",
+    ]

tests/core/engine_adapter/test_bigquery.py

@@ -20,8 +20,10 @@
 
 
 @pytest.fixture
-def adapter(make_mocked_engine_adapter: t.Callable) -> BigQueryEngineAdapter:
-    return make_mocked_engine_adapter(BigQueryEngineAdapter)
+def adapter(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture) -> BigQueryEngineAdapter:
+    mocked_adapter = make_mocked_engine_adapter(BigQueryEngineAdapter)
+    mocker.patch("sqlmesh.core.engine_adapter.bigquery.BigQueryEngineAdapter.execute")
+    return mocked_adapter
 
 
 def test_insert_overwrite_by_time_partition_query(
@@ -575,6 +577,8 @@ def test_begin_end_session(mocker: MockerFixture):
 
 
 def _to_sql_calls(execute_mock: t.Any, identify: bool = True) -> t.List[str]:
+    if isinstance(execute_mock, BigQueryEngineAdapter):
+        execute_mock = execute_mock.execute
     output = []
     for call in execute_mock.call_args_list:
         value = call[0][0]
@@ -1150,3 +1154,22 @@ def test_job_cancellation_on_keyboard_interrupt_job_already_done(mocker: MockerF
     # Verify job status was checked but cancellation was NOT called
     mock_job.done.assert_called_once()
     mock_job.cancel.assert_not_called()
+
+
+def test_drop_cascade(adapter: BigQueryEngineAdapter):
+    adapter.drop_table("foo", cascade=True)
+    adapter.drop_table("foo", cascade=False)
+
+    # BigQuery doesnt support DROP CASCADE for tables
+    # ref: https://cloud.google.com/bigquery/docs/reference/standard-sql/data-definition-language#drop_table_statement
+    assert _to_sql_calls(adapter) == ["DROP TABLE IF EXISTS `foo`", "DROP TABLE IF EXISTS `foo`"]
+    adapter.execute.reset_mock()  # type: ignore
+
+    # But, it does for schemas
+    adapter.drop_schema("foo", cascade=True)
+    adapter.drop_schema("foo", cascade=False)
+
+    assert _to_sql_calls(adapter) == [
+        "DROP SCHEMA IF EXISTS `foo` CASCADE",
+        "DROP SCHEMA IF EXISTS `foo`",
+    ]