Snowflake: Only drop databases if SQLMesh was the one who created them

Author: erindru

sqlmesh/core/constants.py

@@ -7,6 +7,7 @@
 from pathlib import Path
 
 SQLMESH = "sqlmesh"
+SQLMESH_MANAGED = "sqlmesh_managed"
 SQLMESH_PATH = Path.home() / ".sqlmesh"
 
 PROD = "prod"

sqlmesh/core/engine_adapter/snowflake.py

@@ -9,6 +9,7 @@
 from sqlglot.optimizer.normalize_identifiers import normalize_identifiers
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
+import sqlmesh.core.constants as c
 from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.mixins import (
     GetCurrentCatalogFromFunctionMixin,
@@ -43,6 +44,7 @@
         "_get_data_objects": CatalogSupport.REQUIRES_SET_CATALOG,
         "create_schema": CatalogSupport.REQUIRES_SET_CATALOG,
         "drop_schema": CatalogSupport.REQUIRES_SET_CATALOG,
+        "drop_catalog": CatalogSupport.REQUIRES_SET_CATALOG,  # needs a catalog to issue a query to information_schema.databases even though the result is global
     }
 )
 class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, ClusteredByMixin, RowDiffMixin):
@@ -125,10 +127,34 @@ def catalog_support(self) -> CatalogSupport:
         return CatalogSupport.FULL_SUPPORT
 
     def _create_catalog(self, catalog_name: exp.Identifier) -> None:
-        self.execute(exp.Create(this=exp.Table(this=catalog_name), kind="DATABASE", exists=True))
+        props = exp.Properties(
+            expressions=[exp.SchemaCommentProperty(this=exp.Literal.string(c.SQLMESH_MANAGED))]
+        )
+        self.execute(
+            exp.Create(
+                this=exp.Table(this=catalog_name), kind="DATABASE", exists=True, properties=props
+            )
+        )
 
     def _drop_catalog(self, catalog_name: exp.Identifier) -> None:
-        self.execute(exp.Drop(this=exp.Table(this=catalog_name), kind="DATABASE", exists=True))
+        # only drop the catalog if it was created by SQLMesh, which is indicated by its comment matching {c.SQLMESH_MANAGED}
+        exists_check = (
+            exp.select(exp.Literal.number(1))
+            .from_(exp.to_table("information_schema.databases"))
+            .where(
+                exp.and_(
+                    exp.column("database_name").eq(exp.Literal.string(catalog_name)),
+                    exp.column("comment").eq(exp.Literal.string(c.SQLMESH_MANAGED)),
+                )
+            )
+        )
+        normalize_identifiers(exists_check, dialect=self.dialect)
+        if self.fetchone(exists_check, quote_identifiers=True) is not None:
+            self.execute(exp.Drop(this=exp.Table(this=catalog_name), kind="DATABASE", exists=True))
+        else:
+            logger.warning(
+                f"Not dropping database {catalog_name.sql(dialect=self.dialect)} because there is no indication it is '{c.SQLMESH_MANAGED}'"
+            )
 
     def _create_table(
         self,

tests/core/engine_adapter/integration/__init__.py

@@ -205,6 +205,9 @@ def __init__(
         self._schemas: t.List[
             str
         ] = []  # keep track of any schemas returned from self.schema() / self.table() so we can drop them at the end
+        self._catalogs: t.List[
+            str
+        ] = []  # keep track of any catalogs created via self.create_catalog() so we can drop them at the end
 
     @property
     def test_type(self) -> str:
@@ -685,6 +688,8 @@ def create_catalog(self, catalog_name: str):
             except Exception:
                 pass
 
+        self._catalogs.append(catalog_name)
+
     def drop_catalog(self, catalog_name: str):
         if self.dialect == "bigquery":
             return  # bigquery cannot create/drop catalogs
@@ -707,6 +712,9 @@ def cleanup(self, ctx: t.Optional[Context] = None):
                 schema_name=schema_name, ignore_if_not_exists=True, cascade=True
             )
 
+        for catalog_name in set(self._catalogs):
+            self.drop_catalog(catalog_name)
+
         self.engine_adapter.close()
 
     def upsert_sql_model(self, model_definition: str) -> t.Tuple[Context, SqlModel]:

tests/core/engine_adapter/integration/test_integration_snowflake.py

@@ -181,7 +181,7 @@ def _get_data_object(table: exp.Table) -> DataObject:
     assert not metadata.is_clustered
 
 
-def test_create_iceberg_table(ctx: TestContext, engine_adapter: SnowflakeEngineAdapter) -> None:
+def test_create_iceberg_table(ctx: TestContext) -> None:
     # Note: this test relies on a default Catalog and External Volume being configured in Snowflake
     # ref: https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-catalog-integration#set-a-default-catalog-at-the-account-database-or-schema-level
     # ref: https://docs.snowflake.com/en/user-guide/tables-iceberg-configure-external-volume#set-a-default-external-volume-at-the-account-database-or-schema-level
@@ -271,3 +271,39 @@ def execute(context: ExecutionContext, start: datetime, **kwargs) -> DataFrame:
     query = exp.select("*").from_(table)
     df = ctx.engine_adapter.fetchdf(query, quote_identifiers=True)
     assert len(df) == 10
+
+
+def test_create_drop_catalog(ctx: TestContext, engine_adapter: SnowflakeEngineAdapter):
+    non_sqlmesh_managed_catalog = ctx.add_test_suffix("external_catalog")
+    sqlmesh_managed_catalog = ctx.add_test_suffix("env_dev")
+
+    initial_catalog = engine_adapter.get_current_catalog()
+    assert initial_catalog
+
+    ctx.create_catalog(
+        non_sqlmesh_managed_catalog
+    )  # create via TestContext so the sqlmesh_managed comment doesnt get added
+    ctx._catalogs.append(sqlmesh_managed_catalog)  # so it still gets cleaned up if the test fails
+
+    engine_adapter.create_catalog(
+        sqlmesh_managed_catalog
+    )  # create via EngineAdapter so the sqlmesh_managed comment is added
+
+    def fetch_database_names() -> t.Set[str]:
+        engine_adapter.set_current_catalog(initial_catalog)
+        return {
+            str(r[0])
+            for r in engine_adapter.fetchall(
+                f"select database_name from information_schema.databases where database_name like '%{ctx.test_id}'"
+            )
+        }
+
+    assert fetch_database_names() == {non_sqlmesh_managed_catalog, sqlmesh_managed_catalog}
+
+    engine_adapter.drop_catalog(
+        non_sqlmesh_managed_catalog
+    )  # no-op: catalog is not SQLMesh-managed
+    assert fetch_database_names() == {non_sqlmesh_managed_catalog, sqlmesh_managed_catalog}
+
+    engine_adapter.drop_catalog(sqlmesh_managed_catalog)  # works, catalog is SQLMesh-managed
+    assert fetch_database_names() == {non_sqlmesh_managed_catalog}

tests/core/engine_adapter/test_snowflake.py

@@ -820,11 +820,16 @@ def test_create_catalog(snowflake_mocked_engine_adapter: SnowflakeEngineAdapter)
     adapter = snowflake_mocked_engine_adapter
     adapter.create_catalog(exp.to_identifier("foo"))
 
-    assert to_sql_calls(adapter) == ['CREATE DATABASE IF NOT EXISTS "foo"']
+    assert to_sql_calls(adapter) == [
+        "CREATE DATABASE IF NOT EXISTS \"foo\" COMMENT='sqlmesh_managed'"
+    ]
 
 
 def test_drop_catalog(snowflake_mocked_engine_adapter: SnowflakeEngineAdapter) -> None:
     adapter = snowflake_mocked_engine_adapter
     adapter.drop_catalog(exp.to_identifier("foo"))
 
-    assert to_sql_calls(adapter) == ['DROP DATABASE IF EXISTS "foo"']
+    assert to_sql_calls(adapter) == [
+        """SELECT 1 FROM "INFORMATION_SCHEMA"."DATABASES" WHERE "DATABASE_NAME" = 'foo' AND "COMMENT" = 'sqlmesh_managed'""",
+        'DROP DATABASE IF EXISTS "foo"',
+    ]