diff --git a/databend_sqlalchemy/databend_dialect.py b/databend_sqlalchemy/databend_dialect.py index 9af3ef3..2d751d8 100644 --- a/databend_sqlalchemy/databend_dialect.py +++ b/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 diff --git a/tests/unit/test_databend_dialect.py b/tests/unit/test_databend_dialect.py index 122e037..a242cdc 100644 --- a/tests/unit/test_databend_dialect.py +++ b/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) ):