From 46688febbbd0c96351b0f62f60b6561daa3cd001 Mon Sep 17 00:00:00 2001
From: Khushi Kathuria <khushi.kathuria@37848692.corp-partner.google.com>
Date: Tue, 14 Oct 2025 11:29:38 +0530
Subject: [PATCH] fix(client_library): PSC endpoint fix

---
 .../internal/dataset_rest_stub_factory.cc        | 16 ++++++++++++++--
 .../v2/minimal/internal/job_rest_stub_factory.cc | 14 +++++++++++++-
 .../internal/project_rest_stub_factory.cc        | 14 +++++++++++++-
 .../minimal/internal/table_rest_stub_factory.cc  | 14 +++++++++++++-
 4 files changed, 53 insertions(+), 5 deletions(-)

diff --git a/google/cloud/bigquery/v2/minimal/internal/dataset_rest_stub_factory.cc b/google/cloud/bigquery/v2/minimal/internal/dataset_rest_stub_factory.cc
index a1d735ef40a81..3e2fe13c5e696 100644
--- a/google/cloud/bigquery/v2/minimal/internal/dataset_rest_stub_factory.cc
+++ b/google/cloud/bigquery/v2/minimal/internal/dataset_rest_stub_factory.cc
@@ -31,9 +31,21 @@ std::shared_ptr<DatasetRestStub> CreateDefaultDatasetRestStub(
     local_opts.set<UnifiedCredentialsOption>(MakeGoogleDefaultCredentials());
   }
 
-  auto curl_rest_client = rest_internal::MakePooledRestClient(
-      opts.get<EndpointOption>(), local_opts);
+ std::string endpoint = local_opts.get<EndpointOption>();
+  auto psc_uris = std::getenv("PrivateServiceConnectUris"); 
+  if (psc_uris) {
+      std::string psc_prefix = "BIGQUERY=";
+      auto pos = std::string(psc_uris).find(psc_prefix);
+      if (pos != std::string::npos) {
+          auto start = pos + psc_prefix.size();
+          auto end = std::string(psc_uris).find(',', start);
+          endpoint = std::string(psc_uris).substr(start, end - start);
+      }
+  }
 
+  auto curl_rest_client = rest_internal::MakePooledRestClient(
+      endpoint, local_opts);
+  
   std::shared_ptr<DatasetRestStub> stub =
       std::make_shared<DefaultDatasetRestStub>(std::move(curl_rest_client));
 
diff --git a/google/cloud/bigquery/v2/minimal/internal/job_rest_stub_factory.cc b/google/cloud/bigquery/v2/minimal/internal/job_rest_stub_factory.cc
index 93bcef44bd9c6..7ef9a27cd4d32 100644
--- a/google/cloud/bigquery/v2/minimal/internal/job_rest_stub_factory.cc
+++ b/google/cloud/bigquery/v2/minimal/internal/job_rest_stub_factory.cc
@@ -33,8 +33,20 @@ std::shared_ptr<BigQueryJobRestStub> CreateDefaultBigQueryJobRestStub(
     local_opts.set<UnifiedCredentialsOption>(MakeGoogleDefaultCredentials());
   }
 
+ std::string endpoint = local_opts.get<EndpointOption>();
+  auto psc_uris = std::getenv("PrivateServiceConnectUris"); 
+  if (psc_uris) {
+      std::string psc_prefix = "BIGQUERY=";
+      auto pos = std::string(psc_uris).find(psc_prefix);
+      if (pos != std::string::npos) {
+          auto start = pos + psc_prefix.size();
+          auto end = std::string(psc_uris).find(',', start);
+          endpoint = std::string(psc_uris).substr(start, end - start);
+      }
+  }
+
   auto curl_rest_client = rest_internal::MakePooledRestClient(
-      opts.get<EndpointOption>(), local_opts);
+      endpoint, local_opts);
 
   std::shared_ptr<BigQueryJobRestStub> stub =
       std::make_shared<DefaultBigQueryJobRestStub>(std::move(curl_rest_client));
diff --git a/google/cloud/bigquery/v2/minimal/internal/project_rest_stub_factory.cc b/google/cloud/bigquery/v2/minimal/internal/project_rest_stub_factory.cc
index 3413b599c587b..75dbc35f5a33b 100644
--- a/google/cloud/bigquery/v2/minimal/internal/project_rest_stub_factory.cc
+++ b/google/cloud/bigquery/v2/minimal/internal/project_rest_stub_factory.cc
@@ -33,8 +33,20 @@ std::shared_ptr<ProjectRestStub> CreateDefaultProjectRestStub(
     local_opts.set<UnifiedCredentialsOption>(MakeGoogleDefaultCredentials());
   }
 
+ std::string endpoint = local_opts.get<EndpointOption>();
+  auto psc_uris = std::getenv("PrivateServiceConnectUris"); 
+  if (psc_uris) {
+      std::string psc_prefix = "BIGQUERY=";
+      auto pos = std::string(psc_uris).find(psc_prefix);
+      if (pos != std::string::npos) {
+          auto start = pos + psc_prefix.size();
+          auto end = std::string(psc_uris).find(',', start);
+          endpoint = std::string(psc_uris).substr(start, end - start);
+      }
+  }
+
   auto curl_rest_client = rest_internal::MakePooledRestClient(
-      opts.get<EndpointOption>(), local_opts);
+      endpoint, local_opts);
 
   std::shared_ptr<ProjectRestStub> stub =
       std::make_shared<DefaultProjectRestStub>(std::move(curl_rest_client));
diff --git a/google/cloud/bigquery/v2/minimal/internal/table_rest_stub_factory.cc b/google/cloud/bigquery/v2/minimal/internal/table_rest_stub_factory.cc
index 82acf8077ab83..e802811dfcc24 100644
--- a/google/cloud/bigquery/v2/minimal/internal/table_rest_stub_factory.cc
+++ b/google/cloud/bigquery/v2/minimal/internal/table_rest_stub_factory.cc
@@ -30,8 +30,20 @@ std::shared_ptr<TableRestStub> CreateDefaultTableRestStub(Options const& opts) {
     local_opts.set<UnifiedCredentialsOption>(MakeGoogleDefaultCredentials());
   }
 
+ std::string endpoint = local_opts.get<EndpointOption>();
+  auto psc_uris = std::getenv("PrivateServiceConnectUris"); 
+  if (psc_uris) {
+      std::string psc_prefix = "BIGQUERY=";
+      auto pos = std::string(psc_uris).find(psc_prefix);
+      if (pos != std::string::npos) {
+          auto start = pos + psc_prefix.size();
+          auto end = std::string(psc_uris).find(',', start);
+          endpoint = std::string(psc_uris).substr(start, end - start);
+      }
+  }
+
   auto curl_rest_client = rest_internal::MakePooledRestClient(
-      opts.get<EndpointOption>(), local_opts);
+      endpoint, local_opts);
 
   std::shared_ptr<TableRestStub> stub =
       std::make_shared<DefaultTableRestStub>(std::move(curl_rest_client));