fix(c/driver/postgresql): don't unnecessarily COMMIT (#2412)

Fixes #2410.
apache · Jan 7, 2025 · 5536325 · 5536325
1 parent ef03777
commit 5536325
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 22 deletions.
diff --git 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
@@ -56,6 +56,7 @@ struct BindStream {
   Handle<struct ArrowBuffer> 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<PostgresCopyFieldWriter> 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
@@ -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()