From d627125e9b3d01d512b8d66f6ab9e7000f6a0bb6 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Fri, 14 Aug 2026 22:19:43 +0000 Subject: [PATCH 1/2] impl: PartnerToken configuration property --- .../bq_client_interface/odbc_authentication.h | 1 + .../bq_client_interface/odbc_bq_client.cc | 16 ++++- google/cloud/odbc/bq_client_interface/utils.h | 31 ++++++++- .../odbc/bq_client_interface/utils_test.cc | 63 +++++++++++++++++++ .../bq_driver/internal/odbc_conn_handle.cc | 1 + .../bq_driver/internal/odbc_conn_handle.h | 1 + .../cloud/odbc/bq_driver/odbc_connection.cc | 1 + 7 files changed, 110 insertions(+), 4 deletions(-) create mode 100644 google/cloud/odbc/bq_client_interface/utils_test.cc diff --git a/google/cloud/odbc/bq_client_interface/odbc_authentication.h b/google/cloud/odbc/bq_client_interface/odbc_authentication.h index b3e14c2cec..e899434e8e 100644 --- a/google/cloud/odbc/bq_client_interface/odbc_authentication.h +++ b/google/cloud/odbc/bq_client_interface/odbc_authentication.h @@ -91,6 +91,7 @@ struct ConnProps { std::string psc; GCD gcd; std::string impersonated_email; + std::string partner_token; }; // Returns true if all required BYOID properties are set. diff --git a/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc b/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc index 473de66126..2c26d0f857 100644 --- a/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc +++ b/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc @@ -19,6 +19,7 @@ #include "google/cloud/odbc/bq_client_interface/projects.h" #include "google/cloud/odbc/bq_client_interface/storage.h" #include "google/cloud/odbc/bq_client_interface/tables.h" +#include "google/cloud/odbc/bq_client_interface/utils.h" #include "google/cloud/odbc/internal/status_record_or.h" #include "google/cloud/odbc/internal/version.h" #include "google/cloud/common_options.h" @@ -192,8 +193,17 @@ StatusRecordOr> ODBCBQClient::CreateBQClient( .set_password(conn_props.proxy_options.password) .set_scheme("http")); + auto partner_token_or = + ::google::cloud::odbc_bigquery_client_interface::ParsePartnerToken( + conn_props.partner_token); + if (!partner_token_or) { + LOG(ERROR) << "CreateBQClient::ParsePartnerToken:: " + << partner_token_or.GetStatusRecord().message; + return partner_token_or.GetStatusRecord(); + } + std::string partner_token = *partner_token_or; options.set( - {"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION)}); + {"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token}); if (conn_props.gcd.enable_gcd && conn_props.gcd.universe_domain != "googleapis.com") { @@ -282,8 +292,8 @@ StatusRecordOr> ODBCBQClient::CreateBQClient( // account impersonation) will hang waiting on completion events. grpc::ChannelArguments channel_arguments; - channel_arguments.SetUserAgentPrefix("Google-Bigquery-ODBC/" + - std::string(DRIVER_VERSION)); + channel_arguments.SetUserAgentPrefix( + "Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token); channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, std::chrono::minutes(1).count()); channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, diff --git a/google/cloud/odbc/bq_client_interface/utils.h b/google/cloud/odbc/bq_client_interface/utils.h index a7d30483e5..def7d9ef15 100644 --- a/google/cloud/odbc/bq_client_interface/utils.h +++ b/google/cloud/odbc/bq_client_interface/utils.h @@ -15,10 +15,14 @@ #ifndef CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H #define CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H +#include "google/cloud/odbc/internal/sql_state_constants.h" +#include "google/cloud/odbc/internal/status_record_or.h" +#include "google/cloud/internal/backoff_policy.h" #include "absl/log/log.h" #include #include #include +#include #include #include #include @@ -56,7 +60,7 @@ auto RetryLoop(Functor&& functor, std::string const& operation_name, using ReturnType = decltype(functor()); ReturnType response; - google::cloud::ExponentialBackoffPolicy backoff_policy( + google::cloud::internal::ExponentialBackoffPolicy backoff_policy( std::chrono::milliseconds(initial_delay_ms), std::chrono::milliseconds(max_delay_ms), backoff_multiplier); @@ -95,5 +99,30 @@ auto RetryLoop(Functor&& functor, std::string const& operation_name, return response; } + +inline google::cloud::odbc_internal::StatusRecordOr +ParsePartnerToken(std::string const& raw_token) { + if (raw_token.empty()) { + return std::string(""); + } + std::regex pattern(R"(\(\s*(GPN:[^;]*?)\s*(?:;\s*([^)]*?))?\s*\))"); + std::smatch match; + if (std::regex_search(raw_token, match, pattern)) { + std::string gpn_part = match[1].str(); + std::string env_part = match[2].str(); + std::string partner_token = " ("; + partner_token += gpn_part; + if (!env_part.empty()) { + partner_token += "; "; + partner_token += env_part; + } + partner_token += ")"; + return partner_token; + } + return google::cloud::odbc_internal::StatusRecord{ + google::cloud::odbc_internal::SQLStates::k_HY024(), + "Invalid PartnerToken format."}; +} + } // namespace google::cloud::odbc_bigquery_client_interface #endif // CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H diff --git a/google/cloud/odbc/bq_client_interface/utils_test.cc b/google/cloud/odbc/bq_client_interface/utils_test.cc new file mode 100644 index 0000000000..b3fd6f1960 --- /dev/null +++ b/google/cloud/odbc/bq_client_interface/utils_test.cc @@ -0,0 +1,63 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// https://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#include "google/cloud/odbc/bq_client_interface/utils.h" +#include +#include + +namespace google::cloud::odbc_bigquery_client_interface { +namespace { + +TEST(ParsePartnerToken, ValidFull) { + std::string raw = "(GPN:PartnerName; Environment)"; + auto result = ParsePartnerToken(raw); + EXPECT_TRUE(result.Ok()); + EXPECT_EQ(*result, " (GPN:PartnerName; Environment)"); +} + +TEST(ParsePartnerToken, ValidShort) { + std::string raw = "(GPN:PartnerName)"; + auto result = ParsePartnerToken(raw); + EXPECT_TRUE(result.Ok()); + EXPECT_EQ(*result, " (GPN:PartnerName)"); +} + +TEST(ParsePartnerToken, ValidWithSpaces) { + std::string raw = " ( GPN:PartnerName ; Environment ) "; + auto result = ParsePartnerToken(raw); + EXPECT_TRUE(result.Ok()); + EXPECT_EQ(*result, " (GPN:PartnerName; Environment)"); +} + +TEST(ParsePartnerToken, InvalidNoGpn) { + std::string raw = "(PartnerName; Environment)"; + auto result = ParsePartnerToken(raw); + EXPECT_FALSE(result.Ok()); +} + +TEST(ParsePartnerToken, Empty) { + std::string raw; + auto result = ParsePartnerToken(raw); + EXPECT_TRUE(result.Ok()); + EXPECT_EQ(*result, ""); +} + +TEST(ParsePartnerToken, NoMatch) { + std::string raw = "invalid"; + auto result = ParsePartnerToken(raw); + EXPECT_FALSE(result.Ok()); +} + +} // namespace +} // namespace google::cloud::odbc_bigquery_client_interface diff --git a/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc b/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc index b08a04f1bf..171c2d1a3b 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc +++ b/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc @@ -169,6 +169,7 @@ void ConnectionHandle::SetUp(Section& dsn_section, } dsn_.impersonated_email = dsn_section["SERVICEACCOUNTIMPERSONATIONEMAIL"]; + dsn_.partner_token = dsn_section["PARTNERTOKEN"]; // Populate HTAPI related configurations std::string use_default_large_results_dataset = diff --git a/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h b/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h index 04a2085138..55313d23b1 100644 --- a/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h +++ b/google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h @@ -108,6 +108,7 @@ struct Dsn { bool enable_gcd; std::string universe_domain; std::string impersonated_email; + std::string partner_token; }; class EnvironmentHandle; diff --git a/google/cloud/odbc/bq_driver/odbc_connection.cc b/google/cloud/odbc/bq_driver/odbc_connection.cc index ceddacd470..cc181ebd90 100644 --- a/google/cloud/odbc/bq_driver/odbc_connection.cc +++ b/google/cloud/odbc/bq_driver/odbc_connection.cc @@ -93,6 +93,7 @@ Authentication CreateAuth(Dsn const& dsn) { auth.conn_props.gcd.enable_gcd = dsn.enable_gcd; auth.conn_props.gcd.universe_domain = dsn.universe_domain; auth.conn_props.impersonated_email = dsn.impersonated_email; + auth.conn_props.partner_token = dsn.partner_token; return auth; } From c2bbe40f1f917b065633ef14c4aeefb02d144f53 Mon Sep 17 00:00:00 2001 From: Sachin Purohit Date: Fri, 14 Aug 2026 22:35:41 +0000 Subject: [PATCH 2/2] refactor(bq_driver): address review comments for partner token --- google/cloud/odbc/bq_client_interface/odbc_bq_client.cc | 8 ++++---- google/cloud/odbc/bq_client_interface/utils_test.cc | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc b/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc index 2c26d0f857..b1d3602424 100644 --- a/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc +++ b/google/cloud/odbc/bq_client_interface/odbc_bq_client.cc @@ -202,8 +202,9 @@ StatusRecordOr> ODBCBQClient::CreateBQClient( return partner_token_or.GetStatusRecord(); } std::string partner_token = *partner_token_or; - options.set( - {"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token}); + std::string user_agent = + "Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token; + options.set({user_agent}); if (conn_props.gcd.enable_gcd && conn_props.gcd.universe_domain != "googleapis.com") { @@ -292,8 +293,7 @@ StatusRecordOr> ODBCBQClient::CreateBQClient( // account impersonation) will hang waiting on completion events. grpc::ChannelArguments channel_arguments; - channel_arguments.SetUserAgentPrefix( - "Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token); + channel_arguments.SetUserAgentPrefix(user_agent); channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS, std::chrono::minutes(1).count()); channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS, diff --git a/google/cloud/odbc/bq_client_interface/utils_test.cc b/google/cloud/odbc/bq_client_interface/utils_test.cc index b3fd6f1960..7a97a3b5c2 100644 --- a/google/cloud/odbc/bq_client_interface/utils_test.cc +++ b/google/cloud/odbc/bq_client_interface/utils_test.cc @@ -20,7 +20,7 @@ namespace google::cloud::odbc_bigquery_client_interface { namespace { TEST(ParsePartnerToken, ValidFull) { - std::string raw = "(GPN:PartnerName; Environment)"; + std::string raw = "(GPN:PartnerName;Environment)"; auto result = ParsePartnerToken(raw); EXPECT_TRUE(result.Ok()); EXPECT_EQ(*result, " (GPN:PartnerName; Environment)");