Fix: Fix Transactions with Truncate (#1880)

* fix transactions with truncate

* fix tests

(cherry picked from commit a6d8d90)

Author: eakmanrq

sqlmesh/core/config/connection.py

@@ -265,6 +265,10 @@ def _connection_kwargs_keys(self) -> t.Set[str]:
     def _engine_adapter(self) -> t.Type[EngineAdapter]:
         return engine_adapter.SnowflakeEngineAdapter
 
+    @property
+    def _static_connection_kwargs(self) -> t.Dict[str, t.Any]:
+        return {"autocommit": False}
+
     @property
     def _connection_factory(self) -> t.Callable:
         from snowflake import connector

sqlmesh/core/engine_adapter/base.py

@@ -1551,9 +1551,9 @@ def _add_where_to_query(
 
         return query
 
-    def _truncate_table(self, table_name: TableName) -> str:
-        table = quote_identifiers(exp.to_table(table_name))
-        return f"TRUNCATE {table.sql(dialect=self.dialect)}"
+    def _truncate_table(self, table_name: TableName) -> None:
+        table = exp.to_table(table_name)
+        self.execute(f"TRUNCATE TABLE {table.sql(dialect=self.dialect, identify=True)}")
 
 
 class EngineAdapterWithIndexSupport(EngineAdapter):

sqlmesh/core/engine_adapter/mixins.py

@@ -100,7 +100,7 @@ def replace_table(
             temp_query = query.transform(
                 replace_table, curr_table=target_table, new_table=temp_table
             )
-            engine_adapter.execute(engine_adapter._truncate_table(target_table))
+            engine_adapter._truncate_table(target_table)
             return engine_adapter._insert_append_query(target_table, temp_query, columns_to_types)
 
     def replace_query(
@@ -138,7 +138,7 @@ def replace_query(
                     return self.overwrite_target_from_temp(
                         self, query, columns_to_types, target_table
                     )
-                self.execute(self._truncate_table(table_name))
+                self._truncate_table(table_name)
                 return self._insert_append_query(table_name, query, columns_to_types)
 
 
@@ -260,3 +260,12 @@ def get_current_catalog(self) -> t.Optional[str]:
         if result:
             return result[0]
         return None
+
+
+class NonTransactionalTruncateMixin(EngineAdapter):
+    def _truncate_table(self, table_name: TableName) -> None:
+        # Truncate forces a commit of the current transaction so we want to do an unconditional delete to
+        # preserve the transaction if one exists otherwise we can truncate
+        if self._connection_pool.is_transaction_active:
+            return self.execute(exp.Delete(this=exp.to_table(table_name)))
+        super()._truncate_table(table_name)

sqlmesh/core/engine_adapter/mssql.py

@@ -8,7 +8,6 @@
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
-from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.engine_adapter.base import (
     CatalogSupport,
@@ -243,10 +242,6 @@ def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
             for row in dataframe.itertuples()
         ]
 
-    def _truncate_table(self, table_name: TableName) -> str:
-        table = quote_identifiers(exp.to_table(table_name))
-        return f"TRUNCATE TABLE {table.sql(dialect=self.dialect)}"
-
     def _to_sql(self, expression: exp.Expression, quote: bool = True, **kwargs: t.Any) -> str:
         sql = super()._to_sql(expression, quote=quote, **kwargs)
         return f"{sql};"

sqlmesh/core/engine_adapter/mysql.py

@@ -5,6 +5,7 @@
 from sqlmesh.core.engine_adapter.mixins import (
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
+    NonTransactionalTruncateMixin,
     PandasNativeFetchDFSupportMixin,
 )
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType, set_catalog
@@ -17,6 +18,7 @@ class MySQLEngineAdapter(
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
     PandasNativeFetchDFSupportMixin,
+    NonTransactionalTruncateMixin,
 ):
     DEFAULT_BATCH_SIZE = 200
     DIALECT = "mysql"

sqlmesh/core/engine_adapter/redshift.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import typing as t
-import uuid
 
 import pandas as pd
 from sqlglot import exp
@@ -12,6 +11,7 @@
     GetCurrentCatalogFromFunctionMixin,
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
+    NonTransactionalTruncateMixin,
 )
 from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType, set_catalog
 
@@ -25,6 +25,7 @@ class RedshiftEngineAdapter(
     LogicalReplaceQueryMixin,
     LogicalMergeMixin,
     GetCurrentCatalogFromFunctionMixin,
+    NonTransactionalTruncateMixin,
 ):
     DIALECT = "redshift"
     ESCAPE_JSON = True
@@ -180,12 +181,8 @@ def replace_query(
         columns_to_types = columns_to_types or self.columns(table_name)
         target_table = exp.to_table(table_name)
         with self.transaction():
-            temp_table_name = f"{target_table.alias_or_name}_temp_{self._short_hash()}"
-            temp_table = target_table.copy()
-            temp_table.set("this", exp.to_identifier(temp_table_name))
-            old_table_name = f"{target_table.alias_or_name}_old_{self._short_hash()}"
-            old_table = target_table.copy()
-            old_table.set("this", exp.to_identifier(old_table_name))
+            temp_table = self._get_temp_table(target_table)
+            old_table = self._get_temp_table(target_table)
             self.create_table(temp_table, columns_to_types, exists=False, **kwargs)
             self._insert_append_source_queries(temp_table, source_queries, columns_to_types)
             self.rename_table(target_table, old_table)
@@ -200,17 +197,18 @@ def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
         """
         Returns all the data objects that exist in the given schema and optionally catalog.
         """
+        catalog_name = self.get_current_catalog()
         query = f"""
             SELECT
-                null AS catalog_name,
+                '{catalog_name}' AS catalog_name,
                 tablename AS name,
                 schemaname AS schema_name,
                 'TABLE' AS type
             FROM pg_tables
             WHERE schemaname ILIKE '{schema_name}'
             UNION ALL
             SELECT
-                null AS catalog_name,
+                '{catalog_name}' AS catalog_name,
                 viewname AS name,
                 schemaname AS schema_name,
                 'VIEW' AS type
@@ -219,7 +217,7 @@ def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
             AND definition not ilike '%create materialized view%'
             UNION ALL
             SELECT
-                null AS catalog_name,
+                '{catalog_name}' AS catalog_name,
                 viewname AS name,
                 schemaname AS schema_name,
                 'MATERIALIZED_VIEW' AS type
@@ -236,9 +234,6 @@ def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
             for row in df.itertuples()
         ]
 
-    def _short_hash(self) -> str:
-        return uuid.uuid4().hex[:8]
-
 
 def parse_plan(plan: str) -> t.Optional[t.Dict]:
     """Parse the output of a redshift explain verbose query plan into a Python dict."""

sqlmesh/core/engine_adapter/spark.py

@@ -500,10 +500,6 @@ def wap_publish(self, table_name: TableName, wap_id: str) -> None:
         )
         self.execute(f"ALTER TABLE {fqn.sql(dialect=self.dialect)} DROP BRANCH {branch_name}")
 
-    def _truncate_table(self, table_name: TableName) -> str:
-        table = quote_identifiers(exp.to_table(table_name))
-        return f"TRUNCATE TABLE {table.sql(dialect=self.dialect)}"
-
     def _ensure_fqn(self, table_name: TableName) -> exp.Table:
         if isinstance(table_name, exp.Table):
             table_name = table_name.copy()

sqlmesh/core/engine_adapter/trino.py

@@ -5,7 +5,6 @@
 import pandas as pd
 from pandas.api.types import is_datetime64_any_dtype  # type: ignore
 from sqlglot import exp
-from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.dialect import to_schema
 from sqlmesh.core.engine_adapter.base import (
@@ -68,9 +67,10 @@ def _insert_overwrite_by_condition(
             f"SET SESSION {self.get_current_catalog()}.insert_existing_partitions_behavior='APPEND'"
         )
 
-    def _truncate_table(self, table_name: TableName) -> str:
-        table = quote_identifiers(exp.to_table(table_name))
-        return f"DELETE FROM {table.sql(dialect=self.dialect)}"
+    def _truncate_table(self, table_name: TableName) -> None:
+        table = exp.to_table(table_name)
+        # Some trino connectors don't support truncate so we use delete.
+        self.execute(f"DELETE FROM {table.sql(dialect=self.dialect, identify=True)}")
 
     def _get_data_objects(self, schema_name: SchemaName) -> t.List[DataObject]:
         """

tests/core/engine_adapter/test_integration.py

@@ -883,6 +883,53 @@ def test_scd_type_2(ctx: TestContext):
     )
 
 
+def test_truncate_table(ctx: TestContext):
+    if ctx.test_type != "query":
+        pytest.skip("Truncate table test does not change based on input data type")
+
+    ctx.init()
+    table = ctx.table("test_table")
+    ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+    input_data = pd.DataFrame(
+        [
+            {"id": 1, "ds": "2022-01-01"},
+            {"id": 2, "ds": "2022-01-02"},
+            {"id": 3, "ds": "2022-01-03"},
+        ]
+    )
+    ctx.engine_adapter.insert_append(table, ctx.input_data(input_data))
+    ctx.compare_with_current(table, input_data)
+    ctx.engine_adapter._truncate_table(table)
+    assert ctx.engine_adapter.fetchone(exp.select("count(*)").from_(table))[0] == 0
+
+
+def test_transaction(ctx: TestContext):
+    if ctx.engine_adapter.SUPPORTS_TRANSACTIONS is False:
+        pytest.skip(f"Engine adapter {ctx.engine_adapter.dialect} doesn't support transactions")
+    if ctx.test_type != "query":
+        pytest.skip("Transaction test can just run for query")
+
+    ctx.init()
+    table = ctx.table("test_table")
+    input_data = pd.DataFrame(
+        [
+            {"id": 1, "ds": "2022-01-01"},
+            {"id": 2, "ds": "2022-01-02"},
+            {"id": 3, "ds": "2022-01-03"},
+        ]
+    )
+    with ctx.engine_adapter.transaction():
+        ctx.engine_adapter.create_table(table, ctx.columns_to_types)
+        ctx.engine_adapter.insert_append(
+            table, ctx.input_data(input_data, ctx.columns_to_types), ctx.columns_to_types
+        )
+    ctx.compare_with_current(table, input_data)
+    with ctx.engine_adapter.transaction():
+        ctx.engine_adapter._truncate_table(table)
+        ctx.engine_adapter._connection_pool.rollback()
+    ctx.compare_with_current(table, input_data)
+
+
 def test_sushi(ctx: TestContext):
     if ctx.test_type != "query":
         pytest.skip("Sushi end-to-end tests only need to run for query and pyspark tests")

tests/core/engine_adapter/test_mixins.py

@@ -8,6 +8,7 @@
 from sqlmesh.core.engine_adapter.mixins import (
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
+    NonTransactionalTruncateMixin,
 )
 from tests.core.engine_adapter import to_sql_calls
 
@@ -30,7 +31,7 @@ def test_logical_replace_query_already_exists(
     adapter.replace_query("db.table", parse_one("SELECT col FROM db.other_table"))
 
     assert to_sql_calls(adapter) == [
-        'TRUNCATE "db"."table"',
+        'TRUNCATE TABLE "db"."table"',
         'INSERT INTO "db"."table" ("col") SELECT "col" FROM "db"."other_table"',
     ]
 
@@ -76,7 +77,7 @@ def test_logical_replace_self_reference(
     assert to_sql_calls(adapter) == [
         f'CREATE SCHEMA IF NOT EXISTS "db"',
         f'CREATE TABLE IF NOT EXISTS "db"."__temp_table_{temp_table_id}" AS SELECT "col" FROM "db"."table"',
-        'TRUNCATE "db"."table"',
+        'TRUNCATE TABLE "db"."table"',
         f'INSERT INTO "db"."table" ("col") SELECT "col" + 1 AS "col" FROM "db"."__temp_table_{temp_table_id}"',
         f'DROP TABLE IF EXISTS "db"."__temp_table_{temp_table_id}"',
     ]
@@ -135,3 +136,23 @@ def test_logical_merge(make_mocked_engine_adapter: t.Callable, mocker: MockerFix
             call('''DROP TABLE IF EXISTS "temporary"'''),
         ]
     )
+
+
+def test_non_transaction_truncate_mixin(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(NonTransactionalTruncateMixin, "redshift")
+    adapter._truncate_table(table_name="test_table")
+
+    assert to_sql_calls(adapter) == ['TRUNCATE TABLE "test_table"']
+
+
+def test_non_transaction_truncate_mixin_within_transaction(
+    make_mocked_engine_adapter: t.Callable, mocker: MockerFixture, make_temp_table_name: t.Callable
+):
+    adapter = make_mocked_engine_adapter(NonTransactionalTruncateMixin, "redshift")
+    adapter._connection_pool = mocker.MagicMock()
+    adapter._connection_pool.is_transaction_active = True
+    adapter._truncate_table(table_name="test_table")
+
+    assert to_sql_calls(adapter) == ['DELETE FROM "test_table"']