(arrow-adbc) branch main updated: fix(c/driver/postgresql): handle connection options before Init (#2701)

[lidavidm]

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 93baecb70 fix(c/driver/postgresql): handle connection options before 
Init (#2701)
93baecb70 is described below

commit 93baecb70ad819a627955479b1d1dc3be8e30fe2
Author: David Li <li.david...@gmail.com>
AuthorDate: Tue Apr 15 11:48:12 2025 +0900

    fix(c/driver/postgresql): handle connection options before Init (#2701)
---
 c/driver/postgresql/connection.cc      | 22 +++++++++++++++++++-
 c/driver/postgresql/connection.h       |  4 ++++
 c/driver/postgresql/postgresql_test.cc | 38 ++++++++++++++++++++++++++++++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/c/driver/postgresql/connection.cc 
b/c/driver/postgresql/connection.cc
index 9e6c41c6f..d3b691138 100644
--- a/c/driver/postgresql/connection.cc
+++ b/c/driver/postgresql/connection.cc
@@ -32,6 +32,7 @@
 #include <vector>
 
 #include <arrow-adbc/adbc.h>
+#include <fmt/format.h>
 #include <libpq-fe.h>
 
 #include "database.h"
@@ -1108,6 +1109,11 @@ AdbcStatusCode PostgresConnection::Init(struct 
AdbcDatabase* database,
 
   std::ignore = PQsetNoticeProcessor(conn_, SilentNoticeProcessor, nullptr);
 
+  for (const auto& [key, value] : post_init_options_) {
+    RAISE_ADBC(SetOption(key.data(), value.data(), error));
+  }
+  post_init_options_.clear();
+
   return ADBC_STATUS_OK;
 }
 
@@ -1159,6 +1165,11 @@ AdbcStatusCode PostgresConnection::SetOption(const char* 
key, const char* value,
       return ADBC_STATUS_INVALID_ARGUMENT;
     }
 
+    if (!conn_) {
+      post_init_options_.emplace_back(key, value);
+      return ADBC_STATUS_OK;
+    }
+
     if (autocommit != autocommit_) {
       const char* query = autocommit ? "COMMIT" : "BEGIN TRANSACTION";
 
@@ -1174,9 +1185,18 @@ AdbcStatusCode PostgresConnection::SetOption(const char* 
key, const char* value,
     }
     return ADBC_STATUS_OK;
   } else if (std::strcmp(key, ADBC_CONNECTION_OPTION_CURRENT_DB_SCHEMA) == 0) {
+    if (!conn_) {
+      post_init_options_.emplace_back(key, value);
+      return ADBC_STATUS_OK;
+    }
+
     // PostgreSQL doesn't accept a parameter here
     char* value_esc = PQescapeIdentifier(conn_, value, strlen(value));
-    std::string query = std::string("SET search_path TO ") + value_esc;
+    if (!value_esc) {
+      SetError(error, "[libpq] Could not escape identifier: %s", 
PQerrorMessage(conn_));
+      return ADBC_STATUS_INTERNAL;
+    }
+    std::string query = fmt::format("SET search_path TO {}", value_esc);
     PQfreemem(value_esc);
 
     PqResultHelper result_helper{conn_, query};
diff --git a/c/driver/postgresql/connection.h b/c/driver/postgresql/connection.h
index 7683875b5..02e0c4f1b 100644
--- a/c/driver/postgresql/connection.h
+++ b/c/driver/postgresql/connection.h
@@ -20,6 +20,9 @@
 #include <array>
 #include <cstdint>
 #include <memory>
+#include <string>
+#include <utility>
+#include <vector>
 
 #include <arrow-adbc/adbc.h>
 #include <libpq-fe.h>
@@ -83,5 +86,6 @@ class PostgresConnection {
   PGconn* conn_;
   PGcancel* cancel_;
   bool autocommit_;
+  std::vector<std::pair<std::string, std::string>> post_init_options_;
 };
 }  // namespace adbcpq
diff --git a/c/driver/postgresql/postgresql_test.cc 
b/c/driver/postgresql/postgresql_test.cc
index b45fc5183..5167dbef7 100644
--- a/c/driver/postgresql/postgresql_test.cc
+++ b/c/driver/postgresql/postgresql_test.cc
@@ -691,6 +691,44 @@ TEST_F(PostgresConnectionTest, MetadataSetCurrentDbSchema) 
{
   ASSERT_THAT(AdbcStatementRelease(&statement.value, &error), 
IsOkStatus(&error));
 }
 
+TEST_F(PostgresConnectionTest, MetadataSetCurrentDbSchemaInit) {
+  // Regression test: setting the schema before Init (which Python does)
+
+  // 1. Create the schema
+  {
+    ASSERT_THAT(AdbcConnectionNew(&connection, &error), IsOkStatus(&error));
+    ASSERT_THAT(AdbcConnectionInit(&connection, &database, &error), 
IsOkStatus(&error));
+
+    adbc_validation::Handle<struct AdbcStatement> statement;
+    ASSERT_THAT(AdbcStatementNew(&connection, &statement.value, &error),
+                IsOkStatus(&error));
+
+    ASSERT_THAT(
+        AdbcStatementSetSqlQuery(&statement.value,
+                                 "CREATE SCHEMA IF NOT EXISTS regtestschema", 
&error),
+        IsOkStatus(&error));
+    ASSERT_THAT(AdbcStatementExecuteQuery(&statement.value, nullptr, nullptr, 
&error),
+                IsOkStatus(&error));
+
+    ASSERT_THAT(AdbcStatementRelease(&statement.value, &error), 
IsOkStatus(&error));
+    ASSERT_THAT(AdbcConnectionRelease(&connection, &error), 
IsOkStatus(&error));
+  }
+
+  // 2. Initialize a connection with the schema
+  {
+    ASSERT_THAT(AdbcConnectionNew(&connection, &error), IsOkStatus(&error));
+    ASSERT_THAT(
+        AdbcConnectionSetOption(&connection, 
ADBC_CONNECTION_OPTION_CURRENT_DB_SCHEMA,
+                                "regtestschema", &error),
+        IsOkStatus(&error));
+    ASSERT_THAT(AdbcConnectionInit(&connection, &database, &error), 
IsOkStatus(&error));
+
+    ASSERT_THAT(adbc_validation::ConnectionGetOption(
+                    &connection, ADBC_CONNECTION_OPTION_CURRENT_DB_SCHEMA, 
&error),
+                ::testing::Optional("regtestschema"s));
+  }
+}
+
 TEST_F(PostgresConnectionTest, MetadataGetSchemaCaseSensitiveTable) {
   ASSERT_THAT(AdbcConnectionNew(&connection, &error), IsOkStatus(&error));
   ASSERT_THAT(AdbcConnectionInit(&connection, &database, &error), 
IsOkStatus(&error));