From 5536325b5ee2478bc51d33ca4bac7206f9fcb794 Mon Sep 17 00:00:00 2001 From: David Li Date: Mon, 6 Jan 2025 19:30:21 -0500 Subject: [PATCH] fix(c/driver/postgresql): don't unnecessarily COMMIT (#2412) Fixes #2410. --- .github/workflows/integration.yml | 29 ++++++------------- c/driver/postgresql/bind_stream.h | 12 ++++++-- .../tests/test_dbapi.py | 24 +++++++++++++++ 3 files changed, 43 insertions(+), 22 deletions(-) diff --git a/.github/workflows/integration.yml b/.github/workflows/integration.yml index 352f3cca30..e09a644136 100644 --- a/.github/workflows/integration.yml +++ b/.github/workflows/integration.yml @@ -230,62 +230,51 @@ jobs: echo "ADBC_USE_ASAN=ON" >> $GITHUB_ENV echo "ADBC_USE_UBSAN=ON" >> $GITHUB_ENV - - name: Test PostgreSQL Driver - postgres 11 - shell: bash -l {0} - env: - BUILD_ALL: "0" - BUILD_DRIVER_POSTGRESQL: "1" - PYTEST_ADDOPTS: "--error-for-skips" - run: | - env POSTGRES_VERSION=11 docker compose up --wait --detach postgres-test - ./ci/scripts/cpp_test.sh "$(pwd)/build" - ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" - docker compose down - - name: Test PostgreSQL Driver - postgres 12 + - name: Test PostgreSQL Driver - postgres 13 shell: bash -l {0} env: BUILD_ALL: "0" BUILD_DRIVER_POSTGRESQL: "1" PYTEST_ADDOPTS: "--error-for-skips" run: | - env POSTGRES_VERSION=12 docker compose up --wait --detach postgres-test + env POSTGRES_VERSION=13 docker compose up --wait --detach postgres-test ./ci/scripts/cpp_test.sh "$(pwd)/build" ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" docker compose down - - name: Test PostgreSQL Driver - postgres 13 + - name: Test PostgreSQL Driver - postgres 14 shell: bash -l {0} env: BUILD_ALL: "0" BUILD_DRIVER_POSTGRESQL: "1" PYTEST_ADDOPTS: "--error-for-skips" run: | - env POSTGRES_VERSION=13 docker compose up --wait --detach postgres-test + env POSTGRES_VERSION=14 docker compose up --wait --detach postgres-test ./ci/scripts/cpp_test.sh "$(pwd)/build" ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" docker compose down - - name: Test PostgreSQL Driver - postgres 14 + - name: Test PostgreSQL Driver - postgres 15 shell: bash -l {0} env: BUILD_ALL: "0" BUILD_DRIVER_POSTGRESQL: "1" PYTEST_ADDOPTS: "--error-for-skips" run: | - env POSTGRES_VERSION=14 docker compose up --wait --detach postgres-test + env POSTGRES_VERSION=15 docker compose up --wait --detach postgres-test ./ci/scripts/cpp_test.sh "$(pwd)/build" ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" docker compose down - - name: Test PostgreSQL Driver - postgres 15 + - name: Test PostgreSQL Driver - postgres 16 shell: bash -l {0} env: BUILD_ALL: "0" BUILD_DRIVER_POSTGRESQL: "1" PYTEST_ADDOPTS: "--error-for-skips" run: | - env POSTGRES_VERSION=15 docker compose up --wait --detach postgres-test + env POSTGRES_VERSION=16 docker compose up --wait --detach postgres-test ./ci/scripts/cpp_test.sh "$(pwd)/build" ./ci/scripts/python_test.sh "$(pwd)" "$(pwd)/build" docker compose down - - name: Test PostgreSQL Driver - postgres 16 + - name: Test PostgreSQL Driver - postgres 17 shell: bash -l {0} env: BUILD_ALL: "0" diff --git a/c/driver/postgresql/bind_stream.h b/c/driver/postgresql/bind_stream.h index df0b9d2ca5..5d1389a3c2 100644 --- a/c/driver/postgresql/bind_stream.h +++ b/c/driver/postgresql/bind_stream.h @@ -56,6 +56,7 @@ struct BindStream { Handle param_buffer; bool has_tz_field = false; + bool autocommit = false; std::string tz_setting; struct ArrowError na_error; @@ -119,6 +120,7 @@ struct BindStream { if (!has_tz_field && type.type_id() == PostgresTypeId::kTimestamptz) { UNWRAP_STATUS(SetDatabaseTimezoneUTC(pg_conn, autocommit)); has_tz_field = true; + this->autocommit = autocommit; } std::unique_ptr writer; @@ -254,8 +256,14 @@ struct BindStream { PqResultHelper reset(pg_conn, "SET TIME ZONE '" + tz_setting + "'"); UNWRAP_STATUS(reset.Execute()); - PqResultHelper commit(pg_conn, "COMMIT"); - UNWRAP_STATUS(reset.Execute()); + if (autocommit) { + // SetDatabaseTimezoneUTC will start a transaction if autocommit is + // enabled (so the timezone setting will roll back if we error), so we + // need to commit here. Otherwise we should not commit and let the + // user commit. + PqResultHelper commit(pg_conn, "COMMIT"); + UNWRAP_STATUS(commit.Execute()); + } } return Status::Ok(); diff --git a/python/adbc_driver_postgresql/tests/test_dbapi.py b/python/adbc_driver_postgresql/tests/test_dbapi.py index 9e2661d542..4efd833259 100644 --- a/python/adbc_driver_postgresql/tests/test_dbapi.py +++ b/python/adbc_driver_postgresql/tests/test_dbapi.py @@ -15,6 +15,7 @@ # specific language governing permissions and limitations # under the License. +import datetime import string from pathlib import Path from typing import Generator @@ -424,3 +425,26 @@ def test_ingest_large(postgres: dbapi.Connection) -> None: table = pyarrow.Table.from_batches([batch] * 4) with postgres.cursor() as cur: cur.adbc_ingest("test_ingest_large", table, mode="replace") + + +def test_timestamp_txn(postgres: dbapi.Connection) -> None: + """Regression test for #2410.""" + ts = datetime.datetime.now() + ts_with_tz = datetime.datetime.now(tz=datetime.UTC) + + with postgres.cursor() as cur: + cur.execute("DROP TABLE IF EXISTS ts_txn") + cur.execute("CREATE TABLE ts_txn (ts TIMESTAMP WITH TIME ZONE);") + postgres.commit() + + with postgres.cursor() as cur: + cur.execute("INSERT INTO ts_txn VALUES ($1)", parameters=[ts]) + cur.execute("SELECT pg_current_xact_id_if_assigned()") + assert cur.fetchone() != (None,) + postgres.commit() + + with postgres.cursor() as cur: + cur.execute("INSERT INTO ts_txn VALUES ($1)", parameters=[ts_with_tz]) + cur.execute("SELECT pg_current_xact_id_if_assigned()") + assert cur.fetchone() != (None,) + postgres.commit()