From 25a975fee888b299dcbe13c9ae246e99f74aaed4 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Mon, 19 Oct 2026 22:40:13 -0400 Subject: [PATCH 1/3] [Python] Support custom BigQuery endpoints and universe domains BigQueryWrapper always sent BigQuery REST calls to bigquery.googleapis.com, so it could not be used in Trusted Partner Cloud (TPC) or other universes. Add a --bigquery_endpoint option, matching Java's bigQueryEndpoint. The BigqueryV2 URL now comes from that option, then from GOOGLE_CLOUD_UNIVERSE_DOMAIN, then the default URL. The option is also passed to google.cloud.bigquery.Client, which reads the universe domain on its own when the option is unset. Workers fall back to the running pipeline's options. --- .../apache_beam/io/gcp/bigquery_tools.py | 96 ++++++++++++-- .../apache_beam/io/gcp/bigquery_tools_test.py | 117 ++++++++++++++++++ .../apache_beam/options/pipeline_options.py | 9 ++ 3 files changed, 214 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 d1412d0f855e..84b293313073 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 io import json import logging +import os import re import sys import time @@ -123,6 +124,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(GoogleCloudOptions).bigquery_endpoint + if not endpoint: + endpoint = value_provider.RuntimeValueProvider.get_value( + 'bigquery_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.``, 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 +582,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 +629,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 +1671,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 +1687,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 +1701,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__)) 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 6df9f6c5c143..269b5a7a729a 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,8 @@ 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 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 +243,121 @@ 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.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( + 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( + {'bigquery_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'}) + + @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 da3aae4e6a17..189135e7673d 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1044,6 +1044,15 @@ def _add_argparse_args(cls, parser): 'resolved by the Google Cloud client library, which honors ' 'GOOGLE_CLOUD_UNIVERSE_DOMAIN and otherwise uses the default ' 'public URL (%s).' % cls.DATAFLOW_ENDPOINT)) + parser.add_argument( + '--bigquery_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.')) # Remote execution must check that this option is not None. parser.add_argument( '--project', From 0c31ef642052e943a1bcced281069319d8462804 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 21 Oct 2026 13:46:37 -0400 Subject: [PATCH 2/3] Move the new flag into BigQueryOptions --- .../apache_beam/io/gcp/bigquery_tools.py | 3 ++- .../apache_beam/io/gcp/bigquery_tools_test.py | 4 ++++ .../apache_beam/options/pipeline_options.py | 24 ++++++++++++------- 3 files changed, 21 insertions(+), 10 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 84b293313073..8f22f9b64a6b 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -61,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 @@ -143,7 +144,7 @@ def _get_bigquery_endpoint( """ endpoint = None if pipeline_options is not None: - endpoint = pipeline_options.view_as(GoogleCloudOptions).bigquery_endpoint + endpoint = pipeline_options.view_as(BigQueryOptions).bigquery_endpoint if not endpoint: endpoint = value_provider.RuntimeValueProvider.get_value( 'bigquery_endpoint', str, None) 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 269b5a7a729a..3b89e4476a53 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,7 @@ 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 @@ -260,6 +261,7 @@ def tearDown(self): 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/') @@ -289,6 +291,8 @@ def test_universe_domain_from_env(self): ]) 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/') diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 189135e7673d..2b58dd83248e 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1044,15 +1044,6 @@ def _add_argparse_args(cls, parser): 'resolved by the Google Cloud client library, which honors ' 'GOOGLE_CLOUD_UNIVERSE_DOMAIN and otherwise uses the default ' 'public URL (%s).' % cls.DATAFLOW_ENDPOINT)) - parser.add_argument( - '--bigquery_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.')) # Remote execution must check that this option is not None. parser.add_argument( '--project', @@ -1397,6 +1388,21 @@ def get_cloud_profiler_service_name(self): return None +class BigQueryOptions(PipelineOptions): + """BigQuery-specific options, mirroring Java's ``BigQueryOptions``.""" + @classmethod + def _add_argparse_args(cls, parser): + parser.add_argument( + '--bigquery_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.')) + + class AzureOptions(PipelineOptions): """Azure Blob Storage options.""" @classmethod From 6414324b0a64c9f4c7768877e1917fc42303e993 Mon Sep 17 00:00:00 2001 From: Shunping Huang Date: Wed, 21 Oct 2026 14:08:05 -0400 Subject: [PATCH 3/3] Pass bigquery_endpoint to Java as bigQueryEndpoint Store the option as big_query_endpoint so it is serialized as beam:option:big_query_endpoint:v1, which Java maps to the bigQueryEndpoint option used by cross-language BigQuery transforms. Both --bigquery_endpoint and --big_query_endpoint, and both attribute names on BigQueryOptions, are accepted. --- .../apache_beam/io/gcp/bigquery_tools.py | 4 +- .../apache_beam/io/gcp/bigquery_tools_test.py | 9 ++- .../apache_beam/options/pipeline_options.py | 22 +++++- .../options/pipeline_options_test.py | 67 +++++++++++++++++++ 4 files changed, 97 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 8f22f9b64a6b..ea3aa1bbbd40 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -144,10 +144,10 @@ def _get_bigquery_endpoint( """ endpoint = None if pipeline_options is not None: - endpoint = pipeline_options.view_as(BigQueryOptions).bigquery_endpoint + endpoint = pipeline_options.view_as(BigQueryOptions).big_query_endpoint if not endpoint: endpoint = value_provider.RuntimeValueProvider.get_value( - 'bigquery_endpoint', str, None) + 'big_query_endpoint', str, None) if not endpoint: return None endpoint = endpoint.strip().rstrip('/') 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 3b89e4476a53..be5441c263b9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -313,7 +313,7 @@ def test_explicit_endpoint_takes_precedence_over_universe_domain(self): def test_endpoint_from_runtime_options(self): # On workers, clients may be created without the pipeline options. RuntimeValueProvider.set_runtime_options( - {'bigquery_endpoint': 'https://bq.example.com'}) + {'big_query_endpoint': 'https://bq.example.com'}) self.assertEqual( bigquery_tools.get_bigquery_api_url(PipelineOptions([])), 'https://bq.example.com/bigquery/v2/') @@ -321,6 +321,13 @@ def test_endpoint_from_runtime_options(self): 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']) diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 2b58dd83248e..5b73b22476b4 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1389,18 +1389,36 @@ def get_cloud_profiler_service_name(self): class BigQueryOptions(PipelineOptions): - """BigQuery-specific options, mirroring Java's ``BigQueryOptions``.""" + """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.')) + '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): diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py index 2814995b067b..94083a55c6f4 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 @@ 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 @@ def test_options_are_registered_once(self): 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::v1, and Java derives the property name from . + 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()