This is an automated email from the ASF dual-hosted git repository.
shunping pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new d593d389e0b [Python] Support custom BigQuery endpoints and universe
domains (#40675)
d593d389e0b is described below
commit d593d389e0b57a43d79ca816c00982519cfe05ad
Author: Shunping Huang <[email protected]>
AuthorDate: Thu Oct 22 18:40:58 2026 -0400
[Python] Support custom BigQuery endpoints and universe domains (#40675)
---
sdks/python/apache_beam/io/gcp/bigquery_tools.py | 97 ++++++++++++++--
.../apache_beam/io/gcp/bigquery_tools_test.py | 128 +++++++++++++++++++++
.../python/apache_beam/options/pipeline_options.py | 33 ++++++
.../apache_beam/options/pipeline_options_test.py | 67 +++++++++++
4 files changed, 317 insertions(+), 8 deletions(-)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
index d1412d0f855..ea3aa1bbbd4 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py
@@ -33,6 +33,7 @@ import decimal
import io
import json
import logging
+import os
import re
import sys
import time
@@ -60,6 +61,7 @@ from apache_beam.io.gcp.internal.clients import bigquery
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 @@ _PROJECT_NAME_SEGMENT_PATTERN = r'[-a-z0-9]*[a-z0-9]'
# 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 @@ class BigQueryWrapper(object):
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 @@ class BigQueryWrapper(object):
"""
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 @@ class BigQueryWrapper(object):
# 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 @@ class BigQueryWrapper(object):
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 @@ class BigQueryWrapper(object):
})
@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__))
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
index 6df9f6c5c14..be5441c263b 100644
--- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
+++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py
@@ -53,6 +53,9 @@ from apache_beam.io.gcp.bigquery_tools import
parse_table_schema_from_json
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 @@ class TestTableReferenceParser(unittest.TestCase):
self.assertEqual(parsed_ref.tableId, tableId)
[email protected](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):
diff --git a/sdks/python/apache_beam/options/pipeline_options.py
b/sdks/python/apache_beam/options/pipeline_options.py
index da3aae4e6a1..5b73b22476b 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1388,6 +1388,39 @@ class GoogleCloudOptions(PipelineOptions):
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
diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py
b/sdks/python/apache_beam/options/pipeline_options_test.py
index 2814995b067..94083a55c6f 100644
--- a/sdks/python/apache_beam/options/pipeline_options_test.py
+++ b/sdks/python/apache_beam/options/pipeline_options_test.py
@@ -29,6 +29,7 @@ import mock
from parameterized import parameterized
from apache_beam.options.pipeline_options import AwsOptions
+from apache_beam.options.pipeline_options import BigQueryOptions
from apache_beam.options.pipeline_options import CrossLanguageOptions
from apache_beam.options.pipeline_options import DebugOptions
from apache_beam.options.pipeline_options import GoogleCloudOptions
@@ -1269,6 +1270,72 @@ class AwsOptionsTest(unittest.TestCase):
self.assertIn('s3_api_version', all_options)
+def _java_option_name(python_option_name):
+ """Converts a Python option name the way Java's PipelineOptionsTranslation
+ does (Guava CaseFormat.LOWER_UNDERSCORE to LOWER_CAMEL)."""
+ first, *rest = python_option_name.lower().split('_')
+ return first + ''.join(part.capitalize() for part in rest)
+
+
+class BigQueryOptionsTest(unittest.TestCase):
+ # Python option name (dest) -> property name in Java's
+ # org.apache.beam.sdk.io.gcp.bigquery.BigQueryOptions. Options are passed to
+ # Java (e.g. cross-language BigQuery transforms) as
+ # beam:option:<dest>:v1, and Java derives the property name from <dest>.
+ JAVA_OPTION_NAMES = {
+ 'big_query_endpoint': 'bigQueryEndpoint',
+ }
+
+ def test_option_names_map_to_java_option_names(self):
+ parser = _BeamArgumentParser()
+ BigQueryOptions._add_argparse_args(parser)
+ dests = {action.dest for action in parser._actions if action.dest !=
'help'}
+ # A new option must be added to JAVA_OPTION_NAMES, so that its name is
+ # checked against Java.
+ self.assertEqual(dests, set(self.JAVA_OPTION_NAMES))
+ for dest, java_name in self.JAVA_OPTION_NAMES.items():
+ self.assertEqual(_java_option_name(dest), java_name)
+
+ def test_options_serialized_with_java_compatible_names(self):
+ options = PipelineOptions(['--bigquery_endpoint=https://bq.example.com'])
+ fields = options.to_runner_api().fields
+ self.assertEqual(
+ fields['beam:option:big_query_endpoint:v1'].string_value,
+ 'https://bq.example.com')
+ self.assertNotIn('beam:option:bigquery_endpoint:v1', fields)
+
+ @parameterized.expand([
+ ('--bigquery_endpoint', ),
+ ('--big_query_endpoint', ),
+ ])
+ def test_endpoint_flags(self, flag):
+ options = PipelineOptions([flag, 'https://bq.example.com'])
+ self.assertEqual(
+ options.get_all_options()['big_query_endpoint'],
+ 'https://bq.example.com')
+ bq_options = options.view_as(BigQueryOptions)
+ self.assertEqual(bq_options.big_query_endpoint, 'https://bq.example.com')
+ self.assertEqual(bq_options.bigquery_endpoint, 'https://bq.example.com')
+
+ def test_endpoint_default(self):
+ bq_options = PipelineOptions([]).view_as(BigQueryOptions)
+ self.assertIsNone(bq_options.big_query_endpoint)
+ self.assertIsNone(bq_options.bigquery_endpoint)
+
+ def test_endpoint_attribute_alias_setter(self):
+ options = PipelineOptions([])
+ options.view_as(
+ BigQueryOptions).bigquery_endpoint = 'https://bq.example.com'
+ self.assertEqual(
+ options.view_as(BigQueryOptions).big_query_endpoint,
+ 'https://bq.example.com')
+ self.assertNotIn('bigquery_endpoint', options.get_all_options())
+
+ def test_unknown_attribute_still_raises(self):
+ with self.assertRaises(AttributeError):
+ getattr(PipelineOptions([]).view_as(BigQueryOptions), 'not_an_option')
+
+
if __name__ == '__main__':
logging.getLogger().setLevel(logging.INFO)
unittest.main()
