diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index d1412d0f855e..ea3aa1bbbd40 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 @@ -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.``, 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__)) 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..be5441c263b9 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.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): diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index da3aae4e6a17..5b73b22476b4 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/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 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()