Fix(mssql): execute within transactions to prevent hanging (#1532)

* Remove transactiontype

* Wrap execute in a transaction

* disable transactions duckdb

* Remove duckdb transaction test

* Fix _transaction typing

---------

Co-authored-by: eakmanrq <6326532+eakmanrq@users.noreply.github.com>

Author: treysp

sqlmesh/core/engine_adapter/__init__.py

@@ -13,7 +13,6 @@
 from sqlmesh.core.engine_adapter.mysql import MySQLEngineAdapter
 from sqlmesh.core.engine_adapter.postgres import PostgresEngineAdapter
 from sqlmesh.core.engine_adapter.redshift import RedshiftEngineAdapter
-from sqlmesh.core.engine_adapter.shared import TransactionType
 from sqlmesh.core.engine_adapter.snowflake import SnowflakeEngineAdapter
 from sqlmesh.core.engine_adapter.spark import SparkEngineAdapter
 

sqlmesh/core/engine_adapter/base.py

@@ -26,7 +26,7 @@
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
 from sqlmesh.core.dialect import add_table, select_from_values_for_batch_range
-from sqlmesh.core.engine_adapter.shared import DataObject, TransactionType
+from sqlmesh.core.engine_adapter.shared import DataObject
 from sqlmesh.core.model.kind import TimeColumn
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils import double_escape
@@ -117,6 +117,7 @@ class EngineAdapter:
     DEFAULT_BATCH_SIZE = 10000
     DEFAULT_SQL_GEN_KWARGS: t.Dict[str, str | bool | int] = {}
     ESCAPE_JSON = False
+    SUPPORTS_TRANSACTIONS = True
     SUPPORTS_INDEXES = False
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.DELETE_INSERT
     SUPPORTS_MATERIALIZED_VIEWS = False
@@ -538,7 +539,7 @@ def alter_table(
         """
         Performs the required alter statements to change the current table into the structure of the target table.
         """
-        with self.transaction(TransactionType.DDL):
+        with self.transaction():
             for alter_expression in self.SCHEMA_DIFFER.compare_columns(
                 current_table_name,
                 self.columns(current_table_name),
@@ -1216,13 +1217,12 @@ def fetch_pyspark_df(
     @contextlib.contextmanager
     def transaction(
         self,
-        transaction_type: TransactionType = TransactionType.DML,
         condition: t.Optional[bool] = None,
     ) -> t.Iterator[None]:
         """A transaction context manager."""
         if (
             self._connection_pool.is_transaction_active
-            or not self.supports_transactions(transaction_type)
+            or not self.SUPPORTS_TRANSACTIONS
             or (condition is not None and not condition)
         ):
             yield
@@ -1236,10 +1236,6 @@ def transaction(
         else:
             self._connection_pool.commit()
 
-    def supports_transactions(self, transaction_type: TransactionType) -> bool:
-        """Whether or not the engine adapter supports transactions for the given transaction type."""
-        return True
-
     @contextlib.contextmanager
     def session(self) -> t.Iterator[None]:
         """A session context manager."""
@@ -1275,14 +1271,15 @@ def execute(
             {"unsupported_level": ErrorLevel.IGNORE} if ignore_unsupported_errors else {}
         )
 
-        for e in ensure_list(expressions):
-            sql = (
-                self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs)
-                if isinstance(e, exp.Expression)
-                else e
-            )
-            logger.debug(f"Executing SQL:\n{sql}")
-            self.cursor.execute(sql, **kwargs)
+        with self.transaction():
+            for e in ensure_list(expressions):
+                sql = (
+                    self._to_sql(e, quote=quote_identifiers, **to_sql_kwargs)
+                    if isinstance(e, exp.Expression)
+                    else e
+                )
+                logger.debug(f"Executing SQL:\n{sql}")
+                self.cursor.execute(sql, **kwargs)
 
     @contextlib.contextmanager
     def temp_table(
@@ -1307,7 +1304,7 @@ def temp_table(
         source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
             query_or_df, columns_to_types=columns_to_types, target_table=name
         )
-        with self.transaction(TransactionType.DDL):
+        with self.transaction():
             table = self._get_temp_table(name)
             if table.db:
                 self.create_schema(table.db)

sqlmesh/core/engine_adapter/base_postgres.py

@@ -4,20 +4,16 @@
 
 from sqlglot import exp
 
-from sqlmesh.core.engine_adapter.mixins import CommitOnExecuteMixin
-from sqlmesh.core.engine_adapter.shared import (
-    DataObject,
-    DataObjectType,
-    TransactionType,
-)
+from sqlmesh.core.engine_adapter.mixins import EngineAdapter
+from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
     from sqlmesh.core._typing import TableName
     from sqlmesh.core.engine_adapter.base import QueryOrDF
 
 
-class BasePostgresEngineAdapter(CommitOnExecuteMixin):
+class BasePostgresEngineAdapter(EngineAdapter):
     COLUMNS_TABLE = "information_schema.columns"
 
     def columns(
@@ -85,7 +81,7 @@ def create_view(
 
         Reference: https://www.postgresql.org/docs/current/sql-createview.html
         """
-        with self.transaction(TransactionType.DDL):
+        with self.transaction():
             if replace:
                 self.drop_view(view_name, materialized=materialized)
             super().create_view(

sqlmesh/core/engine_adapter/bigquery.py

@@ -11,11 +11,7 @@
 
 from sqlmesh.core.engine_adapter.base import SourceQuery
 from sqlmesh.core.engine_adapter.mixins import InsertOverwriteWithMergeMixin
-from sqlmesh.core.engine_adapter.shared import (
-    DataObject,
-    DataObjectType,
-    TransactionType,
-)
+from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.schema_diff import SchemaDiffer
 from sqlmesh.utils.date import to_datetime
@@ -45,6 +41,7 @@ class BigQueryEngineAdapter(InsertOverwriteWithMergeMixin):
     DIALECT = "bigquery"
     DEFAULT_BATCH_SIZE = 1000
     ESCAPE_JSON = True
+    SUPPORTS_TRANSACTIONS = False
     SUPPORTS_MATERIALIZED_VIEWS = True
     SUPPORTS_CLONING = True
 
@@ -478,9 +475,6 @@ def create_state_table(
             columns_to_types,
         )
 
-    def supports_transactions(self, transaction_type: TransactionType) -> bool:
-        return False
-
     def _db_call(self, func: t.Callable[..., t.Any], *args: t.Any, **kwargs: t.Any) -> t.Any:
         return func(
             retry=self.__retry,

sqlmesh/core/engine_adapter/duckdb.py

@@ -16,6 +16,7 @@
 
 class DuckDBEngineAdapter(LogicalMergeMixin):
     DIALECT = "duckdb"
+    SUPPORTS_TRANSACTIONS = False
 
     def _df_to_source_queries(
         self,

sqlmesh/core/engine_adapter/mixins.py

@@ -6,7 +6,7 @@
 from sqlglot import exp
 from sqlglot.optimizer.qualify_columns import quote_identifiers
 
-from sqlmesh.core.engine_adapter.base import EngineAdapter, SourceQuery, TransactionType
+from sqlmesh.core.engine_adapter.base import EngineAdapter, SourceQuery
 from sqlmesh.utils.errors import SQLMeshError
 
 if t.TYPE_CHECKING:
@@ -42,7 +42,7 @@ def merge(
         unique_exp = exp.func("CONCAT_WS", "'__SQLMESH_DELIM__'", *unique_key)
         column_names = list(columns_to_types or [])
 
-        with self.transaction(TransactionType.DML):
+        with self.transaction():
             self.ctas(temp_table, source_table, columns_to_types=columns_to_types, exists=False)
             self.execute(
                 exp.delete(target_table).where(
@@ -77,7 +77,7 @@ def replace_query(
 
         if not self.table_exists(table_name):
             return self.ctas(table_name, query_or_df, columns_to_types, exists=False, **kwargs)
-        with self.transaction(TransactionType.DDL):
+        with self.transaction():
             # TODO: remove quote_identifiers when sqlglot has an expression to represent TRUNCATE
             source_queries, columns_to_types = self._get_source_queries_and_columns_to_types(
                 query_or_df, columns_to_types, table_name
@@ -125,30 +125,6 @@ def _fetch_native_df(
         return read_sql_query(sql, self._connection_pool.get())
 
 
-class CommitOnExecuteMixin(EngineAdapter):
-    def execute(
-        self,
-        expressions: t.Union[str, exp.Expression, t.Sequence[exp.Expression]],
-        ignore_unsupported_errors: bool = False,
-        quote_identifiers: bool = True,
-        **kwargs: t.Any,
-    ) -> None:
-        """
-        To make sure that inserts and updates take effect we need to commit explicitly unless the
-        statement is executed as part of an active transaction.
-
-        Reference: https://www.psycopg.org/psycopg3/docs/basic/transactions.html
-        """
-        super().execute(
-            expressions,
-            ignore_unsupported_errors=ignore_unsupported_errors,
-            quote_identifiers=quote_identifiers,
-            **kwargs,
-        )
-        if not self._connection_pool.is_transaction_active:
-            self._connection_pool.commit()
-
-
 class InsertOverwriteWithMergeMixin(EngineAdapter):
     def _insert_overwrite_by_condition(
         self,

sqlmesh/core/engine_adapter/mysql.py

@@ -3,7 +3,6 @@
 import typing as t
 
 from sqlmesh.core.engine_adapter.mixins import (
-    CommitOnExecuteMixin,
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
     PandasNativeFetchDFSupportMixin,
@@ -15,7 +14,6 @@
 
 
 class MySQLEngineAdapter(
-    CommitOnExecuteMixin,
     LogicalMergeMixin,
     LogicalReplaceQueryMixin,
     PandasNativeFetchDFSupportMixin,
@@ -54,8 +52,8 @@ def _get_data_objects(
                 null AS catalog_name,
                 table_name AS name,
                 table_schema AS schema_name,
-                CASE 
-                    WHEN table_type = 'BASE TABLE' THEN 'table' 
+                CASE
+                    WHEN table_type = 'BASE TABLE' THEN 'table'
                     WHEN table_type = 'VIEW' THEN 'view'
                     ELSE table_type
                 END AS type

sqlmesh/core/engine_adapter/shared.py

@@ -8,19 +8,6 @@
 from sqlmesh.utils.pydantic import PydanticModel
 
 
-class TransactionType(str, Enum):
-    DDL = "DDL"
-    DML = "DML"
-
-    @property
-    def is_ddl(self) -> bool:
-        return self == TransactionType.DDL
-
-    @property
-    def is_dml(self) -> bool:
-        return self == TransactionType.DML
-
-
 class DataObjectType(str, Enum):
     UNKNOWN = "unknown"
     TABLE = "table"

sqlmesh/core/engine_adapter/spark.py

@@ -12,11 +12,7 @@
     InsertOverwriteStrategy,
     SourceQuery,
 )
-from sqlmesh.core.engine_adapter.shared import (
-    DataObject,
-    DataObjectType,
-    TransactionType,
-)
+from sqlmesh.core.engine_adapter.shared import DataObject, DataObjectType
 from sqlmesh.utils import classproperty, nullsafe_join
 from sqlmesh.utils.errors import SQLMeshError
 
@@ -40,6 +36,7 @@
 class SparkEngineAdapter(EngineAdapter):
     DIALECT = "spark"
     ESCAPE_JSON = True
+    SUPPORTS_TRANSACTIONS = False
     INSERT_OVERWRITE_STRATEGY = InsertOverwriteStrategy.INSERT_OVERWRITE
 
     @property
@@ -401,6 +398,3 @@ def __table_properties_to_expressions(
         return [
             exp.Property(this=key, value=value.copy()) for key, value in table_properties.items()
         ]
-
-    def supports_transactions(self, transaction_type: TransactionType) -> bool:
-        return False

sqlmesh/core/snapshot/evaluator.py

@@ -32,7 +32,7 @@
 from sqlglot.executor import execute
 
 from sqlmesh.core.audit import Audit, AuditResult
-from sqlmesh.core.engine_adapter import EngineAdapter, TransactionType
+from sqlmesh.core.engine_adapter import EngineAdapter
 from sqlmesh.core.engine_adapter.base import InsertOverwriteStrategy
 from sqlmesh.core.model import IncrementalUnmanagedKind, Model, SCDType2Kind, ViewKind
 from sqlmesh.core.snapshot import (
@@ -160,11 +160,7 @@ def apply(query_or_df: QueryOrDF, index: int = 0) -> None:
             **common_render_kwargs,
         )
 
-        with self.adapter.transaction(
-            transaction_type=TransactionType.DDL
-            if model.kind.is_view or model.kind.is_full
-            else TransactionType.DML
-        ), self.adapter.session():
+        with self.adapter.transaction(), self.adapter.session():
             if not limit:
                 self.adapter.execute(model.render_pre_statements(**render_statements_kwargs))
 
@@ -418,7 +414,7 @@ def _create_snapshot(
 
         evaluation_strategy = _evaluation_strategy(snapshot, self.adapter)
 
-        with self.adapter.transaction(TransactionType.DDL), self.adapter.session():
+        with self.adapter.transaction(), self.adapter.session():
             self.adapter.execute(snapshot.model.render_pre_statements(**render_kwargs))
 
             if is_dev and not snapshot.previous_versions: