From ed057b7f99482b7164721abc1860bc0bf259b5d9 Mon Sep 17 00:00:00 2001
From: David Li
Date: Fri, 21 Aug 2026 12:07:34 +0900
Subject: [PATCH] fix(c/driver/postgresql): drain all results post-query
Closes #4695.
---
c/driver/postgresql/connection.cc | 2 +-
c/driver/postgresql/statement.cc | 12 ++++--
.../tests/test_dbapi.py | 6 ++-
.../tests/test_lowlevel.py | 40 +++++++++++++++++++
4 files changed, 54 insertions(+), 6 deletions(-)
diff --git a/c/driver/postgresql/connection.cc b/c/driver/postgresql/connection.cc
index 262c047d74..f4a59aa71e 100644
--- a/c/driver/postgresql/connection.cc
+++ b/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,
diff --git a/c/driver/postgresql/statement.cc b/c/driver/postgresql/statement.cc
index 37e61979c4..91e39e93f3 100644
--- a/c/driver/postgresql/statement.cc
+++ b/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;
diff --git a/python/adbc_driver_postgresql/tests/test_dbapi.py b/python/adbc_driver_postgresql/tests/test_dbapi.py
index 4f6456047a..8109243d72 100644
--- a/python/adbc_driver_postgresql/tests/test_dbapi.py
+++ b/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")
diff --git a/python/adbc_driver_postgresql/tests/test_lowlevel.py b/python/adbc_driver_postgresql/tests/test_lowlevel.py
index 98a355bb12..92e4659dbb 100644
--- a/python/adbc_driver_postgresql/tests/test_lowlevel.py
+++ b/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