fix(dialect): delegate do_rollback to dbapi connection (#77)

databend_sqlalchemy/databend_dialect.py

@@ -27,6 +27,7 @@
 import re
 import operator
 import datetime
+import logging
 from types import NoneType
 
 import sqlalchemy.engine.reflection
@@ -43,6 +44,9 @@
     Select,
     Subquery,
 )
+from databend_sqlalchemy.errors import NotSupportedError
+
+logger = logging.getLogger(__name__)
 from sqlalchemy.dialects.postgresql.base import PGCompiler, PGIdentifierPreparer
 from sqlalchemy import Table, MetaData, Column
 from sqlalchemy.types import (
@@ -1817,8 +1821,18 @@ def get_multi_table_comment(
         )
 
     def do_rollback(self, dbapi_connection):
-        # No transactions
-        pass
+        dbapi = getattr(self, "dbapi", None)
+        not_supported = (
+            getattr(dbapi, "NotSupportedError", NotSupportedError),
+            NotSupportedError,
+            AttributeError,
+        )
+        try:
+            dbapi_connection.rollback()
+        except not_supported as e:
+            logger.debug(
+                "Database rollback not supported by underlying connection: %s", e
+            )
 
     def _check_unicode_returns(self, connection, additional_tests=None):
         # We decode everything as UTF-8

tests/unit/test_databend_dialect.py

@@ -3,12 +3,14 @@
 from unittest import mock
 
 import sqlalchemy
+import pytest
 from conftest import MockCursor, MockDBApi
 from pytest import mark
 from sqlalchemy.engine import url
 from sqlalchemy.sql import text
 
 import databend_sqlalchemy  # SQLAlchemy package
+from databend_sqlalchemy.errors import NotSupportedError
 from databend_sqlalchemy.databend_dialect import (
     DatabendCompiler,
     DatabendDialect,
@@ -56,6 +58,29 @@ def test_do_execute(
         cursor.execute.reset_mock()
         dialect.do_execute(cursor, "SELECT *", (1, 22), None)
 
+    def test_do_rollback_propagates_rollback(self, dialect: DatabendDialect):
+        mock_conn = mock.Mock()
+        dialect.do_rollback(mock_conn)
+        mock_conn.rollback.assert_called_once()
+
+    def test_do_rollback_safely_handles_not_supported_error(
+        self, dialect: DatabendDialect
+    ):
+        mock_conn = mock.Mock()
+        mock_conn.rollback.side_effect = NotSupportedError(
+            "Transactions are not supported"
+        )
+        # Must not raise to ensure pool checkout/checkin lifecycle succeeds
+        dialect.do_rollback(mock_conn)
+        mock_conn.rollback.assert_called_once()
+
+    def test_do_rollback_safely_handles_attribute_error(
+        self, dialect: DatabendDialect
+    ):
+        mock_conn = object()
+        # Must not raise if connection object lacks rollback method
+        dialect.do_rollback(mock_conn)
+
     def test_table_names(
             self, dialect: DatabendDialect, connection: mock.Mock(spec=MockDBApi)
     ):