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