Fix: Support of the table renaming in Snowflake and Databricks (#2509)

Author: izeigerman

sqlmesh/core/engine_adapter/databricks.py

@@ -6,6 +6,7 @@
 import pandas as pd
 from sqlglot import exp
 
+from sqlmesh.core.engine_adapter.mixins import RenameTableFullTargetNameMixin
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
     DataObject,
@@ -29,7 +30,7 @@
         "_get_data_objects": CatalogSupport.REQUIRES_SET_CATALOG,
     }
 )
-class DatabricksEngineAdapter(SparkEngineAdapter):
+class DatabricksEngineAdapter(SparkEngineAdapter, RenameTableFullTargetNameMixin):
     DIALECT = "databricks"
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.INSERT_OVERWRITE
     SUPPORTS_CLONING = True

sqlmesh/core/engine_adapter/mixins.py

@@ -246,3 +246,18 @@ def _truncate_table(self, table_name: TableName) -> None:
         if self._connection_pool.is_transaction_active:
             return self.execute(exp.Delete(this=exp.to_table(table_name)))
         super()._truncate_table(table_name)
+
+
+class RenameTableFullTargetNameMixin(EngineAdapter):
+    def _rename_table(
+        self,
+        old_table_name: TableName,
+        new_table_name: TableName,
+    ) -> None:
+        old_table = exp.to_table(old_table_name)
+        new_table = exp.to_table(new_table_name)
+        if not new_table.db and old_table.db:
+            # In MySQL and Snowflake you need to provide the full target table name.
+            # Therefore we use the old schema name if one is not provided explicitly in the new name.
+            new_table.set("db", old_table.args["db"])
+        return super()._rename_table(old_table, new_table)

sqlmesh/core/engine_adapter/mysql.py

@@ -10,6 +10,7 @@
     LogicalMergeMixin,
     NonTransactionalTruncateMixin,
     PandasNativeFetchDFSupportMixin,
+    RenameTableFullTargetNameMixin,
 )
 from sqlmesh.core.engine_adapter.shared import (
     CommentCreationTable,
@@ -30,6 +31,7 @@ class MySQLEngineAdapter(
     LogicalMergeMixin,
     PandasNativeFetchDFSupportMixin,
     NonTransactionalTruncateMixin,
+    RenameTableFullTargetNameMixin,
 ):
     DEFAULT_BATCH_SIZE = 200
     DIALECT = "mysql"
@@ -142,16 +144,3 @@ def _create_column_comments(
                 f"Column comments for table '{table.alias_or_name}' not registered - this may be due to limited permissions.",
                 exc_info=True,
             )
-
-    def _rename_table(
-        self,
-        old_table_name: TableName,
-        new_table_name: TableName,
-    ) -> None:
-        old_table = exp.to_table(old_table_name)
-        new_table = exp.to_table(new_table_name)
-        if not new_table.db and old_table.db:
-            # In MySQL you need to provide the full target table name.
-            # Therefore we use the old schema name if one is not provided explicitly in the new name.
-            new_table.set("db", old_table.args["db"])
-        return super()._rename_table(old_table, new_table)

sqlmesh/core/engine_adapter/snowflake.py

@@ -10,7 +10,10 @@
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.dialect import to_schema
-from sqlmesh.core.engine_adapter.mixins import GetCurrentCatalogFromFunctionMixin
+from sqlmesh.core.engine_adapter.mixins import (
+    GetCurrentCatalogFromFunctionMixin,
+    RenameTableFullTargetNameMixin,
+)
 from sqlmesh.core.engine_adapter.shared import (
     CatalogSupport,
     DataObject,
@@ -32,7 +35,7 @@
         "drop_schema": CatalogSupport.REQUIRES_SET_CATALOG,
     }
 )
-class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin):
+class SnowflakeEngineAdapter(GetCurrentCatalogFromFunctionMixin, RenameTableFullTargetNameMixin):
     DIALECT = "snowflake"
     SUPPORTS_MATERIALIZED_VIEWS = True
     SUPPORTS_MATERIALIZED_VIEW_SCHEMA = True

tests/core/engine_adapter/test_databricks.py

@@ -103,3 +103,18 @@ def test_get_current_database(make_mocked_engine_adapter: t.Callable):
 
     assert adapter.get_current_database() == "test_database"
     assert to_sql_calls(adapter) == ["SELECT CURRENT_DATABASE()"]
+
+
+def test_rename_table(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(DatabricksEngineAdapter)
+
+    adapter.rename_table("test_schema.old_name", "new_name")
+    adapter.rename_table("test_schema.old_name", "new_test_schema.new_name")
+    adapter.rename_table("old_name", "new_name")
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        "ALTER TABLE `test_schema`.`old_name` RENAME TO `test_schema`.`new_name`",
+        "ALTER TABLE `test_schema`.`old_name` RENAME TO `new_test_schema`.`new_name`",
+        "ALTER TABLE `old_name` RENAME TO `new_name`",
+    ]

tests/core/engine_adapter/test_snowflake.py

@@ -143,3 +143,18 @@ def test_comments(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture)
         """COMMENT ON TABLE "test_table" IS 'table description'""",
         """ALTER TABLE "test_table" ALTER COLUMN "a" COMMENT 'a column description'""",
     ]
+
+
+def test_rename_table(make_mocked_engine_adapter: t.Callable, mocker: MockerFixture):
+    adapter = make_mocked_engine_adapter(SnowflakeEngineAdapter)
+
+    adapter.rename_table("test_schema.old_name", "new_name")
+    adapter.rename_table("test_schema.old_name", "new_test_schema.new_name")
+    adapter.rename_table("old_name", "new_name")
+
+    sql_calls = to_sql_calls(adapter)
+    assert sql_calls == [
+        'ALTER TABLE "test_schema"."old_name" RENAME TO "test_schema"."new_name"',
+        'ALTER TABLE "test_schema"."old_name" RENAME TO "new_test_schema"."new_name"',
+        'ALTER TABLE "old_name" RENAME TO "new_name"',
+    ]