This is an automated email from the ASF dual-hosted git repository.
lidavidm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/arrow-adbc.git
The following commit(s) were added to refs/heads/main by this push:
new 7616599b0 fix(c/driver/postgresql): drain all results post-query
(#4710)
7616599b0 is described below
commit 7616599b07fabf679854cc956b86ddc4fd87bf48
Author: David Li <[email protected]>
AuthorDate: Mon Aug 24 12:10:02 2026 +0900
fix(c/driver/postgresql): drain all results post-query (#4710)
Closes #4695.
---
c/driver/postgresql/connection.cc | 2 +-
c/driver/postgresql/statement.cc | 12 +++++--
python/adbc_driver_postgresql/tests/test_dbapi.py | 6 ++--
.../adbc_driver_postgresql/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 262c047d7..f4a59aa71 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 37e61979c..91e39e93f 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 4f6456047..8109243d7 100644
--- a/python/adbc_driver_postgresql/tests/test_dbapi.py
+++ b/python/adbc_driver_postgresql/tests/test_dbapi.py
@@ -575,11 +575,13 @@ def test_txn_status(postgres: dbapi.Connection) -> None:
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 98a355bb1..92e4659db 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")
+
+
[email protected]("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
