fix(c/driver/postgresql): drain all results post-query

c/driver/postgresql/connection.cc

@@ -503,7 +503,7 @@ AdbcStatusCode PostgresConnection::EnsureTransaction(struct AdbcError* error) {
     return ADBC_STATUS_OK;
   }
   auto txstatus = PQtransactionStatus(conn_);
-  if (txstatus == PQTRANS_ACTIVE || txstatus == PQTRANS_INTRANS) {
+  if (txstatus == PQTRANS_INTRANS) {
     return ADBC_STATUS_OK;
   } else if (txstatus == PQTRANS_INERROR) {
     InternalAdbcSetError(error,

c/driver/postgresql/statement.cc

@@ -96,14 +96,20 @@ int TupleReader::GetCopyData() {
     PQclear(result_);
     result_ = PQgetResult(conn_);
     const ExecStatusType pq_status = PQresultStatus(result_);
+    int errno_result = ENODATA;
     if (pq_status != PGRES_COMMAND_OK) {
       status_ = MakeStatus(result_, "[libpq] Execution error [{}]: {}",
                            PQresStatus(pq_status), PQresultErrorMessage(result_))
                     .ToAdbc(&error_);
-      return InternalAdbcStatusCodeToErrno(status_);
-    } else {
-      return ENODATA;
+      errno_result = InternalAdbcStatusCodeToErrno(status_);
+    }
+
+    // Drain remaining responses
+    PQclear(result_);
+    while ((result_ = PQgetResult(conn_)) != nullptr) {
+      PQclear(result_);
     }
+    return errno_result;
   }
 
   data_.size_bytes = get_copy_res;

python/adbc_driver_postgresql/tests/test_dbapi.py

@@ -575,11 +575,13 @@ def status() -> str:
 
     assert status() == "idle"
     postgres.rollback()
-    assert status() == "idle"
+    assert status() == "idle"  # because txn is lazily started
 
     with postgres.cursor() as cur:
         cur.execute("SELECT 1")
-        assert status() == "active"
+        assert status() == "active"  # because result is unread
+        cur.fetchall()
+        assert status() == "intrans"
         postgres.commit()
         assert status() == "idle"
         cur.execute("SELECT 1")

python/adbc_driver_postgresql/tests/test_lowlevel.py

@@ -55,3 +55,43 @@ def test_failed_connection() -> None:
         adbc_driver_manager.OperationalError, match=".*libpq.*Failed to connect.*"
     ):
         adbc_driver_postgresql.connect("invalid")
+
+
+@pytest.mark.parametrize("drain", [False, True])
+def test_transaction(postgres_uri: str, drain: bool) -> None:
+    # regression test for https://github.com/apache/arrow-adbc/issues/4695
+    status = adbc_driver_postgresql.ConnectionOptions.TRANSACTION_STATUS.value
+    with adbc_driver_postgresql.connect(postgres_uri) as db:
+        with adbc_driver_manager.AdbcConnection(db) as conn:
+            with adbc_driver_manager.AdbcStatement(conn) as stmt:
+                stmt.set_sql_query("DROP TABLE IF EXISTS test_transaction")
+                stmt.execute_update()
+                stmt.set_sql_query("CREATE TABLE test_transaction (id INT)")
+                stmt.execute_update()
+
+        with adbc_driver_manager.AdbcConnection(db) as conn:
+            with adbc_driver_manager.AdbcStatement(conn) as stmt:
+                stmt.set_sql_query("SELECT COUNT(*) FROM test_transaction")
+                handle, _ = stmt.execute_query()
+                with pyarrow.RecordBatchReader._import_from_c(handle.address) as reader:
+                    if drain:
+                        result = reader.read_all()
+                        assert result[0][0].as_py() == 0
+            assert conn.get_option(status) == ("idle" if drain else "active")
+
+            conn.set_autocommit(False)
+            assert conn.get_option(status) == ("idle" if drain else "active")
+            with adbc_driver_manager.AdbcStatement(conn) as stmt:
+                stmt.set_sql_query("INSERT INTO test_transaction (id) VALUES (1)")
+                stmt.execute_update()
+            assert conn.get_option(status) == "intrans"
+            conn.rollback()
+            assert conn.get_option(status) == "idle"
+
+        with adbc_driver_manager.AdbcConnection(db) as conn:
+            with adbc_driver_manager.AdbcStatement(conn) as stmt:
+                stmt.set_sql_query("SELECT COUNT(*) FROM test_transaction")
+                handle, _ = stmt.execute_query()
+                with pyarrow.RecordBatchReader._import_from_c(handle.address) as reader:
+                    result = reader.read_all()
+                assert result[0][0].as_py() == 0