fix: rollback transaction on is_dialect exceptions (#341)

karenc-bq · web-flow · commit 3b33c1815ce6 · 2023-12-20T15:12:15.000-08:00
diff --git a/aws_advanced_python_wrapper/database_dialect.py b/aws_advanced_python_wrapper/database_dialect.py
@@ -130,7 +130,7 @@ def exception_handler(self) -> Optional[ExceptionHandler]:
         ...
 
     @abstractmethod
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
         ...
 
     @abstractmethod
@@ -184,7 +184,8 @@ def exception_handler(self) -> Optional[ExceptionHandler]:
     def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:
         return MysqlDatabaseDialect._DIALECT_UPDATE_CANDIDATES
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute(self.server_version_query)
@@ -193,7 +194,8 @@ def is_dialect(self, conn: Connection) -> bool:
                         if "mysql" in column_value.lower():
                             return True
         except Exception:
-            pass
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -232,16 +234,16 @@ def exception_handler(self) -> Optional[ExceptionHandler]:
                 "aws_advanced_python_wrapper.utils.pg_exception_handler.SingleAzPgExceptionHandler")
         return PgDatabaseDialect._exception_handler
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute('SELECT 1 FROM pg_proc LIMIT 1')
                 if cursor.fetchone() is not None:
                     return True
         except Exception:
-            # Executing the select statements will start a transaction, if the queries failed due to invalid syntax,
-            # the transaction will be aborted, no further commands can be executed. We need to call rollback here.
-            conn.rollback()
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -255,7 +257,8 @@ def prepare_conn_props(self, props: Properties):
 class RdsMysqlDialect(MysqlDatabaseDialect):
     _DIALECT_UPDATE_CANDIDATES = (DialectCode.AURORA_MYSQL, DialectCode.MULTI_AZ_MYSQL)
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute(self.server_version_query)
@@ -264,7 +267,8 @@ def is_dialect(self, conn: Connection) -> bool:
                         if "source distribution" in column_value.lower():
                             return True
         except Exception:
-            pass
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -280,8 +284,9 @@ class RdsPgDialect(PgDatabaseDialect):
                          "WHERE name='rds.extensions'")
     _DIALECT_UPDATE_CANDIDATES = (DialectCode.AURORA_PG, DialectCode.MULTI_AZ_PG)
 
-    def is_dialect(self, conn: Connection) -> bool:
-        if not super().is_dialect(conn):
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
+        if not super().is_dialect(conn, driver_dialect):
             return False
 
         try:
@@ -296,9 +301,8 @@ def is_dialect(self, conn: Connection) -> bool:
                         return True
 
         except Exception:
-            # Executing the select statements will start a transaction, if the queries failed due to invalid syntax,
-            # the transaction will be aborted, no further commands can be executed. We need to call rollback here.
-            conn.rollback()
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
         return False
 
     @property
@@ -320,15 +324,17 @@ class AuroraMysqlDialect(MysqlDatabaseDialect, TopologyAwareDatabaseDialect):
     def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:
         return AuroraMysqlDialect._DIALECT_UPDATE_CANDIDATES
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute("SHOW VARIABLES LIKE 'aurora_version'")
                 # If variable with such a name is presented then it means it's an Aurora cluster
                 if cursor.fetchone() is not None:
                     return True
         except Exception:
-            pass
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -358,13 +364,14 @@ class AuroraPgDialect(PgDatabaseDialect, TopologyAwareDatabaseDialect):
     def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:
         return AuroraPgDialect._DIALECT_UPDATE_CANDIDATES
 
-    def is_dialect(self, conn: Connection) -> bool:
-        if not super().is_dialect(conn):
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        if not super().is_dialect(conn, driver_dialect):
             return False
 
         has_extensions: bool = False
         has_topology: bool = False
 
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute(self._EXTENSIONS_QUERY)
@@ -381,9 +388,8 @@ def is_dialect(self, conn: Connection) -> bool:
 
             return has_extensions and has_topology
         except Exception:
-            # Executing the select statements will start a transaction, if the queries failed due to invalid syntax,
-            # the transaction will be aborted, no further commands can be executed. We need to call rollback here.
-            conn.rollback()
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -402,15 +408,17 @@ class MultiAzMysqlDialect(MysqlDatabaseDialect, TopologyAwareDatabaseDialect):
     def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:
         return None
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute(MultiAzMysqlDialect._TOPOLOGY_QUERY)
                 records = cursor.fetchall()
                 if records is not None and len(records) > 0:
                     return True
         except Exception:
-            pass
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -459,14 +467,16 @@ def exception_handler(self) -> Optional[ExceptionHandler]:
                 "aws_advanced_python_wrapper.utils.pg_exception_handler.MultiAzPgExceptionHandler")
         return MultiAzPgDialect._exception_handler
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
+        initial_transaction_status: bool = driver_dialect.is_in_transaction(conn)
         try:
             with closing(conn.cursor()) as cursor:
                 cursor.execute(MultiAzPgDialect._WRITER_HOST_QUERY)
                 if cursor.fetchone() is not None:
                     return True
         except Exception:
-            pass
+            if not initial_transaction_status and driver_dialect.is_in_transaction(conn):
+                conn.rollback()
 
         return False
 
@@ -511,7 +521,7 @@ def dialect_update_candidates(self) -> Optional[Tuple[DialectCode, ...]]:
     def exception_handler(self) -> Optional[ExceptionHandler]:
         return None
 
-    def is_dialect(self, conn: Connection) -> bool:
+    def is_dialect(self, conn: Connection, driver_dialect: DriverDialect) -> bool:
         return False
 
     def get_host_list_provider_supplier(self) -> Callable:
@@ -662,7 +672,7 @@ def query_for_dialect(self, url: str, host_info: Optional[HostInfo], conn: Conne
                         timeout_sec,
                         driver_dialect,
                         conn)(dialect_candidate.is_dialect)
-                    is_dialect = cursor_execute_func_with_timeout(conn)
+                    is_dialect = cursor_execute_func_with_timeout(conn, driver_dialect)
                 except TimeoutError as e:
                     raise QueryTimeoutError("DatabaseDialectManager.QueryForDialectTimeout") from e
 
diff --git a/aws_advanced_python_wrapper/pg_driver_dialect.py b/aws_advanced_python_wrapper/pg_driver_dialect.py
@@ -38,8 +38,7 @@ class PgDriverDialect(DriverDialect):
     TARGET_DRIVER_CODE: str = "psycopg"
 
     # https://www.psycopg.org/psycopg3/docs/api/pq.html#psycopg.pq.TransactionStatus
-    PSYCOPG_ACTIVE_TRANSACTION_STATUS = 1
-    PSYCOPG_IN_TRANSACTION_STATUS = 2
+    PSYCOPG_IDLE_TRANSACTION_STATUS = 0
 
     _dialect_code: str = DriverDialectCodes.PSYCOPG
     _network_bound_methods: Set[str] = {
@@ -78,7 +77,7 @@ def abort_connection(self, conn: Connection):
     def is_in_transaction(self, conn: Connection) -> bool:
         if isinstance(conn, psycopg.Connection):
             status: int = conn.info.transaction_status
-            return status == self.PSYCOPG_ACTIVE_TRANSACTION_STATUS or status == self.PSYCOPG_IN_TRANSACTION_STATUS
+            return status != self.PSYCOPG_IDLE_TRANSACTION_STATUS
 
         raise UnsupportedOperationError(Messages.get_formatted(
             "DriverDialect.UnsupportedOperationError",
diff --git a/tests/integration/container/test_basic_connectivity.py b/tests/integration/container/test_basic_connectivity.py
@@ -110,7 +110,6 @@ def test_proxied_wrapper_connection_failed(
             # That is expected exception. Test pass.
             assert True
 
-    @pytest.mark.skip("Currently failing, will be fixed in another PR")
     @enable_on_num_instances(min_instances=2)
     @enable_on_deployments([DatabaseEngineDeployment.AURORA, DatabaseEngineDeployment.MULTI_AZ])
     @enable_on_features([TestEnvironmentFeatures.ABORT_CONNECTION_SUPPORTED])
diff --git a/tests/unit/test_dialect.py b/tests/unit/test_dialect.py
@@ -99,96 +99,96 @@ def mock_default_behavior(mock_conn, mock_cursor, mock_fetchone_row):
     mock_cursor.fetchone.return_value = mock_fetchone_row
 
 
-def test_pg_is_dialect(mock_conn, mock_cursor, mock_session, pg_dialect):
-    assert pg_dialect.is_dialect(mock_conn)
+def test_pg_is_dialect(mock_conn, mock_cursor, mock_session, pg_dialect, mock_driver_dialect):
+    assert pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.fetchone.return_value = None
-    assert not pg_dialect.is_dialect(mock_conn)
+    assert not pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
 
-def test_mysql_is_dialect(mock_conn, mock_cursor, mock_session, mysql_dialect):
+def test_mysql_is_dialect(mock_conn, mock_cursor, mock_session, mysql_dialect, mock_driver_dialect):
     records = [("some_value", "some_value"), ("some_value", "mysql")]
     mock_cursor.__iter__.return_value = records
 
-    assert mysql_dialect.is_dialect(mock_conn)
+    assert mysql_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     records = [("some_value", "some_value"), ("some_value", "some_value")]
     mock_cursor.__iter__.return_value = records
 
-    assert not mysql_dialect.is_dialect(mock_conn)
+    assert not mysql_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
 
 @patch('aws_advanced_python_wrapper.database_dialect.super')
-def test_rds_mysql_is_dialect(mock_super, mock_cursor, mock_conn, rds_mysql_dialect):
+def test_rds_mysql_is_dialect(mock_super, mock_cursor, mock_conn, rds_mysql_dialect, mock_driver_dialect):
     mock_super().is_dialect.return_value = True
 
     records = [("some_value", "some_value"), ("some_value", "source distribution")]
     mock_cursor.__iter__.return_value = records
 
-    assert rds_mysql_dialect.is_dialect(mock_conn)
+    assert rds_mysql_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     records = [("some_value", "some_value"), ("some_value", "some_value")]
     mock_cursor.__iter__.return_value = records
 
-    assert not rds_mysql_dialect.is_dialect(mock_conn)
+    assert not rds_mysql_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_super().is_dialect.return_value = False
 
-    assert not rds_mysql_dialect.is_dialect(mock_conn)
+    assert not rds_mysql_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
 
-def test_aurora_mysql_is_dialect(mock_conn, mock_cursor):
+def test_aurora_mysql_is_dialect(mock_conn, mock_cursor, mock_driver_dialect):
     mock_conn.cursor.return_value = mock_cursor
     mock_cursor.fetchone.return_value = None
 
     dialect = AuroraMysqlDialect()
-    assert dialect.is_dialect(mock_conn) is False
+    assert dialect.is_dialect(mock_conn, mock_driver_dialect) is False
 
     mock_cursor.fetchone.return_value = ('aurora_version', '3.0.0')
-    assert dialect.is_dialect(mock_conn) is True
+    assert dialect.is_dialect(mock_conn, mock_driver_dialect) is True
 
 
 @patch('aws_advanced_python_wrapper.database_dialect.super')
-def test_aurora_pg_is_dialect(mock_super, mock_conn, mock_cursor):
+def test_aurora_pg_is_dialect(mock_super, mock_conn, mock_cursor, mock_driver_dialect):
     aurora_pg_dialect = AuroraPgDialect()
     mock_conn.cursor.return_value = mock_cursor
     mock_super().is_dialect.return_value = True
 
     records = [("aurora_stat_utils", "aurora_stat_utils"), ("some_value", "source distribution")]
     mock_cursor.__iter__.return_value = records
 
-    assert aurora_pg_dialect.is_dialect(mock_conn)
+    assert aurora_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.fetchone.return_value = None
 
-    assert not aurora_pg_dialect.is_dialect(mock_conn)
+    assert not aurora_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.fetchone.return_value = records
     mock_super().is_dialect.return_value = False
 
-    assert not aurora_pg_dialect.is_dialect(mock_conn)
+    assert not aurora_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
 
 @patch('aws_advanced_python_wrapper.database_dialect.super')
-def test_rds_pg_is_dialect(mock_super, mock_cursor, mock_conn, rds_pg_dialect):
+def test_rds_pg_is_dialect(mock_super, mock_cursor, mock_conn, rds_pg_dialect, mock_driver_dialect):
     mock_super().is_dialect.return_value = True
 
     mock_cursor.__iter__.return_value = [(False, False), (True, False)]
 
-    assert rds_pg_dialect.is_dialect(mock_conn)
+    assert rds_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.__iter__.return_value = [(False, False), (False, False)]
 
-    assert not rds_pg_dialect.is_dialect(mock_conn)
+    assert not rds_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.__iter__.return_value = []
 
-    assert not rds_pg_dialect.is_dialect(mock_conn)
+    assert not rds_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
     mock_cursor.fetchone.return_value = [(False, False), (True, False)]
     mock_super().is_dialect.return_value = False
 
-    assert not rds_pg_dialect.is_dialect(mock_conn)
+    assert not rds_pg_dialect.is_dialect(mock_conn, mock_driver_dialect)
 
 
 def test_get_dialect_custom_dialect(mock_custom_dialect, mock_driver_dialect):
