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 1da8aa0be feat(c/driver/postgresql): set use_copy on connect (#4750)
1da8aa0be is described below
commit 1da8aa0be18d7ba32423c4f88ebf4e42f712319f
Author: David Li <[email protected]>
AuthorDate: Thu Sep 3 10:19:57 2026 +0900
feat(c/driver/postgresql): set use_copy on connect (#4750)
This makes it easier to work with databases that are mostly
PostgreSQL-compatible but don't support COPY.
---
c/driver/postgresql/connection.cc | 20 ++++++++++
c/driver/postgresql/connection.h | 8 +++-
c/driver/postgresql/database.cc | 25 ++++++++++++-
c/driver/postgresql/database.h | 2 +
c/driver/postgresql/postgres_util.h | 2 +
c/driver/postgresql/postgresql_test.cc | 67 ++++++++++++++++++++++++++++++++++
c/driver/postgresql/statement.cc | 7 ++--
c/driver/postgresql/statement.h | 10 ++---
c/validation/adbc_validation_util.cc | 14 +++++++
c/validation/adbc_validation_util.h | 4 ++
10 files changed, 146 insertions(+), 13 deletions(-)
diff --git a/c/driver/postgresql/connection.cc
b/c/driver/postgresql/connection.cc
index f4a59aa71..3a6ecd928 100644
--- a/c/driver/postgresql/connection.cc
+++ b/c/driver/postgresql/connection.cc
@@ -45,6 +45,7 @@
#include "driver/framework/objects.h"
#include "driver/framework/utility.h"
#include "error.h"
+#include "postgres_util.h"
#include "result_helper.h"
using adbc::driver::Result;
@@ -668,6 +669,8 @@ AdbcStatusCode PostgresConnection::GetOption(const char*
option, char* value,
output = "unknown";
break;
}
+ } else if (std::strcmp(option, ADBC_POSTGRESQL_OPTION_USE_COPY) == 0) {
+ output = use_copy_ ? ADBC_OPTION_VALUE_ENABLED :
ADBC_OPTION_VALUE_DISABLED;
} else {
return ADBC_STATUS_NOT_FOUND;
}
@@ -1130,6 +1133,7 @@ AdbcStatusCode PostgresConnection::Init(struct
AdbcDatabase* database,
database_ =
*reinterpret_cast<std::shared_ptr<PostgresDatabase>*>(database->private_data);
type_resolver_ = database_->type_resolver();
+ use_copy_ = database_->use_copy();
RAISE_ADBC(database_->Connect(&conn_, error));
@@ -1265,6 +1269,22 @@ AdbcStatusCode PostgresConnection::SetOption(const char*
key, const char* value,
PqResultHelper result_helper{conn_, query};
RAISE_STATUS(error, result_helper.Execute());
return ADBC_STATUS_OK;
+ } else if (std::strcmp(key, ADBC_POSTGRESQL_OPTION_USE_COPY) == 0) {
+ if (!conn_) {
+ post_init_options_.emplace_back(key, value);
+ return ADBC_STATUS_OK;
+ }
+
+ if (std::strcmp(value, ADBC_OPTION_VALUE_ENABLED) == 0) {
+ use_copy_ = true;
+ } else if (std::strcmp(value, ADBC_OPTION_VALUE_DISABLED) == 0) {
+ use_copy_ = false;
+ } else {
+ InternalAdbcSetError(error, "%s%s%s%s", "[libpq] Invalid value for
option ", key,
+ ": ", value);
+ return ADBC_STATUS_INVALID_ARGUMENT;
+ }
+ return ADBC_STATUS_OK;
}
InternalAdbcSetError(error, "%s%s", "[libpq] Unknown option ", key);
return ADBC_STATUS_NOT_IMPLEMENTED;
diff --git a/c/driver/postgresql/connection.h b/c/driver/postgresql/connection.h
index ff9d2effe..6eeef7132 100644
--- a/c/driver/postgresql/connection.h
+++ b/c/driver/postgresql/connection.h
@@ -34,7 +34,11 @@ class PostgresDatabase;
class PostgresConnection {
public:
PostgresConnection()
- : database_(nullptr), conn_(nullptr), cancel_(nullptr),
autocommit_(true) {}
+ : database_(nullptr),
+ conn_(nullptr),
+ cancel_(nullptr),
+ autocommit_(true),
+ use_copy_(true) {}
AdbcStatusCode Cancel(struct AdbcError* error);
AdbcStatusCode Commit(struct AdbcError* error);
@@ -77,6 +81,7 @@ class PostgresConnection {
return type_resolver_;
}
bool autocommit() const { return autocommit_; }
+ bool use_copy() const { return use_copy_; }
std::string_view VendorName();
const std::array<int, 3>& VendorVersion();
@@ -90,6 +95,7 @@ class PostgresConnection {
PGconn* conn_;
PGcancel* cancel_;
bool autocommit_;
+ bool use_copy_;
std::vector<std::pair<std::string, std::string>> post_init_options_;
};
} // namespace adbcpq
diff --git a/c/driver/postgresql/database.cc b/c/driver/postgresql/database.cc
index efe1ffadd..a329b36bc 100644
--- a/c/driver/postgresql/database.cc
+++ b/c/driver/postgresql/database.cc
@@ -36,6 +36,7 @@
#include <nanoarrow/nanoarrow.h>
#include "driver/common/utils.h"
+#include "postgres_util.h"
#include "result_helper.h"
namespace adbcpq {
@@ -47,7 +48,18 @@ PostgresDatabase::~PostgresDatabase() = default;
AdbcStatusCode PostgresDatabase::GetOption(const char* option, char* value,
size_t* length, struct AdbcError*
error) {
- return ADBC_STATUS_NOT_FOUND;
+ std::string output;
+ if (std::strcmp(option, ADBC_POSTGRESQL_OPTION_USE_COPY) == 0) {
+ output = use_copy_ ? ADBC_OPTION_VALUE_ENABLED :
ADBC_OPTION_VALUE_DISABLED;
+ } else {
+ InternalAdbcSetError(error, "[libpq] unknown database option '%s'",
option);
+ return ADBC_STATUS_NOT_FOUND;
+ }
+ if (output.size() + 1 <= *length) {
+ std::memcpy(value, output.c_str(), output.size() + 1);
+ }
+ *length = output.size() + 1;
+ return ADBC_STATUS_OK;
}
AdbcStatusCode PostgresDatabase::GetOptionBytes(const char* option, uint8_t*
value,
size_t* length, struct
AdbcError* error) {
@@ -89,8 +101,17 @@ AdbcStatusCode PostgresDatabase::Release(struct AdbcError*
error) {
AdbcStatusCode PostgresDatabase::SetOption(const char* key, const char* value,
struct AdbcError* error) {
- if (strcmp(key, "uri") == 0) {
+ if (std::strcmp(key, "uri") == 0) {
uri_ = value;
+ } else if (strcmp(key, ADBC_POSTGRESQL_OPTION_USE_COPY) == 0) {
+ if (strcmp(value, ADBC_OPTION_VALUE_ENABLED) == 0) {
+ use_copy_ = true;
+ } else if (strcmp(value, ADBC_OPTION_VALUE_DISABLED) == 0) {
+ use_copy_ = false;
+ } else {
+ InternalAdbcSetError(error, "[libpq] Invalid value for option %s=%s",
key, value);
+ return ADBC_STATUS_INVALID_ARGUMENT;
+ }
} else {
InternalAdbcSetError(error, "%s%s", "[libpq] Unknown database option ",
key);
return ADBC_STATUS_NOT_IMPLEMENTED;
diff --git a/c/driver/postgresql/database.h b/c/driver/postgresql/database.h
index d9333d307..1bb421cdb 100644
--- a/c/driver/postgresql/database.h
+++ b/c/driver/postgresql/database.h
@@ -66,11 +66,13 @@ class PostgresDatabase {
Status RebuildTypeResolver(PGconn* conn);
std::string_view VendorName() { return "PostgreSQL"; }
const std::array<int, 3>& VendorVersion() { return postgres_server_version_;
}
+ bool use_copy() const { return use_copy_; }
private:
int32_t open_connections_;
std::string uri_;
std::shared_ptr<PostgresTypeResolver> type_resolver_;
std::array<int, 3> postgres_server_version_{};
+ bool use_copy_ = true;
};
} // namespace adbcpq
diff --git a/c/driver/postgresql/postgres_util.h
b/c/driver/postgresql/postgres_util.h
index fbd609848..2591750c1 100644
--- a/c/driver/postgresql/postgres_util.h
+++ b/c/driver/postgresql/postgres_util.h
@@ -37,6 +37,8 @@
namespace adbcpq {
+#define ADBC_POSTGRESQL_OPTION_USE_COPY "adbc.postgresql.use_copy"
+
#if defined(_WIN32) && defined(_MSC_VER)
static inline uint16_t SwapNetworkToHost(uint16_t x) { return ntohs(x); }
static inline uint16_t SwapHostToNetwork(uint16_t x) { return htons(x); }
diff --git a/c/driver/postgresql/postgresql_test.cc
b/c/driver/postgresql/postgresql_test.cc
index 5bb9d3061..d202a430f 100644
--- a/c/driver/postgresql/postgresql_test.cc
+++ b/c/driver/postgresql/postgresql_test.cc
@@ -229,6 +229,73 @@ class PostgresDatabaseTest : public ::testing::Test,
};
ADBCV_TEST_DATABASE(PostgresDatabaseTest)
+TEST_F(PostgresDatabaseTest, UseCopyOption) {
+ const char* key = "adbc.postgresql.use_copy";
+ std::optional<std::string> use_copy;
+
+ adbc_validation::Handle<struct AdbcDatabase> database;
+ adbc_validation::Handle<struct AdbcConnection> connection;
+ adbc_validation::Handle<struct AdbcStatement> statement;
+
+ ASSERT_THAT(AdbcDatabaseNew(&database.value, &error), IsOkStatus(&error));
+ ASSERT_THAT(quirks_.SetupDatabase(&database.value, &error),
IsOkStatus(&error));
+ ASSERT_THAT(AdbcDatabaseSetOption(&database.value, key, "false", &error),
+ IsOkStatus(&error));
+ ASSERT_THAT(AdbcDatabaseInit(&database.value, &error), IsOkStatus(&error));
+ use_copy = adbc_validation::DatabaseGetOption(&database.value, key, &error);
+ EXPECT_THAT(use_copy, ::testing::Optional("false"s));
+
+ ASSERT_THAT(AdbcConnectionNew(&connection.value, &error),
IsOkStatus(&error));
+ ASSERT_THAT(AdbcConnectionInit(&connection.value, &database.value, &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::ConnectionGetOption(&connection.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("false"s));
+
+ ASSERT_THAT(AdbcStatementNew(&connection.value, &statement.value, &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::StatementGetOption(&statement.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("false"s));
+ ASSERT_THAT(AdbcStatementRelease(&statement.value, &error),
IsOkStatus(&error));
+
+ ASSERT_THAT(AdbcConnectionSetOption(&connection.value, key, "true", &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::ConnectionGetOption(&connection.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("true"s));
+
+ ASSERT_THAT(AdbcStatementNew(&connection.value, &statement.value, &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::StatementGetOption(&statement.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("true"s));
+ ASSERT_THAT(AdbcStatementRelease(&statement.value, &error),
IsOkStatus(&error));
+}
+
+TEST_F(PostgresDatabaseTest, UseCopyOptionDefault) {
+ const char* key = "adbc.postgresql.use_copy";
+ std::optional<std::string> use_copy;
+
+ adbc_validation::Handle<struct AdbcDatabase> database;
+ adbc_validation::Handle<struct AdbcConnection> connection;
+ adbc_validation::Handle<struct AdbcStatement> statement;
+
+ ASSERT_THAT(AdbcDatabaseNew(&database.value, &error), IsOkStatus(&error));
+ ASSERT_THAT(quirks_.SetupDatabase(&database.value, &error),
IsOkStatus(&error));
+ ASSERT_THAT(AdbcDatabaseInit(&database.value, &error), IsOkStatus(&error));
+ use_copy = adbc_validation::DatabaseGetOption(&database.value, key, &error);
+ EXPECT_THAT(use_copy, ::testing::Optional("true"s));
+
+ ASSERT_THAT(AdbcConnectionNew(&connection.value, &error),
IsOkStatus(&error));
+ ASSERT_THAT(AdbcConnectionInit(&connection.value, &database.value, &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::ConnectionGetOption(&connection.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("true"s));
+
+ ASSERT_THAT(AdbcStatementNew(&connection.value, &statement.value, &error),
+ IsOkStatus(&error));
+ use_copy = adbc_validation::StatementGetOption(&statement.value, key,
&error);
+ EXPECT_THAT(use_copy, ::testing::Optional("true"s));
+ ASSERT_THAT(AdbcStatementRelease(&statement.value, &error),
IsOkStatus(&error));
+}
+
int Canary(const struct AdbcError*) { return 0; }
TEST_F(PostgresDatabaseTest, AdbcDriverBackwardsCompatibility) {
diff --git a/c/driver/postgresql/statement.cc b/c/driver/postgresql/statement.cc
index 91e39e93f..8e215ade5 100644
--- a/c/driver/postgresql/statement.cc
+++ b/c/driver/postgresql/statement.cc
@@ -331,6 +331,7 @@ AdbcStatusCode PostgresStatement::New(struct
AdbcConnection* connection,
connection_ =
*reinterpret_cast<std::shared_ptr<PostgresConnection>*>(connection->private_data);
type_resolver_ = connection_->type_resolver();
+ use_copy_ = connection_->use_copy();
ClearResult();
return ADBC_STATUS_OK;
}
@@ -607,7 +608,7 @@ AdbcStatusCode PostgresStatement::ExecuteQuery(struct
ArrowArrayStream* stream,
// If we have been requested to avoid COPY or there is no output requested,
// execute using the PqResultArrayReader.
- if (!stream || !UseCopy()) {
+ if (!stream || !use_copy()) {
PqResultArrayReader reader(connection_->conn(), type_resolver_, query_);
reader.SetVendorName(connection_->VendorName());
RAISE_STATUS(error, reader.ToArrayStream(rows_affected, stream));
@@ -804,7 +805,7 @@ AdbcStatusCode PostgresStatement::GetOption(const char*
key, char* value, size_t
} else if (std::strcmp(key, ADBC_POSTGRESQL_OPTION_BATCH_SIZE_HINT_BYTES) ==
0) {
result = std::to_string(reader_->batch_size_hint_bytes_);
} else if (std::strcmp(key, ADBC_POSTGRESQL_OPTION_USE_COPY) == 0) {
- if (UseCopy()) {
+ if (use_copy()) {
result = "true";
} else {
result = "false";
@@ -1011,6 +1012,4 @@ void PostgresStatement::ClearResult() {
reader_->batch_size_hint_bytes_ = batch_size_hint_bytes_;
}
-int PostgresStatement::UseCopy() { return use_copy_; }
-
} // namespace adbcpq
diff --git a/c/driver/postgresql/statement.h b/c/driver/postgresql/statement.h
index fe77443c2..0ab4456c8 100644
--- a/c/driver/postgresql/statement.h
+++ b/c/driver/postgresql/statement.h
@@ -33,8 +33,6 @@
#define ADBC_POSTGRESQL_OPTION_BATCH_SIZE_HINT_BYTES \
"adbc.postgresql.batch_size_hint_bytes"
-#define ADBC_POSTGRESQL_OPTION_USE_COPY "adbc.postgresql.use_copy"
-
// This is not a public-facing PostgreSQL driver option.
#define ADBC_POSTGRESQL_OPTION_DISABLE_DECIMAL_FAST_PATH \
"adbc.postgresql.disable_decimal_fast_path"
@@ -103,7 +101,7 @@ class PostgresStatement {
: connection_(nullptr),
query_(),
prepared_(false),
- use_copy_(-1),
+ use_copy_(true),
disable_decimal_fast_path_(false),
reader_(nullptr),
batch_size_hint_bytes_(kDefaultBatchSizeHintBytes) {
@@ -156,6 +154,8 @@ class PostgresStatement {
AdbcStatusCode ExecuteBind(struct ArrowArrayStream* stream, int64_t*
rows_affected,
struct AdbcError* error);
+ bool use_copy() const { return use_copy_; }
+
private:
std::shared_ptr<PostgresTypeResolver> type_resolver_;
std::shared_ptr<PostgresConnection> connection_;
@@ -174,7 +174,7 @@ class PostgresStatement {
};
// Options
- int use_copy_;
+ bool use_copy_;
bool disable_decimal_fast_path_;
struct {
@@ -186,7 +186,5 @@ class PostgresStatement {
std::shared_ptr<TupleReader> reader_;
int64_t batch_size_hint_bytes_;
-
- int UseCopy();
};
} // namespace adbcpq
diff --git a/c/validation/adbc_validation_util.cc
b/c/validation/adbc_validation_util.cc
index 9fcf2c598..04f14420c 100644
--- a/c/validation/adbc_validation_util.cc
+++ b/c/validation/adbc_validation_util.cc
@@ -27,6 +27,20 @@
namespace adbc_validation {
+std::optional<std::string> DatabaseGetOption(struct AdbcDatabase* database,
+ std::string_view option,
+ struct AdbcError* error) {
+ char buffer[128];
+ size_t buffer_size = sizeof(buffer);
+ AdbcStatusCode status =
+ AdbcDatabaseGetOption(database, option.data(), buffer, &buffer_size,
error);
+ EXPECT_THAT(status, IsOkStatus(error));
+ if (status != ADBC_STATUS_OK) return std::nullopt;
+ EXPECT_GT(buffer_size, 0);
+ if (buffer_size == 0) return std::nullopt;
+ return std::string(buffer, buffer_size - 1);
+}
+
std::optional<std::string> ConnectionGetOption(struct AdbcConnection*
connection,
std::string_view option,
struct AdbcError* error) {
diff --git a/c/validation/adbc_validation_util.h
b/c/validation/adbc_validation_util.h
index 980912f71..cb9c8892e 100644
--- a/c/validation/adbc_validation_util.h
+++ b/c/validation/adbc_validation_util.h
@@ -39,6 +39,10 @@ namespace adbc_validation {
// ------------------------------------------------------------
// ADBC helpers
+std::optional<std::string> DatabaseGetOption(struct AdbcDatabase* database,
+ std::string_view option,
+ struct AdbcError* error);
+
std::optional<std::string> ConnectionGetOption(struct AdbcConnection*
connection,
std::string_view option,
struct AdbcError* error);
