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..b1d3602424 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,18 @@ StatusRecordOr> ODBCBQClient::CreateBQClient( .set_password(conn_props.proxy_options.password) .set_scheme("http")); - options.set( - {"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION)}); + 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; + 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") { @@ -282,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)); + 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.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..7a97a3b5c2 --- /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; }