[Python] Support custom BigQuery endpoints and universe domains

sdks/python/apache_beam/io/gcp/bigquery_tools.py

@@ -33,6 +33,7 @@
 import io
 import json
 import logging
+import os
 import re
 import sys
 import time
@@ -60,6 +61,7 @@
 from apache_beam.metrics import monitoring_infos
 from apache_beam.metrics.metric import Metrics
 from apache_beam.options import value_provider
+from apache_beam.options.pipeline_options import BigQueryOptions
 from apache_beam.options.pipeline_options import GoogleCloudOptions
 from apache_beam.options.pipeline_options import PipelineOptions
 from apache_beam.transforms import DoFn
@@ -123,6 +125,77 @@
 # For CI: temp dataset of this name are automatically deleted
 _TEMP_DATASET_PREFIX = 'beam_temp_dataset_'
 
+_DEFAULT_UNIVERSE_DOMAIN = 'googleapis.com'
+_UNIVERSE_DOMAIN_ENV = 'GOOGLE_CLOUD_UNIVERSE_DOMAIN'
+_BIGQUERY_SERVICE_PATH = '/bigquery/v2'
+
+
+def _get_bigquery_endpoint(
+    pipeline_options: Optional[PipelineOptions] = None) -> Optional[str]:
+  """Returns the BigQuery root URL set via ``--bigquery_endpoint``, or None.
+
+  The given pipeline options are checked first. If they do not set the
+  endpoint, the options of the running pipeline are checked; these are
+  available on workers, where BigQuery clients are often created without
+  access to the pipeline options.
+
+  The returned URL always has a scheme and never has a trailing slash or the
+  ``/bigquery/v2`` service path.
+  """
+  endpoint = None
+  if pipeline_options is not None:
+    endpoint = pipeline_options.view_as(BigQueryOptions).big_query_endpoint
+  if not endpoint:
+    endpoint = value_provider.RuntimeValueProvider.get_value(
+        'big_query_endpoint', str, None)
+  if not endpoint:
+    return None
+  endpoint = endpoint.strip().rstrip('/')
+  if endpoint.endswith(_BIGQUERY_SERVICE_PATH):
+    endpoint = endpoint[:-len(_BIGQUERY_SERVICE_PATH)]
+  if '://' not in endpoint:
+    endpoint = 'https://' + endpoint
+  return endpoint
+
+
+def get_bigquery_api_url(
+    pipeline_options: Optional[PipelineOptions] = None) -> str:
+  """Returns the base URL of the BigQuery v2 REST API.
+
+  The URL is resolved in the following order:
+
+  1. ``--bigquery_endpoint``, if set.
+  2. ``https://bigquery.<domain>``, if the ``GOOGLE_CLOUD_UNIVERSE_DOMAIN``
+     environment variable is set to a non-default universe domain. This
+     matches how the google-cloud-bigquery client library resolves its
+     endpoint.
+  3. The default public URL.
+  """
+  endpoint = _get_bigquery_endpoint(pipeline_options)
+  if endpoint is None:
+    universe_domain = (
+        os.environ.get(_UNIVERSE_DOMAIN_ENV, '').strip() or
+        _DEFAULT_UNIVERSE_DOMAIN)
+    if universe_domain == _DEFAULT_UNIVERSE_DOMAIN:
+      return bigquery.BigqueryV2.BASE_URL
+    endpoint = 'https://bigquery.%s' % universe_domain
+  return endpoint + _BIGQUERY_SERVICE_PATH + '/'
+
+
+def get_bigquery_client_options(
+    pipeline_options: Optional[PipelineOptions] = None) -> Optional[dict]:
+  """Returns ``client_options`` for a ``google.cloud.bigquery.Client``.
+
+  Returns None when ``--bigquery_endpoint`` is not set, so that the client
+  library resolves the endpoint itself (it honors
+  ``GOOGLE_CLOUD_UNIVERSE_DOMAIN``).
+  """
+  endpoint = _get_bigquery_endpoint(pipeline_options)
+  if endpoint is None:
+    return None
+  return {'api_endpoint': endpoint}
+
+
 # TODO(https://github.com/apache/beam/issues/25946): Add support for
 # more Beam portable schema types as Python types
 BIGQUERY_TYPE_TO_PYTHON_TYPE = {
@@ -510,6 +583,7 @@ def __init__(
       pipeline_options=None):
     self.quota_project_id = (
         quota_project_id or _quota_project_id_from_options(pipeline_options))
+    self._pipeline_options = pipeline_options
     self.client = client or BigQueryWrapper._bigquery_client(
         pipeline_options or PipelineOptions(),
         quota_project_id=self.quota_project_id)
@@ -556,7 +630,8 @@ def gcp_bq_client(self):
     """
     if self._gcp_bq_client is None:
       self._gcp_bq_client = BigQueryWrapper._gcp_bigquery_client(
-          quota_project_id=self.quota_project_id)
+          quota_project_id=self.quota_project_id,
+          pipeline_options=self._pipeline_options)
     return self._gcp_bq_client
 
   @property
@@ -1597,6 +1672,7 @@ def _bigquery_client(
     # Use explicit quota_project_id if provided, otherwise get from options
     quota_project_id = quota_project_id or _quota_project_id_from_options(
         pipeline_options)
+    url = get_bigquery_api_url(pipeline_options)
     http = get_new_http()
     if quota_project_id:
       if credentials is None:
@@ -1612,11 +1688,12 @@ def _bigquery_client(
       http = _HttpWithHeaders(
           http, {'x-goog-user-project': quota_project_id},
           hosts=[
-              urllib.parse.urlsplit(url).hostname for url in (
-                  bigquery.BigqueryV2.BASE_URL,
+              urllib.parse.urlsplit(u).hostname for u in (
+                  url, bigquery.BigqueryV2.BASE_URL,
                   bigquery.BigqueryV2.MTLS_BASE_URL)
           ])
     return bigquery.BigqueryV2(
+        url=url,
         http=http,
         credentials=credentials,
         response_encoding='utf8',
@@ -1625,17 +1702,21 @@ def _bigquery_client(
         })
 
   @staticmethod
-  def _gcp_bigquery_client(quota_project_id: str = None):
+  def _gcp_bigquery_client(
+      quota_project_id: str = None,
+      pipeline_options: Optional[PipelineOptions] = None):
     """Create a google-cloud-bigquery Client with optional quota project.
 
     The client applies the quota project to the credentials it resolves
-    itself, so no credentials are derived here.
+    itself, so no credentials are derived here. If ``--bigquery_endpoint`` is
+    set, it is used as the API endpoint; otherwise the client library resolves
+    the endpoint itself, honoring ``GOOGLE_CLOUD_UNIVERSE_DOMAIN``.
     """
-    client_options = None
+    client_options = get_bigquery_client_options(pipeline_options) or {}
     if quota_project_id:
-      client_options = {'quota_project_id': quota_project_id}
+      client_options['quota_project_id'] = quota_project_id
     return gcp_bigquery.Client(
-        client_options=client_options,
+        client_options=client_options or None,
         client_info=ClientInfo(
             user_agent="apache-beam-%s" % apache_beam.__version__))
 

sdks/python/apache_beam/io/gcp/bigquery_tools_test.py

@@ -53,6 +53,9 @@
 from apache_beam.io.gcp.internal.clients import bigquery
 from apache_beam.metrics import monitoring_infos
 from apache_beam.metrics.execution import MetricsEnvironment
+from apache_beam.options.pipeline_options import BigQueryOptions
+from apache_beam.options.pipeline_options import PipelineOptions
+from apache_beam.options.value_provider import RuntimeValueProvider
 from apache_beam.options.value_provider import StaticValueProvider
 from apache_beam.typehints.row_type import RowTypeConstraint
 from apache_beam.utils.timestamp import Timestamp
@@ -241,6 +244,131 @@ def test_calling_with_all_arguments(self):
     self.assertEqual(parsed_ref.tableId, tableId)
 
 
+@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
+class TestBigQueryEndpoint(unittest.TestCase):
+  def setUp(self):
+    self._saved_runtime_options = RuntimeValueProvider.runtime_options
+    self._saved_experiments = RuntimeValueProvider.experiments
+    RuntimeValueProvider.runtime_options = None
+    env_patcher = mock.patch.dict(os.environ)
+    env_patcher.start()
+    self.addCleanup(env_patcher.stop)
+    os.environ.pop('GOOGLE_CLOUD_UNIVERSE_DOMAIN', None)
+
+  def tearDown(self):
+    RuntimeValueProvider.runtime_options = self._saved_runtime_options
+    RuntimeValueProvider.experiments = self._saved_experiments
+
+  def test_default_endpoint(self):
+    options = PipelineOptions([])
+    self.assertIsNone(options.view_as(BigQueryOptions).bigquery_endpoint)
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(options),
+        'https://bigquery.googleapis.com/bigquery/v2/')
+    self.assertIsNone(bigquery_tools.get_bigquery_client_options(options))
+
+  def test_default_universe_domain_uses_default_endpoint(self):
+    os.environ['GOOGLE_CLOUD_UNIVERSE_DOMAIN'] = 'googleapis.com'
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(PipelineOptions([])),
+        'https://bigquery.googleapis.com/bigquery/v2/')
+
+  def test_universe_domain_from_env(self):
+    os.environ['GOOGLE_CLOUD_UNIVERSE_DOMAIN'] = 'example-tpc.goog'
+    options = PipelineOptions([])
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(options),
+        'https://bigquery.example-tpc.goog/bigquery/v2/')
+    # google-cloud-bigquery resolves the universe domain by itself.
+    self.assertIsNone(bigquery_tools.get_bigquery_client_options(options))
+
+  @parameterized.expand([
+      ('https://bq.example.com', ),
+      ('https://bq.example.com/', ),
+      ('https://bq.example.com/bigquery/v2', ),
+      ('https://bq.example.com/bigquery/v2/', ),
+      ('bq.example.com', ),
+  ])
+  def test_explicit_endpoint(self, endpoint):
+    options = PipelineOptions(['--bigquery_endpoint', endpoint])
+    self.assertEqual(
+        options.view_as(BigQueryOptions).bigquery_endpoint, endpoint)
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(options),
+        'https://bq.example.com/bigquery/v2/')
+    self.assertEqual(
+        bigquery_tools.get_bigquery_client_options(options),
+        {'api_endpoint': 'https://bq.example.com'})
+
+  def test_explicit_endpoint_takes_precedence_over_universe_domain(self):
+    os.environ['GOOGLE_CLOUD_UNIVERSE_DOMAIN'] = 'example-tpc.goog'
+    options = PipelineOptions(['--bigquery_endpoint=http://localhost:9050'])
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(options),
+        'http://localhost:9050/bigquery/v2/')
+    self.assertEqual(
+        bigquery_tools.get_bigquery_client_options(options),
+        {'api_endpoint': 'http://localhost:9050'})
+
+  def test_endpoint_from_runtime_options(self):
+    # On workers, clients may be created without the pipeline options.
+    RuntimeValueProvider.set_runtime_options(
+        {'big_query_endpoint': 'https://bq.example.com'})
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(PipelineOptions([])),
+        'https://bq.example.com/bigquery/v2/')
+    self.assertEqual(
+        bigquery_tools.get_bigquery_client_options(None),
+        {'api_endpoint': 'https://bq.example.com'})
+
+  def test_endpoint_attribute_alias_setter(self):
+    bq_options = PipelineOptions([]).view_as(BigQueryOptions)
+    bq_options.bigquery_endpoint = 'https://bq.example.com'
+    self.assertEqual(
+        bigquery_tools.get_bigquery_api_url(bq_options),
+        'https://bq.example.com/bigquery/v2/')
+
+  @mock.patch.object(bigquery_tools.auth, 'get_service_credentials')
+  def test_bigquery_client_uses_endpoint(self, unused_mock_credentials):
+    options = PipelineOptions(['--bigquery_endpoint=https://bq.example.com'])
+    client = bigquery_tools.BigQueryWrapper._bigquery_client(options)
+    self.assertEqual(client.url, 'https://bq.example.com/bigquery/v2/')
+
+  @mock.patch.object(bigquery_tools.auth, 'get_service_credentials')
+  @mock.patch.object(bigquery_tools.gcp_bigquery, 'Client')
+  def test_wrapper_applies_endpoint_to_both_clients(
+      self, mock_gcp_client, unused_mock_credentials):
+    options = PipelineOptions(['--bigquery_endpoint=https://bq.example.com'])
+    wrapper = bigquery_tools.BigQueryWrapper(pipeline_options=options)
+    self.assertEqual(wrapper.client.url, 'https://bq.example.com/bigquery/v2/')
+    # The google-cloud-bigquery client is created on first use.
+    wrapper.gcp_bq_client  # pylint: disable=pointless-statement
+    self.assertEqual(
+        mock_gcp_client.call_args.kwargs['client_options'],
+        {'api_endpoint': 'https://bq.example.com'})
+
+  @mock.patch.object(bigquery_tools.gcp_bigquery, 'Client')
+  def test_gcp_client_merges_endpoint_and_quota_project(self, mock_gcp_client):
+    options = PipelineOptions(['--bigquery_endpoint=https://bq.example.com'])
+    bigquery_tools.BigQueryWrapper._gcp_bigquery_client(
+        quota_project_id='billing', pipeline_options=options)
+    self.assertEqual(
+        mock_gcp_client.call_args.kwargs['client_options'], {
+            'api_endpoint': 'https://bq.example.com',
+            'quota_project_id': 'billing'
+        })
+
+  @mock.patch.object(bigquery_tools.auth, 'get_service_credentials')
+  def test_quota_project_header_sent_to_custom_endpoint(self, mock_credentials):
+    options = PipelineOptions(['--bigquery_endpoint=https://bq.example.com'])
+    bigquery_tools.BigQueryWrapper._bigquery_client(
+        options, quota_project_id='billing')
+    # apitools wraps the given http with credentials.authorize().
+    http = mock_credentials.return_value.authorize.call_args.args[0]
+    self.assertIsInstance(http, bigquery_tools._HttpWithHeaders)
+    self.assertIn('bq.example.com', http._hosts)
+
+
 @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed')
 class TestBigQueryWrapper(unittest.TestCase):
   def test_delete_non_existing_dataset(self):

sdks/python/apache_beam/options/pipeline_options.py

@@ -1388,6 +1388,39 @@ def get_cloud_profiler_service_name(self):
     return None
 
 
+class BigQueryOptions(PipelineOptions):
+  """BigQuery-specific options, mirroring Java's ``BigQueryOptions``.
+
+  Option names (``dest``) are chosen so that they map to the Java option names
+  when pipeline options are passed to Java, e.g. to cross-language BigQuery
+  transforms: Java converts ``big_query_endpoint`` to ``bigQueryEndpoint``.
+  """
+
+  # Alternative attribute names, mapped to the option names they stand for.
+  _ATTRIBUTE_ALIASES = {'bigquery_endpoint': 'big_query_endpoint'}
+
+  @classmethod
+  def _add_argparse_args(cls, parser):
+    parser.add_argument(
+        '--bigquery_endpoint',
+        '--big_query_endpoint',
+        dest='big_query_endpoint',
+        default=None,
+        help=(
+            'The root URL for the BigQuery API, e.g. '
+            '"https://bigquery.googleapis.com". If not set, the endpoint is '
+            'derived from the GOOGLE_CLOUD_UNIVERSE_DOMAIN environment '
+            'variable when present, and otherwise the default public URL is '
+            'used. Also passed to Java cross-language BigQuery transforms as '
+            'bigQueryEndpoint.'))
+
+  def __getattr__(self, name):
+    return super().__getattr__(self._ATTRIBUTE_ALIASES.get(name, name))
+
+  def __setattr__(self, name, value):
+    super().__setattr__(self._ATTRIBUTE_ALIASES.get(name, name), value)
+
+
 class AzureOptions(PipelineOptions):
   """Azure Blob Storage options."""
   @classmethod