Skip to content

Commit af29b28

Browse files
sachinprosachinpro
committed
impl: PartnerToken configuration property
1 parent 0affb0d commit af29b28

7 files changed

Lines changed: 112 additions & 4 deletions

File tree

‎google/cloud/odbc/bq_client_interface/odbc_authentication.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ struct Oauth {
8585
std::string kms_key_name;
8686
std::string psc;
8787
TPC tpc;
88+
std::string partner_token;
8889
};
8990

9091
// Returns true if all required BYOID properties are set.

‎google/cloud/odbc/bq_client_interface/odbc_bq_client.cc‎

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
#include "google/cloud/odbc/bq_client_interface/projects.h"
2020
#include "google/cloud/odbc/bq_client_interface/storage.h"
2121
#include "google/cloud/odbc/bq_client_interface/tables.h"
22+
#include "google/cloud/odbc/bq_client_interface/utils.h"
2223
#include "google/cloud/odbc/internal/status_record_or.h"
2324
#include "google/cloud/odbc/internal/version.h"
2425
#include "google/cloud/completion_queue.h"
@@ -28,6 +29,7 @@
2829
#include <absl/log/log.h>
2930
#include <grpcpp/security/tls_credentials_options.h>
3031
#include <algorithm>
32+
#include <regex>
3133

3234
namespace google::cloud::odbc_bigquery_client_interface {
3335

@@ -96,8 +98,17 @@ StatusRecordOr<std::shared_ptr<ODBCBQClient>> ODBCBQClient::CreateBQClient(
9698
options.set<google::cloud::CARootsFilePathOption>(pem_file);
9799
}
98100

101+
auto partner_token_or =
102+
::google::cloud::odbc_bigquery_client_interface::ParsePartnerToken(
103+
oauth.partner_token);
104+
if (!partner_token_or) {
105+
LOG(ERROR) << "CreateBQClient::ParsePartnerToken:: "
106+
<< partner_token_or.GetStatusRecord().message;
107+
return partner_token_or.GetStatusRecord();
108+
}
109+
std::string partner_token = *partner_token_or;
99110
options.set<google::cloud::UserAgentProductsOption>(
100-
{"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION)});
111+
{"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token});
101112

102113
StatusRecordOr<std::shared_ptr<Credentials>> credentials =
103114
CreateCredentials(oauth, options);
@@ -157,8 +168,8 @@ StatusRecordOr<std::shared_ptr<ODBCBQClient>> ODBCBQClient::CreateBQClient(
157168
read_options.set<GrpcCompletionQueueOption>(cq);
158169

159170
grpc::ChannelArguments channel_arguments;
160-
channel_arguments.SetUserAgentPrefix("Google-Bigquery-ODBC/" +
161-
std::string(DRIVER_VERSION));
171+
channel_arguments.SetUserAgentPrefix(
172+
"Google-Bigquery-ODBC/" + std::string(DRIVER_VERSION) + partner_token);
162173
channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIMEOUT_MS,
163174
std::chrono::minutes(1).count());
164175
channel_arguments.SetInt(GRPC_ARG_KEEPALIVE_TIME_MS,

‎google/cloud/odbc/bq_client_interface/utils.h‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,9 +15,13 @@
1515
#ifndef CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H
1616
#define CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H
1717

18+
#include "google/cloud/odbc/internal/sql_state_constants.h"
19+
#include "google/cloud/odbc/internal/status_record_or.h"
20+
#include "google/cloud/internal/exponential_backoff_policy.h"
1821
#include <absl/log/log.h>
1922
#include <cctype>
2023
#include <iomanip>
24+
#include <regex>
2125
#include <sstream>
2226
#include <string>
2327
#include <thread>
@@ -50,7 +54,7 @@ auto RetryLoop(Functor&& functor, std::string const& operation_name,
5054
using ReturnType = decltype(functor());
5155
ReturnType response;
5256

53-
google::cloud::ExponentialBackoffPolicy backoff_policy(
57+
google::cloud::internal::ExponentialBackoffPolicy backoff_policy(
5458
std::chrono::milliseconds(initial_delay_ms),
5559
std::chrono::milliseconds(max_delay_ms), backoff_multiplier);
5660

@@ -89,5 +93,31 @@ auto RetryLoop(Functor&& functor, std::string const& operation_name,
8993

9094
return response;
9195
}
96+
97+
inline google::cloud::odbc_internal::StatusRecordOr<std::string>
98+
ParsePartnerToken(std::string const& raw_token) {
99+
if (raw_token.empty()) {
100+
return std::string("");
101+
}
102+
std::regex pattern(R"(\(\s*(GPN:[^;]*?)\s*(?:;\s*([^)]*?))?\s*\))");
103+
std::smatch match;
104+
if (std::regex_search(raw_token, match, pattern)) {
105+
std::string gpn_part = match[1].str();
106+
std::string env_part = match[2].str();
107+
std::string partner_token = " (";
108+
partner_token += gpn_part;
109+
if (!env_part.empty()) {
110+
partner_token += "; ";
111+
partner_token += env_part;
112+
}
113+
partner_token += ")";
114+
return partner_token;
115+
}
116+
return google::cloud::odbc_internal::StatusRecord{
117+
google::cloud::odbc_internal::SQLStates::k_HY024(),
118+
"Invalid PartnerToken format."};
119+
}
120+
92121
} // namespace google::cloud::odbc_bigquery_client_interface
122+
93123
#endif // CPP_BIGQUERY_ODBC_GOOGLE_CLOUD_ODBC_BQ_CLIENT_INTERFACE_UTILS_H

‎google/cloud/odbc/bq_client_interface/utils_test.cc‎

Lines changed: 63 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,63 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// https://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
#include "google/cloud/odbc/bq_client_interface/utils.h"
16+
#include <gtest/gtest.h>
17+
#include <string>
18+
19+
namespace google::cloud::odbc_bigquery_client_interface {
20+
namespace {
21+
22+
TEST(ParsePartnerToken, ValidFull) {
23+
std::string raw = "(GPN:PartnerName; Environment)";
24+
auto result = ParsePartnerToken(raw);
25+
EXPECT_TRUE(result.Ok());
26+
EXPECT_EQ(*result, " (GPN:PartnerName; Environment)");
27+
}
28+
29+
TEST(ParsePartnerToken, ValidShort) {
30+
std::string raw = "(GPN:PartnerName)";
31+
auto result = ParsePartnerToken(raw);
32+
EXPECT_TRUE(result.Ok());
33+
EXPECT_EQ(*result, " (GPN:PartnerName)");
34+
}
35+
36+
TEST(ParsePartnerToken, ValidWithSpaces) {
37+
std::string raw = " ( GPN:PartnerName ; Environment ) ";
38+
auto result = ParsePartnerToken(raw);
39+
EXPECT_TRUE(result.Ok());
40+
EXPECT_EQ(*result, " (GPN:PartnerName; Environment)");
41+
}
42+
43+
TEST(ParsePartnerToken, InvalidNoGpn) {
44+
std::string raw = "(PartnerName; Environment)";
45+
auto result = ParsePartnerToken(raw);
46+
EXPECT_FALSE(result.Ok());
47+
}
48+
49+
TEST(ParsePartnerToken, Empty) {
50+
std::string raw = "";
51+
auto result = ParsePartnerToken(raw);
52+
EXPECT_TRUE(result.Ok());
53+
EXPECT_EQ(*result, "");
54+
}
55+
56+
TEST(ParsePartnerToken, NoMatch) {
57+
std::string raw = "invalid";
58+
auto result = ParsePartnerToken(raw);
59+
EXPECT_FALSE(result.Ok());
60+
}
61+
62+
} // namespace
63+
} // namespace google::cloud::odbc_bigquery_client_interface

‎google/cloud/odbc/bq_driver/internal/odbc_conn_handle.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -203,6 +203,7 @@ void ConnectionHandle::SetUp(Section& dsn_section,
203203
if (dsn_.byoid_token_url.empty()) {
204204
dsn_.byoid_token_url = kDefaultTokenUrl;
205205
}
206+
dsn_.partner_token = dsn_section["PARTNERTOKEN"];
206207
}
207208

208209
ConnectionHandle::ConnectionHandle(ConnectionHandle const& connectionHandle)

‎google/cloud/odbc/bq_driver/internal/odbc_conn_handle.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ struct Dsn {
103103
std::string psc;
104104
bool enable_tpc;
105105
std::string universe_domain;
106+
std::string partner_token;
106107
};
107108

108109
class EnvironmentHandle;

‎google/cloud/odbc/bq_driver/odbc_connection.cc‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -88,6 +88,7 @@ Authentication CreateAuth(Dsn const& dsn) {
8888
auth.oauth.psc = dsn.psc;
8989
auth.oauth.tpc.enable_tpc = dsn.enable_tpc;
9090
auth.oauth.tpc.universe_domain = dsn.universe_domain;
91+
auth.oauth.partner_token = dsn.partner_token;
9192
return auth;
9293
}
9394

0 commit comments

Comments
 (0)