From f340e948a2215214ee1b76f08858d6de629ec7d8 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:27:25 +0300 Subject: [PATCH 1/4] [Python] Support a custom quota project for BigQuery IO Add a quota_project_id parameter to ReadFromBigQuery and a --quota_project_id pipeline option (GoogleCloudOptions) that attribute BigQuery API quota and billing to a project other than the one the data resides in. The quota project is sent as the x-goog-user-project header on the apitools BigQuery client, and through client_options on the google-cloud-bigquery Client (streaming inserts) and the BigQuery Storage Read client. The shared pipeline credentials are left untouched instead of being copied with with_quota_project(), which for some credential types drops the cached token. The header is only added to requests to the BigQuery API hosts: the credentials refresh their token through the same Http, and the token endpoint, metadata server and IAM Credentials API must not receive it. It is honored by ReadFromBigQuery (EXPORT and DIRECT_READ, including temporary dataset cleanup), ReadAllFromBigQuery, WriteToBigQuery with FILE_LOADS or STREAMING_INSERTS, and ReadBigQueryChangeHistory. The value is resolved once at expand() time and passed to each DoFn so it reaches separately deserialized SDF copies. STORAGE_WRITE_API is Java-backed and not covered yet. Generated-By: Claude Opus 5.5 --- sdks/python/apache_beam/io/gcp/bigquery.py | 121 ++++-- .../io/gcp/bigquery_change_history.py | 42 ++- .../io/gcp/bigquery_change_history_test.py | 75 ++++ .../apache_beam/io/gcp/bigquery_file_loads.py | 46 ++- .../io/gcp/bigquery_file_loads_test.py | 47 +++ .../io/gcp/bigquery_read_internal.py | 11 +- .../apache_beam/io/gcp/bigquery_test.py | 355 ++++++++++++++++++ .../apache_beam/io/gcp/bigquery_tools.py | 162 +++++++- .../apache_beam/io/gcp/bigquery_tools_test.py | 232 +++++++++++- .../apache_beam/options/pipeline_options.py | 11 + 10 files changed, 1022 insertions(+), 80 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 2a52b31edf7f..f3d64bfcfa6c 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -720,7 +720,8 @@ def __init__( step_name=None, unique_id=None, temp_dataset=None, - query_priority=BigQueryQueryPriority.BATCH): + query_priority=BigQueryQueryPriority.BATCH, + quota_project_id=None): if table is not None and query is not None: raise ValueError( 'Both a BigQuery table and a query were specified.' @@ -754,6 +755,7 @@ def __init__( self.use_json_exports = use_json_exports self.temp_dataset = temp_dataset self.query_priority = query_priority + self.quota_project_id = quota_project_id self._job_name = job_name or 'BQ_EXPORT_JOB' self._step_name = step_name self._source_uuid = unique_id @@ -773,12 +775,15 @@ def display_data(self): 'use_legacy_sql': self.use_legacy_sql, 'bigquery_job_labels': json.dumps(self.bigquery_job_labels), 'export_file_format': export_format, + 'quota_project_id': DisplayDataItem( + self._get_quota_project_id(), + label='BigQuery Quota Project').drop_if_none(), 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs."), } def estimate_size(self): - bq = bigquery_tools.BigQueryWrapper.from_pipeline_options(self.options) + bq = self._create_bq_wrapper() if self.table_reference is not None: table_ref = self.table_reference if (isinstance(self.table_reference, vp.ValueProvider) and @@ -844,6 +849,23 @@ def _get_project(self): project = self.project return project + def _get_quota_project_id(self): + """Returns the explicit quota project ID, else the pipeline option.""" + quota_project_id = bigquery_tools._quota_project_id_from_options( + self.options) + return self.quota_project_id or quota_project_id + + def _create_bq_wrapper(self, **kwargs): + """Creates a BigQueryWrapper for the API calls made by this source. + + Every BigQuery client this source uses is built here, so that quota + attribution cannot be forgotten at an individual call site. + """ + return bigquery_tools.BigQueryWrapper( + pipeline_options=self.options, + quota_project_id=self.quota_project_id, + **kwargs) + def _create_source(self, path, coder): if not self.use_json_exports: return create_avro_source(path, validate=self.validate) @@ -858,10 +880,9 @@ def _create_source(self, path, coder): def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.export_result is None: - bq = bigquery_tools.BigQueryWrapper( + bq = self._create_bq_wrapper( temp_dataset_id=( - self.temp_dataset.datasetId if self.temp_dataset else None), - client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options)) + self.temp_dataset.datasetId if self.temp_dataset else None)) if self.query is not None: self._setup_temporary_dataset(bq) @@ -1051,7 +1072,8 @@ def __init__( temp_dataset: Optional[DatasetReference] = None, temp_table: Optional[TableReference] = None, use_native_datetime: Optional[bool] = False, - timeout: Optional[float] = None): + timeout: Optional[float] = None, + quota_project_id: Optional[str] = None): if table is not None and query is not None: raise ValueError( @@ -1090,6 +1112,7 @@ def __init__( self._job_name = job_name or 'BQ_DIRECT_READ_JOB' self._step_name = step_name self._source_uuid = unique_id + self.quota_project_id = quota_project_id def _get_project(self): """Returns the project that queries and exports will be billed to.""" @@ -1101,6 +1124,23 @@ def _get_project(self): return project return self.project + def _get_quota_project_id(self): + """Returns the explicit quota project ID, else the pipeline option.""" + quota_project_id = bigquery_tools._quota_project_id_from_options( + self.pipeline_options) + return self.quota_project_id or quota_project_id + + def _create_bq_wrapper(self, **kwargs): + """Creates a BigQueryWrapper for the API calls made by this source. + + Every BigQuery client this source uses is built here, so that quota + attribution cannot be forgotten at an individual call site. + """ + return bigquery_tools.BigQueryWrapper( + pipeline_options=self.pipeline_options, + quota_project_id=self.quota_project_id, + **kwargs) + def _get_parent_project(self): """Returns the project that will be billed.""" if self.temp_table: @@ -1188,14 +1228,16 @@ def display_data(self): 'use_native_datetime': self.use_native_datetime, 'selected_fields': str(self.selected_fields), 'row_restriction': str(self.row_restriction), + 'quota_project_id': DisplayDataItem( + self._get_quota_project_id(), + label='BigQuery Quota Project').drop_if_none(), 'launchesBigQueryJobs': DisplayDataItem( True, label="This Dataflow job launches bigquery jobs."), } def estimate_size(self): # Returns the pre-filtering size of the (temporary) table being read. - bq = bigquery_tools.BigQueryWrapper.from_pipeline_options( - self.pipeline_options) + bq = self._create_bq_wrapper() if self.table_reference is not None: table_ref = self.table_reference if (isinstance(self.table_reference, vp.ValueProvider) and @@ -1242,10 +1284,8 @@ def estimate_size(self): def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.split_result is None: - bq = bigquery_tools.BigQueryWrapper( - temp_table_ref=(self.temp_table if self.temp_table else None), - client=bigquery_tools.BigQueryWrapper._bigquery_client( - self.pipeline_options)) + bq = self._create_bq_wrapper( + temp_table_ref=(self.temp_table if self.temp_table else None)) if self.query is not None: self._setup_temporary_dataset(bq) @@ -1278,7 +1318,8 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.row_restriction is not None: requested_session.read_options.row_restriction = self.row_restriction - storage_client = bq_storage.BigQueryReadClient() + storage_client = bigquery_tools.create_bq_storage_client( + self._get_quota_project_id()) stream_count = self._get_stream_count(bq, desired_bundle_size) parent = 'projects/{}'.format(self.table_reference.projectId) @@ -1302,8 +1343,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): self.split_result = [ _CustomBigQueryStorageStreamSource( - stream.name, self.use_native_datetime, self.timeout) - for stream in read_session.streams + stream.name, + self.use_native_datetime, + self.timeout, + self._get_quota_project_id()) for stream in read_session.streams ] for source in self.split_result: @@ -1337,10 +1380,12 @@ def __init__( self, read_stream_name: str, use_native_datetime: Optional[bool] = True, - timeout: Optional[float] = None): + timeout: Optional[float] = None, + quota_project_id: Optional[str] = None): self.read_stream_name = read_stream_name self.use_native_datetime = use_native_datetime self.timeout = timeout + self.quota_project_id = quota_project_id def display_data(self): return { @@ -1363,7 +1408,10 @@ def split(self, desired_bundle_size, start_position=None, stop_position=None): return SourceBundle( weight=1.0, source=_CustomBigQueryStorageStreamSource( - self.read_stream_name, self.use_native_datetime), + self.read_stream_name, + self.use_native_datetime, + self.timeout, + self.quota_project_id), start_position=None, stop_position=None) @@ -1399,7 +1447,8 @@ def retry_delay_callback(delay): def read_arrow(self): - storage_client = bq_storage.BigQueryReadClient() + storage_client = bigquery_tools.create_bq_storage_client( + self.quota_project_id) read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback} if self.timeout is not None: read_rows_kwargs['timeout'] = self.timeout @@ -1418,7 +1467,8 @@ def read_arrow(self): yield py_row def read_avro(self): - storage_client = bq_storage.BigQueryReadClient() + storage_client = bigquery_tools.create_bq_storage_client( + self.quota_project_id) read_rows_kwargs = {'retry_delay_callback': self.retry_delay_callback} if self.timeout is not None: read_rows_kwargs['timeout'] = self.timeout @@ -1499,7 +1549,8 @@ def __init__( with_batched_input=False, ignore_unknown_columns=False, max_retries=MAX_INSERT_RETRIES, - max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE): + max_insert_payload_size=MAX_INSERT_PAYLOAD_SIZE, + quota_project_id=None): """Initialize a WriteToBigQuery transform. Args: @@ -1552,9 +1603,12 @@ def __init__( backoffs (effectively retry forever). max_insert_payload_size: The maximum byte size for a BigQuery legacy streaming insert payload. + quota_project_id: Optional GCP project ID to use for quota and billing of + BigQuery API requests. """ self.schema = schema self.test_client = test_client + self.quota_project_id = quota_project_id self.create_disposition = create_disposition self.write_disposition = write_disposition if write_disposition in (BigQueryDisposition.WRITE_EMPTY, @@ -1633,7 +1687,7 @@ def start_bundle(self): if not self.bigquery_wrapper: self.bigquery_wrapper = bigquery_tools.BigQueryWrapper( - client=self.test_client) + client=self.test_client, quota_project_id=self.quota_project_id) ( bigquery_tools.BigQueryWrapper.HISTOGRAM_METRIC_LOGGER. @@ -1958,6 +2012,8 @@ def process(self, element): yield (key, (value, insert_id)) def expand(self, input): + quota_project_id = bigquery_tools._quota_project_id_from_options( + input.pipeline.options) bigquery_write_fn = BigQueryWriteFn( schema=self.schema, batch_size=self.batch_size, @@ -1971,7 +2027,8 @@ def expand(self, input): ignore_unknown_columns=self.ignore_unknown_columns, with_batched_input=self.with_auto_sharding, max_retries=self._max_retries, - max_insert_payload_size=self._max_insert_payload_size) + max_insert_payload_size=self._max_insert_payload_size, + quota_project_id=quota_project_id) def _add_random_shard(element): key = element[0] @@ -3043,6 +3100,12 @@ class ReadFromBigQuery(PTransform): PCollection with a schema and yielding Beam Rows via the option `BEAM_ROW`. For more information on schemas, see https://beam.apache.org/documentation/programming-guide/#what-is-a-schema) + quota_project_id (str): The GCP project ID to use for quota and billing + of BigQuery API requests issued by this transform, if different from + the project the data resides in. Falls back to the + ``--quota_project_id`` pipeline option if not set. The credentials + used must have the ``serviceusage.services.use`` permission on that + project. query_output_schema: Required when output_type is 'BEAM_ROW' and a query is specified. A BigQuery schema describing the query result columns, since the schema cannot be auto-derived from an existing table when @@ -3143,10 +3206,12 @@ def _expand_output_type(self, output_pcollection): '%s: table must be of type string' '; got a callable instead' % self.__class__.__name__) return output_pcollection | bigquery_schema_tools.convert_to_usertype( - bigquery_tools.BigQueryWrapper().get_table( - project_id=table_details.projectId, - dataset_id=table_details.datasetId, - table_id=table_details.tableId).schema, + bigquery_tools.BigQueryWrapper( + pipeline_options=output_pcollection.pipeline.options, + quota_project_id=self._kwargs.get('quota_project_id')).get_table( + project_id=table_details.projectId, + dataset_id=table_details.datasetId, + table_id=table_details.tableId).schema, self._kwargs.get('selected_fields', None)) else: raise ValueError( @@ -3228,7 +3293,9 @@ def _get_pipeline_details(unused_elm): bigquery_dataset_labels=self.bigquery_dataset_labels, *self._args, **self._kwargs)) - | _PassThroughThenCleanupTempDatasets(project_to_cleanup_pcoll)) + | _PassThroughThenCleanupTempDatasets( + project_to_cleanup_pcoll, + quota_project_id=self._kwargs.get('quota_project_id'))) class ReadFromBigQueryRequest: diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py index 1e90c2974778..9f836190d9ca 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py @@ -416,7 +416,8 @@ def __init__( start_time: Timestamp, stop_time: Timestamp, poll_interval: Duration, - location: Optional[str] = None) -> None: + location: Optional[str] = None, + quota_project_id: Optional[str] = None) -> None: self._table = table self._project = project self._change_function = change_function @@ -425,9 +426,11 @@ def __init__( self._stop_time = stop_time self._poll_interval = poll_interval self._location = location + self._quota_project_id = quota_project_id def setup(self) -> None: - self._bq_wrapper = bigquery_tools.BigQueryWrapper() + self._bq_wrapper = bigquery_tools.BigQueryWrapper( + quota_project_id=self._quota_project_id) if self._location is None: table_ref = bigquery_tools.parse_table_reference( self._table, project=self._project) @@ -570,7 +573,8 @@ def __init__( change_type_column: str = 'change_type', change_timestamp_column: str = 'change_timestamp', columns: Optional[list[str]] = None, - row_filter: Optional[str] = None) -> None: + row_filter: Optional[str] = None, + quota_project_id: Optional[str] = None) -> None: self._table = table self._project = project self._change_function = change_function @@ -580,9 +584,11 @@ def __init__( self._change_timestamp_column = change_timestamp_column self._columns = columns self._row_filter = row_filter + self._quota_project_id = quota_project_id def setup(self) -> None: - self._bq_wrapper = bigquery_tools.BigQueryWrapper() + self._bq_wrapper = bigquery_tools.BigQueryWrapper( + quota_project_id=self._quota_project_id) if self._location is None: table_ref = bigquery_tools.parse_table_reference( self._table, project=self._project) @@ -711,11 +717,13 @@ def __init__( batch_arrow_read: bool = True, change_timestamp_column: str = 'change_timestamp', max_split_rounds: int = 1, - emit_raw_batches: bool = False) -> None: + emit_raw_batches: bool = False, + quota_project_id: Optional[str] = None) -> None: self._batch_arrow_read = batch_arrow_read self._change_timestamp_column = change_timestamp_column self._max_split_rounds = max_split_rounds self._emit_raw_batches = emit_raw_batches + self._quota_project_id = quota_project_id self._storage_client = None def _ensure_client(self) -> None: @@ -727,7 +735,8 @@ def _ensure_client(self) -> None: """ if self._storage_client is None: _LOGGER.info('[Read] creating BigQueryReadClient') - self._storage_client = bq_storage.BigQueryReadClient() + self._storage_client = bigquery_tools.create_bq_storage_client( + self._quota_project_id) def setup(self) -> None: self._ensure_client() @@ -1072,9 +1081,13 @@ class _CleanupTempTablesFn(beam.DoFn): STREAMS_READ = beam.transforms.userstate.CombiningValueStateSpec( 'streams_read', sum) + def __init__(self, quota_project_id: Optional[str] = None) -> None: + self._quota_project_id = quota_project_id + def setup(self) -> None: _LOGGER.info('[Cleanup] setup: creating BigQueryWrapper') - self._bq_wrapper = bigquery_tools.BigQueryWrapper() + self._bq_wrapper = bigquery_tools.BigQueryWrapper( + quota_project_id=self._quota_project_id) def process( self, @@ -1247,6 +1260,9 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: 'project must be specified either in ReadBigQueryChangeHistory ' 'or in pipeline options (--project)') + quota_project_id = bigquery_tools._quota_project_id_from_options( + pbegin.pipeline.options) + start_time = Timestamp(self._start_time or time.time()) stop_time = ( Timestamp(self._stop_time) @@ -1290,7 +1306,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: start_time=start_time, stop_time=stop_time, poll_interval=poll_interval, - location=self._location))) + location=self._location, + quota_project_id=quota_project_id))) # CommitQueryResults: Reshuffle commits _QueryResult (temp table ref) # so that if the Read SDF retries, it re-reads the existing temp table @@ -1311,7 +1328,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: change_type_column=self._change_type_column, change_timestamp_column=self._change_timestamp_column, columns=self._columns, - row_filter=self._row_filter)) + row_filter=self._row_filter, + quota_project_id=quota_project_id)) | 'CommitQueryResults' >> beam.Reshuffle()) emit_raw = self._reshuffle_decompress @@ -1321,7 +1339,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: batch_arrow_read=self._batch_arrow_read, change_timestamp_column=self._change_timestamp_column, max_split_rounds=self._max_split_rounds, - emit_raw_batches=emit_raw)) + emit_raw_batches=emit_raw, + quota_project_id=quota_project_id)) if emit_raw: read_sdf = read_sdf.with_output_types(tuple[bytes, bytes]) else: @@ -1334,7 +1353,8 @@ def expand(self, pbegin: beam.pvalue.PBegin) -> beam.PCollection: _ = ( read_outputs[_CLEANUP_TAG] - | 'CleanupTempTables' >> beam.ParDo(_CleanupTempTablesFn())) + | 'CleanupTempTables' >> beam.ParDo( + _CleanupTempTablesFn(quota_project_id=quota_project_id))) if emit_raw: # Reshuffle raw Arrow batches for fan-out, then decompress and diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py index 11b467f26d49..e1b1eb69ae38 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py @@ -23,12 +23,18 @@ - ReadBigQueryChangeHistory validation """ +import contextlib import datetime import unittest +from unittest import mock +from apache_beam.io.gcp import bigquery_change_history as bqch +from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp.bigquery_change_history import ReadBigQueryChangeHistory from apache_beam.io.gcp.bigquery_change_history import build_changes_query from apache_beam.io.gcp.bigquery_change_history import compute_ranges +from apache_beam.options.pipeline_options import PipelineOptions +from apache_beam.testing.test_pipeline import TestPipeline from apache_beam.utils.timestamp import Duration from apache_beam.utils.timestamp import Timestamp @@ -205,5 +211,74 @@ def test_default_buffer(self): self.assertEqual(t._buffer_sec, 10) +@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') +class QuotaProjectTest(unittest.TestCase): + """Tests that the quota project reaches every BigQuery client.""" + _QUOTA = 'my-billing-project' + + def _assert_wrapper_quota(self, fn): + with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper: + fn.setup() + self.assertEqual( + mock_wrapper.call_args.kwargs['quota_project_id'], self._QUOTA) + + def test_poll_fn_setup_passes_quota_project(self): + fn = bqch._PollChangeHistoryFn( + table='proj:dataset.table', + project='p', + change_function='APPENDS', + buffer=Duration(seconds=10), + start_time=Timestamp(0), + stop_time=Timestamp(100), + poll_interval=Duration(seconds=60), + quota_project_id=self._QUOTA) + self._assert_wrapper_quota(fn) + + def test_execute_query_fn_setup_passes_quota_project(self): + fn = bqch._ExecuteQueryFn( + table='proj:dataset.table', + project='p', + change_function='APPENDS', + temp_dataset='tmp', + location='US', + quota_project_id=self._QUOTA) + self._assert_wrapper_quota(fn) + + def test_cleanup_fn_setup_passes_quota_project(self): + self._assert_wrapper_quota(bqch._CleanupTempTablesFn(self._QUOTA)) + + def test_read_sdf_creates_storage_client_with_quota_project(self): + fn = bqch._ReadStorageStreamsSDF(quota_project_id=self._QUOTA) + with mock.patch.object(bigquery_tools, + 'create_bq_storage_client') as mock_create: + fn.setup() + mock_create.assert_called_once_with(self._QUOTA) + + def test_expand_resolves_quota_project_from_options(self): + options = PipelineOptions( + ['--quota_project_id=' + self._QUOTA, '--no_pipeline_type_check']) + classes = [ + '_PollChangeHistoryFn', + '_ExecuteQueryFn', + '_ReadStorageStreamsSDF', + '_CleanupTempTablesFn' + ] + p = TestPipeline(options=options) + with contextlib.ExitStack() as stack: + mocks = [ + stack.enter_context( + mock.patch.object( + getattr(bqch, name), + '__init__', + autospec=True, + side_effect=getattr(bqch, name).__init__)) for name in classes + ] + _ = p | ReadBigQueryChangeHistory( + table='p:d.t', project='p', start_time=0, temp_dataset='tmp') + for name, m in zip(classes, mocks): + self.assertEqual( + m.call_args.kwargs['quota_project_id'], self._QUOTA, name) + + if __name__ == '__main__': unittest.main() diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py index dabe80c5c27e..79940d109f07 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py @@ -383,16 +383,19 @@ def __init__( test_client=None, additional_bq_parameters=None, step_name=None, - load_job_project_id=None): + load_job_project_id=None, + quota_project_id=None): self.project = project self._test_client = test_client + self._quota_project_id = quota_project_id self._write_disposition = write_disposition self._additional_bq_parameters = additional_bq_parameters or {} self._step_name = step_name self._load_job_project_id = load_job_project_id def start_bundle(self): - self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self._test_client) + self.bq_wrapper = bigquery_tools.BigQueryWrapper( + client=self._test_client, quota_project_id=self._quota_project_id) self._bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] @@ -525,11 +528,13 @@ def __init__( write_disposition=None, test_client=None, step_name=None, - load_job_project_id=None): + load_job_project_id=None, + quota_project_id=None): self.project = project self.create_disposition = create_disposition self.write_disposition = write_disposition self.test_client = test_client + self.quota_project_id = quota_project_id self._observed_tables = set() self.bq_io_metadata = None self._step_name = step_name @@ -545,7 +550,8 @@ def setup(self): self._observed_tables = set() def start_bundle(self): - self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) + self.bq_wrapper = bigquery_tools.BigQueryWrapper( + client=self.test_client, quota_project_id=self.quota_project_id) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] @@ -677,10 +683,12 @@ def __init__( additional_bq_parameters=None, source_format=None, step_name=None, - load_job_project_id=None): + load_job_project_id=None, + quota_project_id=None): self.schema = schema self.project = project self.test_client = test_client + self.quota_project_id = quota_project_id self.temporary_tables = temporary_tables self.additional_bq_parameters = additional_bq_parameters or {} self.source_format = source_format @@ -709,7 +717,8 @@ def display_data(self): return result def start_bundle(self): - self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) + self.bq_wrapper = bigquery_tools.BigQueryWrapper( + client=self.test_client, quota_project_id=self.quota_project_id) if not self.bq_io_metadata: self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) self.pending_jobs = [] @@ -928,11 +937,13 @@ def process(self, element): class DeleteTablesFn(beam.DoFn): - def __init__(self, test_client=None): + def __init__(self, test_client=None, quota_project_id=None): self.test_client = test_client + self.quota_project_id = quota_project_id def start_bundle(self): - self.bq_wrapper = bigquery_tools.BigQueryWrapper(client=self.test_client) + self.bq_wrapper = bigquery_tools.BigQueryWrapper( + client=self.test_client, quota_project_id=self.quota_project_id) def process(self, table_reference): _LOGGER.info("Deleting table %s", table_reference) @@ -998,6 +1009,7 @@ def __init__( self.test_client = test_client self.schema = schema + self.quota_project_id = None self._temp_file_format = temp_file_format or bigquery_tools.FileFormat.JSON # If we have multiple destinations, then we will have multiple load jobs, @@ -1207,7 +1219,8 @@ def _load_data( additional_bq_parameters=self.additional_bq_parameters, source_format=self._temp_file_format, step_name=step_name, - load_job_project_id=self.load_job_project_id), + load_job_project_id=self.load_job_project_id, + quota_project_id=self.quota_project_id), load_job_name_pcv, *self.schema_side_inputs).with_outputs( TriggerLoadJobs.TEMP_TABLES, @@ -1228,7 +1241,8 @@ def _load_data( test_client=self.test_client, additional_bq_parameters=self.additional_bq_parameters, step_name=step_name, - load_job_project_id=self.load_job_project_id), + load_job_project_id=self.load_job_project_id, + quota_project_id=self.quota_project_id), schema_mod_job_name_pcv)) if self.write_disposition in ('WRITE_EMPTY', 'WRITE_TRUNCATE'): @@ -1259,7 +1273,8 @@ def _load_data( write_disposition=self.write_disposition, test_client=self.test_client, step_name=step_name, - load_job_project_id=self.load_job_project_id), + load_job_project_id=self.load_job_project_id, + quota_project_id=self.quota_project_id), copy_job_name_pcv, pvalue.AsIter(schema_mod_job_ids_pc)).with_outputs( TriggerCopyJobs.TRIGGER_DELETE_TEMP_TABLES, main='main')) @@ -1275,7 +1290,9 @@ def _load_data( | "RemoveTempTables/DeduplicateTables" >> beam.GroupByKey() | "RemoveTempTables/GetTableNames" >> beam.Keys() | "RemoveTempTables/Delete" >> beam.ParDo( - DeleteTablesFn(self.test_client))) + DeleteTablesFn( + test_client=self.test_client, + quota_project_id=self.quota_project_id))) # Load data directly to destination table destination_load_job_ids_pc = ( @@ -1290,7 +1307,8 @@ def _load_data( additional_bq_parameters=self.additional_bq_parameters, source_format=self._temp_file_format, step_name=step_name, - load_job_project_id=self.load_job_project_id), + load_job_project_id=self.load_job_project_id, + quota_project_id=self.quota_project_id), load_job_name_pcv, *self.schema_side_inputs).with_outputs( TriggerLoadJobs.ONGOING_JOBS, main='main') @@ -1305,6 +1323,8 @@ def _load_data( def expand(self, pcoll): p = pcoll.pipeline self.project = self.project or p.options.view_as(GoogleCloudOptions).project + self.quota_project_id = bigquery_tools._quota_project_id_from_options( + p.options) try: step_name = self.label except AttributeError: diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py index 7ae0ccc8e32f..7ccda7ecad8d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads_test.py @@ -19,6 +19,7 @@ # pytype: skip-file +import contextlib import logging import os import secrets @@ -1259,6 +1260,52 @@ def __call__(self): assert_that(jobs, equal_to(expected_jobs), label='CheckJobs') +class TestFileLoadsQuotaProject(_TestCaseWithTempDirCleanUp): + """Tests that the file loads DoFns attribute quota to the quota project.""" + @parameterized.expand([ + param(dofn=bqfl.UpdateDestinationSchema, quota_attr='_quota_project_id'), + param(dofn=bqfl.TriggerCopyJobs, quota_attr='quota_project_id'), + param(dofn=bqfl.TriggerLoadJobs, quota_attr='quota_project_id'), + param(dofn=bqfl.DeleteTablesFn, quota_attr='quota_project_id'), + ]) + def test_start_bundle_passes_quota_project(self, dofn, quota_attr): + fn = dofn(quota_project_id='my-billing-project') + self.assertEqual(getattr(fn, quota_attr), 'my-billing-project') + with mock.patch.object(bigquery_tools, 'BigQueryWrapper') as mock_wrapper: + fn.start_bundle() + self.assertEqual( + mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project') + + def test_expand_resolves_quota_project_from_options(self): + options = PipelineOptions(['--quota_project_id=my-billing-project']) + transform = bqfl.BigQueryBatchFileLoads( + 'project:dataset.table', + custom_gcs_temp_location=self._new_tempdir(), + test_client=mock.Mock(), + validate=False, + temp_file_format=bigquery_tools.FileFormat.JSON) + p = TestPipeline(options=options) + dofns = [ + bqfl.UpdateDestinationSchema, + bqfl.TriggerCopyJobs, + bqfl.TriggerLoadJobs, + bqfl.DeleteTablesFn + ] + with contextlib.ExitStack() as stack: + mocks = [ + stack.enter_context( + mock.patch.object( + dofn, '__init__', autospec=True, side_effect=dofn.__init__)) + for dofn in dofns + ] + _ = p | beam.Create(_ELEMENTS) | transform + self.assertEqual(transform.quota_project_id, 'my-billing-project') + for dofn, m in zip(dofns, mocks): + self.assertTrue(m.call_args_list, dofn.__name__) + for c in m.call_args_list: + self.assertEqual(c.kwargs['quota_project_id'], 'my-billing-project') + + class BigQueryFileLoadsIT(unittest.TestCase): BIG_QUERY_SCHEMA = ( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py index 136b3cc56b7e..2fa95a3d1794 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py @@ -142,11 +142,13 @@ class _PassThroughThenCleanupTempDatasets(PTransform): Utilizes readiness of PCollection to trigger DoFn. """ - def __init__(self, side_input=None): + def __init__(self, side_input=None, quota_project_id=None): self.side_input = side_input + self.quota_project_id = quota_project_id def expand(self, input): pipeline_options = input.pipeline.options + quota_project_id = self.quota_project_id class PassThrough(beam.DoFn): def process(self, element): @@ -154,8 +156,9 @@ def process(self, element): class CleanUpProjects(beam.DoFn): def process(self, unused_element, unused_signal, pipeline_details): - bq = bigquery_tools.BigQueryWrapper.from_pipeline_options( - pipeline_options) + bq = bigquery_tools.BigQueryWrapper( + pipeline_options=pipeline_options, + quota_project_id=quota_project_id) pipeline_details = pipeline_details[0] if 'temp_table_ref' in pipeline_details.keys(): temp_table_ref = pipeline_details['temp_table_ref'] @@ -251,7 +254,7 @@ def _get_temp_dataset_project(self): def start_bundle(self): self.bq = bigquery_tools.BigQueryWrapper( temp_dataset_id=self._get_temp_dataset_id(), - client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options)) + pipeline_options=self.options) def process(self, element: 'ReadFromBigQueryRequest') -> Iterable[BoundedSource]: diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index 22b617fa8741..ae03104612c9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -53,6 +53,7 @@ from apache_beam.io.gcp.bigquery import _StreamToBigQuery from apache_beam.io.gcp.bigquery_read_internal import _BigQueryReadSplit from apache_beam.io.gcp.bigquery_read_internal import _JsonToDictCoder +from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanupTempDatasets from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper @@ -65,6 +66,7 @@ from apache_beam.io.gcp.tests.bigquery_matcher import BigQueryTableMatcher from apache_beam.metrics.metric import Lineage from apache_beam.options import value_provider +from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions from apache_beam.options.value_provider import RuntimeValueProvider @@ -338,6 +340,10 @@ def test_repeatable_field_is_properly_converted(self): class TestReadFromBigQuery(unittest.TestCase): @classmethod def setUpClass(cls): + cls.env_patch = mock.patch.dict( + os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}) + cls.env_patch.start() + class UserDefinedOptions(PipelineOptions): @classmethod def _add_argparse_args(cls, parser): @@ -351,6 +357,7 @@ def tearDown(self): @classmethod def tearDownClass(cls): + cls.env_patch.stop() # Unset the option added in setupClass to avoid interfere with other tests. # Force a gc so PipelineOptions.__subclass__() no longer contains it. del cls.UserDefinedOptions @@ -879,6 +886,343 @@ def test_expand_output_type_uses_query_schema(self): mock_convert.assert_called_once_with(schema, None) +@unittest.skipIf( + HttpError is None or gcp_bigquery is None, + 'GCP dependencies are not installed') +class TestReadFromBigQueryQuotaProject(unittest.TestCase): + """Tests for quota_project_id in ReadFromBigQuery sources.""" + def test_quota_project_id_from_pipeline_options(self): + """Test that quota_project_id is read from GoogleCloudOptions.""" + options = PipelineOptions(['--quota_project_id=my-billing-project']) + gcp_options = options.view_as(GoogleCloudOptions) + self.assertEqual(gcp_options.quota_project_id, 'my-billing-project') + + def test_quota_project_id_none_by_default_in_options(self): + """Test that quota_project_id is None by default in options.""" + options = PipelineOptions([]) + gcp_options = options.view_as(GoogleCloudOptions) + self.assertIsNone(gcp_options.quota_project_id) + + def test_export_source_explicit_quota_project(self): + """Test that explicit quota_project_id is stored in + _CustomBigQuerySource.""" + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + quota_project_id='my-billing-project') + self.assertEqual(source.quota_project_id, 'my-billing-project') + self.assertEqual(source._get_quota_project_id(), 'my-billing-project') + + def test_export_source_gets_quota_from_options(self): + """Test that _CustomBigQuerySource falls back to options for + quota_project_id.""" + options = PipelineOptions(['--quota_project_id=my-billing-project']) + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + pipeline_options=options) + self.assertEqual(source._get_quota_project_id(), 'my-billing-project') + + def test_export_source_explicit_overrides_options(self): + """Test that explicit quota_project_id overrides options.""" + options = PipelineOptions(['--quota_project_id=options-project']) + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + pipeline_options=options, + quota_project_id='explicit-project') + self.assertEqual(source._get_quota_project_id(), 'explicit-project') + + def test_storage_source_explicit_quota_project(self): + """Test that explicit quota_project_id is stored in + _CustomBigQueryStorageSource.""" + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + quota_project_id='my-billing-project') + self.assertEqual(source.quota_project_id, 'my-billing-project') + self.assertEqual(source._get_quota_project_id(), 'my-billing-project') + + def test_storage_source_gets_quota_from_options(self): + """Test that _CustomBigQueryStorageSource falls back to options.""" + options = PipelineOptions(['--quota_project_id=my-billing-project']) + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + pipeline_options=options) + self.assertEqual(source._get_quota_project_id(), 'my-billing-project') + + @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client') + def test_export_source_split_passes_quota_to_client(self, mock_bq_client): + """split() must pass the transform-level quota_project_id to the BigQuery + client, not only the --quota_project_id pipeline option.""" + mock_bq_client.side_effect = RuntimeError('stop') + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + quota_project_id='my-billing-project') + with self.assertRaises(RuntimeError): + list(source.split(desired_bundle_size=0)) + mock_bq_client.assert_called_once() + self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions) + self.assertEqual( + mock_bq_client.call_args.kwargs['quota_project_id'], + 'my-billing-project') + + @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client') + def test_storage_source_split_passes_quota_to_client(self, mock_bq_client): + """split() must pass the transform-level quota_project_id to the BigQuery + client, not only the --quota_project_id pipeline option.""" + mock_bq_client.side_effect = RuntimeError('stop') + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + quota_project_id='my-billing-project') + with self.assertRaises(RuntimeError): + list(source.split(desired_bundle_size=0)) + mock_bq_client.assert_called_once() + self.assertIsInstance(mock_bq_client.call_args.args[0], PipelineOptions) + self.assertEqual( + mock_bq_client.call_args.kwargs['quota_project_id'], + 'my-billing-project') + + @mock.patch.object(bigquery_tools, 'create_bq_storage_client') + @mock.patch.object(bigquery_tools.BigQueryWrapper, 'get_table') + @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client') + def test_storage_source_split_passes_quota_to_storage_client( + self, unused_mock_bq_client, mock_get_table, mock_create_storage_client): + """split() must create the Storage Read client with the quota project + and pass it on to each stream source.""" + mock_get_table.return_value = mock.Mock(numBytes=5) + stream = mock.Mock() + stream.name = 'projects/p/locations/l/sessions/s/streams/stream1' + mock_create_storage_client.return_value.create_read_session.return_value = ( + mock.Mock(streams=[stream])) + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + quota_project_id='my-billing-project') + + bundles = list(source.split(desired_bundle_size=1 << 20)) + + mock_create_storage_client.assert_called_once_with('my-billing-project') + self.assertEqual(len(bundles), 1) + self.assertEqual(bundles[0].source.quota_project_id, 'my-billing-project') + + @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client') + def test_export_source_estimate_size_passes_quota_to_client( + self, mock_bq_client): + """estimate_size() must attribute quota too, not just split().""" + mock_bq_client.side_effect = RuntimeError('stop') + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + quota_project_id='my-billing-project') + with self.assertRaises(RuntimeError): + source.estimate_size() + self.assertEqual( + mock_bq_client.call_args.kwargs['quota_project_id'], + 'my-billing-project') + + @mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client') + def test_storage_source_estimate_size_passes_quota_to_client( + self, mock_bq_client): + """estimate_size() must attribute quota too, not just split().""" + mock_bq_client.side_effect = RuntimeError('stop') + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + quota_project_id='my-billing-project') + with self.assertRaises(RuntimeError): + source.estimate_size() + self.assertEqual( + mock_bq_client.call_args.kwargs['quota_project_id'], + 'my-billing-project') + + @mock.patch( + 'apache_beam.io.gcp.bigquery.bigquery_schema_tools' + '.convert_to_usertype') + @mock.patch.object(bigquery_tools, 'BigQueryWrapper') + def test_output_type_schema_fetch_passes_quota_to_client( + self, mock_wrapper, mock_convert): + """The BEAM_ROW schema lookup is a BigQuery API call of this transform, so + it must be attributed to the quota project as well.""" + mock_convert.return_value = beam.Map(lambda x: x) + transform = ReadFromBigQuery( + table='project:dataset.table', + output_type='BEAM_ROW', + quota_project_id='my-billing-project') + transform._expand_output_type(mock.MagicMock()) + + self.assertEqual( + mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project') + + @mock.patch.object(bigquery_tools, 'BigQueryWrapper') + def test_direct_read_cleanup_passes_quota_to_client(self, mock_wrapper): + """The DIRECT_READ temp-dataset cleanup is a BigQuery API call of this + transform, so it must be attributed to the quota project as well.""" + pipeline_details = { + 'project_id': 'project', 'bigquery_dataset_labels': { + 'k': 'v' + } + } + with TestPipeline() as p: + side_input = beam.pvalue.AsList( + p | 'Details' >> beam.Create([pipeline_details])) + _ = ( + p | beam.Create([1]) + | _PassThroughThenCleanupTempDatasets( + side_input, quota_project_id='my-billing-project')) + + self.assertEqual( + mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project') + self.assertIsNotNone(mock_wrapper.call_args.kwargs['pipeline_options']) + + @mock.patch('apache_beam.io.gcp.bigquery._PassThroughThenCleanupTempDatasets') + def test_direct_read_wires_quota_project_into_cleanup(self, mock_cleanup): + mock_cleanup.return_value = beam.Map(lambda x: x) + transform = ReadFromBigQuery( + table='project:dataset.table', + method=ReadFromBigQuery.Method.DIRECT_READ, + quota_project_id='my-billing-project') + p = TestPipeline() + transform._expand_direct_read(beam.pvalue.PBegin(p)) + + self.assertEqual( + mock_cleanup.call_args.kwargs['quota_project_id'], 'my-billing-project') + + def test_quota_project_id_in_export_source_display_data(self): + """Test that quota_project_id appears in display data for export source.""" + source = beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + quota_project_id='my-billing-project') + items = {i.key: i.value for i in DisplayData.create_from(source).items} + self.assertEqual(items['quota_project_id'], 'my-billing-project') + + def test_quota_project_id_in_direct_read_source_display_data(self): + """Test that quota_project_id appears in display data for DIRECT_READ.""" + source = beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + quota_project_id='my-billing-project') + items = {i.key: i.value for i in DisplayData.create_from(source).items} + self.assertEqual(items['quota_project_id'], 'my-billing-project') + + def test_quota_project_id_absent_from_display_data_when_not_set(self): + """Test that quota_project_id is dropped from display data when not set.""" + options = PipelineOptions([]) + sources = [ + beam_bq._CustomBigQuerySource( + method=ReadFromBigQuery.Method.EXPORT, + table='project:dataset.table', + pipeline_options=options), + beam_bq._CustomBigQueryStorageSource( + method=ReadFromBigQuery.Method.DIRECT_READ, + table='project:dataset.table', + pipeline_options=options) + ] + for source in sources: + keys = {i.key for i in DisplayData.create_from(source).items} + self.assertNotIn('quota_project_id', keys) + + def test_stream_source_stores_quota_project_id(self): + """Test that quota_project_id is stored in + _CustomBigQueryStorageStreamSource.""" + stream_source = beam_bq._CustomBigQueryStorageStreamSource( + read_stream_name='projects/p/locations/l/sessions/s/streams/stream1', + use_native_datetime=True, + timeout=30.0, + quota_project_id='my-billing-project') + self.assertEqual(stream_source.quota_project_id, 'my-billing-project') + + def test_stream_source_quota_project_id_none_by_default(self): + """Test that quota_project_id is None by default in stream source.""" + stream_source = beam_bq._CustomBigQueryStorageStreamSource( + read_stream_name='projects/p/locations/l/sessions/s/streams/stream1') + self.assertIsNone(stream_source.quota_project_id) + + def test_stream_source_split_preserves_quota_project_id(self): + """Test that split() preserves quota_project_id.""" + stream_source = beam_bq._CustomBigQueryStorageStreamSource( + read_stream_name='projects/p/locations/l/sessions/s/streams/stream1', + use_native_datetime=True, + timeout=30.0, + quota_project_id='my-billing-project') + bundle = stream_source.split(desired_bundle_size=0) + self.assertEqual(bundle.source.quota_project_id, 'my-billing-project') + self.assertEqual(bundle.source.timeout, 30.0) + self.assertEqual(bundle.source.use_native_datetime, True) + + @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client') + def test_stream_source_read_arrow_uses_quota_project( + self, mock_create_client): + """Test that read_arrow() uses create_bq_storage_client + with quota_project_id.""" + mock_client = mock.MagicMock() + mock_create_client.return_value = mock_client + # Mock read_rows to return empty iterator + mock_client.read_rows.return_value.rows.return_value = iter([]) + + stream_source = beam_bq._CustomBigQueryStorageStreamSource( + read_stream_name='projects/p/locations/l/sessions/s/streams/stream1', + use_native_datetime=True, + quota_project_id='my-billing-project') + # Consume the iterator + list(stream_source.read_arrow()) + + mock_create_client.assert_called_once_with('my-billing-project') + + @mock.patch('apache_beam.io.gcp.bigquery_tools.create_bq_storage_client') + def test_stream_source_read_avro_uses_quota_project(self, mock_create_client): + """Test that read_avro() uses create_bq_storage_client + with quota_project_id.""" + mock_client = mock.MagicMock() + mock_create_client.return_value = mock_client + # Mock read_rows to return empty iterator + mock_client.read_rows.return_value = iter([]) + + stream_source = beam_bq._CustomBigQueryStorageStreamSource( + read_stream_name='projects/p/locations/l/sessions/s/streams/stream1', + use_native_datetime=False, + quota_project_id='my-billing-project') + # Consume the iterator + list(stream_source.read_avro()) + + mock_create_client.assert_called_once_with('my-billing-project') + + +@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') +class TestWriteToBigQueryQuotaProject(unittest.TestCase): + """Tests for quota_project_id in the BigQuery streaming insert path.""" + @mock.patch.object(bigquery_tools, 'BigQueryWrapper') + def test_write_fn_start_bundle_passes_quota_project(self, mock_wrapper): + fn = beam_bq.BigQueryWriteFn( + batch_size=1, quota_project_id='my-billing-project') + fn.start_bundle() + self.assertEqual( + mock_wrapper.call_args.kwargs['quota_project_id'], 'my-billing-project') + + def test_stream_to_bigquery_resolves_quota_project_from_options(self): + options = PipelineOptions(['--quota_project_id=my-billing-project']) + p = TestPipeline(options=options) + with mock.patch.object( + beam_bq.BigQueryWriteFn, + '__init__', + autospec=True, + side_effect=beam_bq.BigQueryWriteFn.__init__) as mock_write_fn: + _ = ( + p + | beam.Create([]) + | WriteToBigQuery( + 'project:dataset.table', + schema='a:STRING', + method=WriteToBigQuery.Method.STREAMING_INSERTS)) + self.assertEqual( + mock_write_fn.call_args.kwargs['quota_project_id'], + 'my-billing-project') + + @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQuerySink(unittest.TestCase): def test_table_spec_display_data(self): @@ -920,10 +1264,14 @@ def _cleanup_files(self): os.remove('insert_calls2') def setUp(self): + self.env_patch = mock.patch.dict( + os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}) + self.env_patch.start() self._cleanup_files() def tearDown(self): self._cleanup_files() + self.env_patch.stop() def test_noop_schema_parsing(self): expected_table_schema = None @@ -1481,6 +1829,13 @@ def test_copy_load_job_exception(self, exception_type, error_message): HttpError is None or exceptions is None, 'GCP dependencies are not installed') class BigQueryStreamingInsertsErrorHandling(unittest.TestCase): + def setUp(self): + self.env_patch = mock.patch.dict( + os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}) + self.env_patch.start() + + def tearDown(self): + self.env_patch.stop() # Running tests with a variety of exceptions from https://googleapis.dev # /python/google-api-core/latest/_modules/google/api_core/exceptions.html. diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index cc1a16a659fe..d1412d0f855e 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -36,6 +36,7 @@ import re import sys import time +import urllib.parse import uuid from json.decoder import JSONDecodeError from typing import Optional @@ -59,6 +60,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 GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.transforms import DoFn from apache_beam.typehints.row_type import RowTypeConstraint @@ -80,6 +82,11 @@ gcp_bigquery = None pass +try: + from google.cloud import bigquery_storage_v1 as bq_storage +except ImportError: + bq_storage = None + try: from orjson import dumps as fast_json_dumps from orjson import loads as fast_json_loads @@ -419,6 +426,58 @@ def _build_dataset_encryption_config(kms_key): return bigquery.EncryptionConfiguration(kmsKeyName=kms_key) +def _quota_project_id_from_options(pipeline_options): + """Returns --quota_project_id from pipeline options, if any.""" + if pipeline_options is None: + return None + return pipeline_options.view_as(GoogleCloudOptions).quota_project_id + + +def create_bq_storage_client(quota_project_id=None): + """Create a BigQueryReadClient with optional quota project. + + Args: + quota_project_id: Optional GCP project ID to use for quota and billing. + The client applies it to the credentials it resolves itself. + + Returns: + A BigQueryReadClient instance. + """ + if not quota_project_id: + return bq_storage.BigQueryReadClient() + return bq_storage.BigQueryReadClient( + client_options={'quota_project_id': quota_project_id}) + + +class _HttpWithHeaders(object): + """Wraps an httplib2.Http to add fixed headers to requests to given hosts. + + Sits below the credentials layer, so the headers are applied on every + attempt and take precedence over headers set by the credentials. The + credentials refresh their token through the same Http, so requests to other + hosts (token endpoint, metadata server, IAM Credentials API) are passed + through unchanged. + """ + def __init__(self, http, headers, hosts): + object.__setattr__(self, '_http', http) + object.__setattr__(self, '_headers', headers) + object.__setattr__(self, '_hosts', frozenset(hosts)) + + def request(self, uri, method='GET', body=None, headers=None, **kwargs): + if urllib.parse.urlsplit(uri).hostname in self._hosts: + headers = dict(headers or {}) + headers.update(self._headers) + return self._http.request(uri, method, body=body, headers=headers, **kwargs) + + def __getattr__(self, name): + if name == '_http': # object created without __init__ (copy/unpickle) + raise AttributeError(name) + return getattr(self._http, name) + + def __setattr__(self, name, value): + setattr(self._http, name, value) + + class BigQueryWrapper(object): """BigQuery client wrapper with utilities for querying. @@ -427,8 +486,10 @@ class BigQueryWrapper(object): In addition, it offers various functions used both in sources and sinks (e.g., find and create tables, query a table, etc.). - Note that client parameter in constructor is only for testing purposes and - should not be used in production code. + Note that the client and gcp_client parameters in the constructor are only + for testing purposes and should not be used in production code. Production + code should pass pipeline_options (and optionally quota_project_id) so that + every client is built here with the right quota attribution. """ # If updating following names, also update the corresponding pydocs in @@ -438,16 +499,29 @@ class BigQueryWrapper(object): HISTOGRAM_METRIC_LOGGER = MetricLogger() - def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): - self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions()) - self.gcp_bq_client = client or gcp_bigquery.Client( - client_info=ClientInfo( - user_agent="apache-beam-%s" % apache_beam.__version__)) + def __init__( + self, + client=None, + temp_dataset_id=None, + temp_table_ref=None, + quota_project_id=None, + *, + gcp_client=None, + pipeline_options=None): + self.quota_project_id = ( + quota_project_id or _quota_project_id_from_options(pipeline_options)) + self.client = client or BigQueryWrapper._bigquery_client( + pipeline_options or PipelineOptions(), + quota_project_id=self.quota_project_id) + # Test-only overrides. `client` also stands in for the google-cloud-bigquery + # client, which several tests rely on; pass `gcp_client` as well as `client` + # to override only that one. + self._gcp_bq_client = gcp_client or client self._unique_row_id = 0 # For testing scenarios where we pass in a client we do not want a # randomized prefix for row IDs. - self._row_id_prefix = '' if client else uuid.uuid4() + self._row_id_prefix = '' if (client or gcp_client) else uuid.uuid4() self._latency_histogram_metric = Metrics.histogram( self.__class__, 'latency_histogram_ms', @@ -474,6 +548,17 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): self.created_temp_dataset = False + @property + def gcp_bq_client(self): + """The google-cloud-bigquery client, created on first use. + + Only the streaming insert path needs it, so it is not built eagerly. + """ + if self._gcp_bq_client is None: + self._gcp_bq_client = BigQueryWrapper._gcp_bigquery_client( + quota_project_id=self.quota_project_id) + return self._gcp_bq_client + @property def unique_row_id(self): """Returns a unique row ID (str) used to avoid multiple insertions. @@ -1490,19 +1575,70 @@ def convert_row_to_dict(self, row, schema): @staticmethod def from_pipeline_options(pipeline_options: PipelineOptions): - return BigQueryWrapper( - client=BigQueryWrapper._bigquery_client(pipeline_options)) + """Create a BigQueryWrapper from pipeline options. + + Args: + pipeline_options: Pipeline options containing GCP configuration. + The quota_project_id is read from GoogleCloudOptions if set. + """ + return BigQueryWrapper(pipeline_options=pipeline_options) @staticmethod - def _bigquery_client(pipeline_options: PipelineOptions): + def _bigquery_client( + pipeline_options: PipelineOptions, quota_project_id: str = None): + """Create a BigQuery API client from pipeline options. + + Args: + pipeline_options: Pipeline options for credentials. + quota_project_id: Optional quota project ID. If not provided, will be + extracted from pipeline_options. + """ + credentials = auth.get_service_credentials(pipeline_options) + # 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) + http = get_new_http() + if quota_project_id: + if credentials is None: + # Ignoring the request would silently bill a different project. + raise ValueError( + 'quota_project_id was set to %r, but no credentials were found to ' + 'apply it to.' % quota_project_id) + # Send the quota project as a request header rather than deriving new + # credentials: credentials.with_quota_project() returns a copy without + # the cached token, and this keeps the shared credentials untouched. + # The header is set below the credentials layer so that it is applied on + # every attempt and wins over a quota project the credentials carry. + http = _HttpWithHeaders( + http, {'x-goog-user-project': quota_project_id}, + hosts=[ + urllib.parse.urlsplit(url).hostname for url in ( + bigquery.BigqueryV2.BASE_URL, + bigquery.BigqueryV2.MTLS_BASE_URL) + ]) return bigquery.BigqueryV2( - http=get_new_http(), - credentials=auth.get_service_credentials(pipeline_options), + http=http, + credentials=credentials, response_encoding='utf8', additional_http_headers={ "user-agent": "apache-beam-%s" % apache_beam.__version__ }) + @staticmethod + def _gcp_bigquery_client(quota_project_id: str = 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. + """ + client_options = None + if quota_project_id: + client_options = {'quota_project_id': quota_project_id} + return gcp_bigquery.Client( + client_options=client_options, + client_info=ClientInfo( + user_agent="apache-beam-%s" % apache_beam.__version__)) + class RowAsDictJsonCoder(coders.Coder): """A coder for a table row (represented as a dict) to/from a JSON string. 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 9f8b78ee5c3b..de4f3e4066af 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -23,6 +23,7 @@ import json import logging import math +import os import re import unittest from typing import Optional @@ -35,10 +36,12 @@ from parameterized import parameterized import apache_beam as beam +from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import AvroRowWriter from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes +from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.bigquery_tools import JsonRowWriter from apache_beam.io.gcp.bigquery_tools import RowAsDictJsonCoder from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict @@ -299,14 +302,16 @@ def test_delete_dataset_retries_for_timeouts(self, patched_time_sleep): @mock.patch('google.cloud._http.JSONConnection.http') def test_user_agent_insert_all( self, http_mock, patched_skip_get_credentials, patched_sleep): - wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper() - try: - wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None) - except: # pylint: disable=bare-except - # Ignore errors. The errors come from the fact that we did not mock - # the response from the API, so the overall insert_all_rows call fails - # soon after the BQ API is called. - pass + # Set GOOGLE_CLOUD_PROJECT to ensure Client creation succeeds in test env + with mock.patch.dict(os.environ, {'GOOGLE_CLOUD_PROJECT': 'test-project'}): + wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper() + try: + wrapper._insert_all_rows('p', 'd', 't', [{'name': 'any'}], None) + except: # pylint: disable=bare-except + # Ignore errors. The errors come from the fact that we did not mock + # the response from the API, so the overall insert_all_rows call fails + # soon after the BQ API is called. + pass call = http_mock.request.mock_calls[-2] self.assertIn('apache-beam-', call[2]['headers']['User-Agent']) @@ -1199,8 +1204,6 @@ def test_geography_in_bigquery_type_mapping(self): def test_geography_field_conversion(self): """Test that GEOGRAPHY fields are converted correctly.""" - from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper - # Create a mock field with GEOGRAPHY type field = bigquery.TableFieldSchema() field.type = 'GEOGRAPHY' @@ -1322,8 +1325,6 @@ def test_geography_json_encoding(self): def test_geography_with_special_characters(self): """Test GEOGRAPHY values with special characters and geometries.""" - from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper - field = bigquery.TableFieldSchema() field.type = 'GEOGRAPHY' field.name = 'complex_geo' @@ -1502,6 +1503,213 @@ def test_type_overrides_json_to_dict(self): self.assertEqual(typehints_dict, [("data", Optional[dict])]) +@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') +class TestBigQueryWrapperQuotaProject(unittest.TestCase): + """Tests for quota_project_id in BigQueryWrapper.""" + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client') + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client') + def test_quota_project_id_stored(self, mock_gcp_client, mock_bq_client): + """Test that quota_project_id is stored in BigQueryWrapper.""" + mock_bq_client.return_value = mock.Mock() + mock_gcp_client.return_value = mock.Mock() + + wrapper = BigQueryWrapper(quota_project_id='my-billing-project') + self.assertEqual(wrapper.quota_project_id, 'my-billing-project') + + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client') + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client') + def test_from_pipeline_options_reads_quota_from_options( + self, mock_gcp_client, mock_bq_client): + """Test from_pipeline_options reads quota_project_id from + GoogleCloudOptions.""" + from apache_beam.options.pipeline_options import PipelineOptions + + mock_bq_client.return_value = mock.Mock() + mock_gcp_client.return_value = mock.Mock() + + options = PipelineOptions(['--quota_project_id=my-billing-project']) + wrapper = BigQueryWrapper.from_pipeline_options(options) + + self.assertEqual(wrapper.quota_project_id, 'my-billing-project') + + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._bigquery_client') + @mock.patch( + 'apache_beam.io.gcp.bigquery_tools.BigQueryWrapper._gcp_bigquery_client') + def test_from_pipeline_options_none_when_not_set( + self, mock_gcp_client, mock_bq_client): + """Test from_pipeline_options returns None when quota_project_id not set.""" + from apache_beam.options.pipeline_options import PipelineOptions + + mock_bq_client.return_value = mock.Mock() + mock_gcp_client.return_value = mock.Mock() + + options = PipelineOptions([]) + wrapper = BigQueryWrapper.from_pipeline_options(options) + + self.assertIsNone(wrapper.quota_project_id) + + def test_http_with_headers_sets_headers_and_delegates(self): + from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders + inner = mock.MagicMock() + inner.request.return_value = ('response', b'content') + inner.redirect_codes = {301, 302} + http = _HttpWithHeaders( + inner, {'x-goog-user-project': 'p'}, hosts=['example.com']) + + result = http.request( + 'https://example.com', + method='GET', + headers={ + 'x-goog-user-project': 'other', 'accept': 'json' + }, + redirections=3) + + self.assertEqual(result, ('response', b'content')) + inner.request.assert_called_once_with( + 'https://example.com', + 'GET', + body=None, + headers={ + 'x-goog-user-project': 'p', 'accept': 'json' + }, + redirections=3) + # Attribute reads and writes reach the wrapped http. + self.assertEqual(http.redirect_codes, {301, 302}) + http.redirect_codes = {301} + self.assertEqual(inner.redirect_codes, {301}) + + def test_http_with_headers_survives_copy(self): + import copy + + from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders + inner = mock.MagicMock() + http = copy.copy( + _HttpWithHeaders( + inner, {'x-goog-user-project': 'p'}, hosts=['example.com'])) + + http.request('https://example.com') + inner.request.assert_called_once_with( + 'https://example.com', + 'GET', + body=None, + headers={'x-goog-user-project': 'p'}) + + def test_http_with_headers_skips_other_hosts(self): + from apache_beam.io.gcp.bigquery_tools import _HttpWithHeaders + inner = mock.MagicMock() + http = _HttpWithHeaders( + inner, {'x-goog-user-project': 'p'}, hosts=['example.com']) + + http.request('https://oauth2.googleapis.com/token', headers={'a': 'b'}) + inner.request.assert_called_once_with( + 'https://oauth2.googleapis.com/token', + 'GET', + body=None, + headers={'a': 'b'}) + + @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage') + def test_create_bq_storage_client_with_quota_project(self, mock_bq_storage): + """The quota project is passed as a client option; the client applies it + to the credentials it resolves itself.""" + bigquery_tools.create_bq_storage_client('my-billing-project') + mock_bq_storage.BigQueryReadClient.assert_called_once_with( + client_options={'quota_project_id': 'my-billing-project'}) + + @mock.patch('apache_beam.io.gcp.bigquery_tools.bq_storage') + def test_create_bq_storage_client_without_quota_project( + self, mock_bq_storage): + """Without a quota project the client is created with defaults.""" + bigquery_tools.create_bq_storage_client(None) + mock_bq_storage.BigQueryReadClient.assert_called_once_with() + + @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials') + def test_bigquery_client_raises_without_credentials( + self, mock_get_credentials): + """An explicit quota project must fail loudly rather than fall back to + apitools' own credential discovery, which would bill another project.""" + from apache_beam.options.pipeline_options import PipelineOptions + mock_get_credentials.return_value = None + + with self.assertRaisesRegex(ValueError, 'no credentials were found'): + BigQueryWrapper._bigquery_client( + PipelineOptions(), quota_project_id='requested-project') + + @mock.patch('apache_beam.io.gcp.bigquery_tools.get_new_http') + @mock.patch('apache_beam.io.gcp.bigquery_tools.auth.get_service_credentials') + def test_bigquery_client_quota_project_wins_over_credentials( + self, mock_get_credentials, mock_get_new_http): + """The requested quota project must reach every BigQuery request, + including the retry after a 401 refresh, even when the credentials carry + their own, but not the token refresh requests.""" + import httplib2 + from google.auth import credentials as ga_credentials + + from apache_beam.internal.gcp import auth + from apache_beam.options.pipeline_options import PipelineOptions + + class CredentialsWithQuotaProject(ga_credentials.Credentials): + def __init__(self): + super().__init__() + self._quota_project_id = 'adc-project' + + def refresh(self, request): + # Token refreshes go through the same Http as the API requests. + request('https://oauth2.googleapis.com/token', method='POST') + self.token = 'token' + + seen_headers = [] + token_request_headers = [] + statuses = iter([401, 200]) + + class FakeHttp(object): + connections = {} + redirect_codes = set() + + def request(self, uri, method='GET', body=None, headers=None, **kwargs): + if uri.startswith('https://oauth2.googleapis.com/'): + token_request_headers.append(dict(headers or {})) + return httplib2.Response({'status': 200}), b'{}' + seen_headers.append(dict(headers)) + response = httplib2.Response({ + 'status': next(statuses), 'content-type': 'application/json' + }) + return response, b'{}' + + mock_get_new_http.return_value = FakeHttp() + mock_get_credentials.return_value = auth._ApitoolsCredentialsAdapter( + CredentialsWithQuotaProject()) + + client = BigQueryWrapper._bigquery_client( + PipelineOptions(), quota_project_id='requested-project') + client.projects.List(bigquery.BigqueryProjectsListRequest()) + + self.assertEqual(len(seen_headers), 2) + for headers in seen_headers: + self.assertEqual(headers['x-goog-user-project'], 'requested-project') + self.assertEqual(headers['authorization'], 'Bearer token') + # The quota project must not be sent to the token endpoint. + self.assertTrue(token_request_headers) + for headers in token_request_headers: + self.assertNotIn('x-goog-user-project', headers) + + @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client') + def test_gcp_bigquery_client_passes_quota_project(self, mock_client): + BigQueryWrapper._gcp_bigquery_client(quota_project_id='my-billing-project') + self.assertEqual( + mock_client.call_args.kwargs['client_options'], + {'quota_project_id': 'my-billing-project'}) + + @mock.patch('apache_beam.io.gcp.bigquery_tools.gcp_bigquery.Client') + def test_gcp_bigquery_client_no_quota_project_by_default(self, mock_client): + BigQueryWrapper._gcp_bigquery_client() + self.assertIsNone(mock_client.call_args.kwargs['client_options']) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 08ba34dabee4..f0406e2505bc 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1197,6 +1197,17 @@ def _add_argparse_args(cls, parser): action='store_true', help='Throttling counter in GcsIO is enabled by default. Set ' '--no_gcsio_throttling_counter to avoid it.') + parser.add_argument( + '--quota_project_id', + default=None, + help='GCP project ID used for quota and billing attribution of Google ' + 'Cloud API requests (sets the x-goog-user-project header), if ' + 'different from the project the data resides in. Currently honored ' + 'by BigQuery IO transforms (ReadFromBigQuery, ReadAllFromBigQuery, ' + 'WriteToBigQuery with FILE_LOADS or STREAMING_INSERTS, ' + 'ReadBigQueryChangeHistory); STORAGE_WRITE_API (Java-backed) and ' + 'other GCP IOs do not use it yet. The credentials used must have the ' + 'serviceusage.services.use permission on this project.') parser.add_argument( '--enable_gcsio_blob_generation', default=False, From 71d345215f286fdabe6f6f38e9b55bc7c5e9b69c Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:27:29 +0300 Subject: [PATCH 2/4] [Java] Add bigQueryQuotaProjectId pipeline option to BigQueryIO Add BigQueryOptions.bigQueryQuotaProjectId to attribute BigQuery API quota and billing to a project other than the one the data resides in. For the REST client, the x-goog-user-project header is set by an execute interceptor, so it is re-applied on every attempt, including after HttpCredentialsAdapter re-initializes headers on a 401 refresh. The Storage Read and Write clients set it through setQuotaProjectId. Generated-By: Claude Opus 5.5 --- .../sdk/io/gcp/bigquery/BigQueryOptions.java | 10 ++++ .../io/gcp/bigquery/BigQueryServicesImpl.java | 23 +++++++- .../bigquery/BigQueryServicesImplTest.java | 54 +++++++++++++++++++ 3 files changed, 86 insertions(+), 1 deletion(-) diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java index 1826f4a1fd21..64382d64bb51 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryOptions.java @@ -197,6 +197,16 @@ public interface BigQueryOptions void setBigQueryProject(String value); + @Description( + "GCP project ID used for quota and billing attribution of BigQuery API requests " + + "(sets the X-Goog-User-Project header), if different from the project the data " + + "resides in. If unspecified, the project associated with the credentials is used. " + + "The credentials used must have the serviceusage.services.use permission on this " + + "project.") + String getBigQueryQuotaProjectId(); + + void setBigQueryQuotaProjectId(String value); + @Description("Maximum (best effort) size of a single append to the storage API.") @Default.Integer(2 * 1024 * 1024) Integer getStorageApiAppendThresholdBytes(); diff --git a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java index 7a924c6aa021..a921be3700ec 100644 --- a/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java +++ b/sdks/java/io/google-cloud-platform/src/main/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImpl.java @@ -1800,7 +1800,8 @@ private static boolean nextBackOff(Sleeper sleeper, BackOff backoff) throws Inte } /** Returns a BigQuery client builder using the specified {@link BigQueryOptions}. */ - private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) { + @VisibleForTesting + static Bigquery.Builder newBigQueryClient(BigQueryOptions options) { // Do not log 404. It clutters the output and is possibly even required by the // caller. RetryHttpRequestInitializer httpRequestInitializer = @@ -1818,6 +1819,18 @@ private static Bigquery.Builder newBigQueryClient(BigQueryOptions options) { initBuilder.add(new LatencyRecordingHttpRequestInitializer(API_METRIC_LABEL)); initBuilder.add(httpRequestInitializer); + // Set the quota project as a request header instead of deriving credentials: a derived + // credential inherits the shared credential's cached request metadata, which lacks + // x-goog-user-project until the next token refresh. Applied by an execute interceptor so it + // is re-applied on every attempt; HttpCredentialsAdapter re-initializes the request headers + // from the credential after a 401 refresh, which would otherwise override it. + @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId(); + if (!Strings.isNullOrEmpty(quotaProjectId)) { + initBuilder.add( + request -> + request.setInterceptor( + r -> r.getHeaders().set("x-goog-user-project", quotaProjectId))); + } HttpRequestInitializer chainInitializer = new ChainingHttpRequestInitializer( Iterables.toArray(initBuilder.build(), HttpRequestInitializer.class)); @@ -1848,6 +1861,10 @@ private static BigQueryWriteClient newBigQueryWriteClient(BigQueryOptions option if (!Strings.isNullOrEmpty(endpoint)) { builder.setEndpoint(trimSchemaIfNecessary(endpoint)); } + @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId(); + if (!Strings.isNullOrEmpty(quotaProjectId)) { + builder.setQuotaProjectId(quotaProjectId); + } return BigQueryWriteClient.create( builder .setCredentialsProvider(() -> options.as(GcpOptions.class).getGcpCredential()) @@ -1973,6 +1990,10 @@ public void onRetryAttempt(Status status, Metadata metadata) { if (!Strings.isNullOrEmpty(endpoint)) { settingsBuilder.setEndpoint(trimSchemaIfNecessary(endpoint)); } + @Nullable String quotaProjectId = options.getBigQueryQuotaProjectId(); + if (!Strings.isNullOrEmpty(quotaProjectId)) { + settingsBuilder.setQuotaProjectId(quotaProjectId); + } UnaryCallSettings.Builder createReadSessionSettings = settingsBuilder.getStubSettingsBuilder().createReadSessionSettings(); diff --git a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java index 90b2e1fb8723..642d35b43905 100644 --- a/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java +++ b/sdks/java/io/google-cloud-platform/src/test/java/org/apache/beam/sdk/io/gcp/bigquery/BigQueryServicesImplTest.java @@ -41,6 +41,8 @@ import com.google.api.client.googleapis.json.GoogleJsonError.ErrorInfo; import com.google.api.client.googleapis.json.GoogleJsonErrorContainer; import com.google.api.client.googleapis.json.GoogleJsonResponseException; +import com.google.api.client.http.GenericUrl; +import com.google.api.client.http.HttpRequest; import com.google.api.client.http.HttpResponseException; import com.google.api.client.http.LowLevelHttpResponse; import com.google.api.client.json.GenericJson; @@ -2311,6 +2313,58 @@ public RetryInfo parseBytes(byte[] serialized) { assertEquals(123456, (long) container.getCounter(metricName).getCumulative()); } + @Test + public void testQuotaProjectIdOverrides() throws IOException { + BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class); + options.setBigQueryQuotaProjectId("my-quota-project"); + + assertEquals( + "my-quota-project", + new BigQueryServicesImpl.StorageClientImpl(options) + .getClient() + .getSettings() + .getQuotaProjectId()); + assertEquals( + "my-quota-project", + new BigQueryServicesImpl.WriteStreamServiceImpl(options) + .getClient() + .getSettings() + .getQuotaProjectId()); + } + + @Test + public void testQuotaProjectIdSetsRequestHeader() throws IOException { + BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class); + options.setGcpCredential(null); + GenericUrl url = new GenericUrl("https://bigquery.googleapis.com/bigquery/v2/projects"); + + HttpRequest request = + BigQueryServicesImpl.newBigQueryClient(options) + .build() + .getRequestFactory() + .buildGetRequest(url); + request.getInterceptor().intercept(request); + assertNull(request.getHeaders().getFirstHeaderStringValue("x-goog-user-project")); + + options.setBigQueryQuotaProjectId("my-quota-project"); + request = + BigQueryServicesImpl.newBigQueryClient(options) + .build() + .getRequestFactory() + .buildGetRequest(url); + // The interceptor runs before every attempt, including retries. + request.getInterceptor().intercept(request); + assertEquals( + "my-quota-project", request.getHeaders().getFirstHeaderStringValue("x-goog-user-project")); + + // Simulate HttpCredentialsAdapter re-initializing headers from a credential that carries + // its own quota project after a 401 refresh: the configured value must win on the retry. + request.getHeaders().set("x-goog-user-project", "credential-project"); + request.getInterceptor().intercept(request); + assertEquals( + "my-quota-project", request.getHeaders().getFirstHeaderStringValue("x-goog-user-project")); + } + @Test public void testEndpointOverrides() throws IOException { BigQueryOptions options = PipelineOptionsFactory.create().as(BigQueryOptions.class); From 39c54325110ce06af42d125c2751904668c27795 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:27:29 +0300 Subject: [PATCH 3/4] [Go] Add WithQuotaProject option to bigqueryio Add a WithQuotaProject query option that attributes BigQuery API quota and billing to a project other than the one the data resides in, via option.WithQuotaProject on the BigQuery client. Read now accepts query options too. Only WithQuotaProject applies to it; UseStandardSQL and WithQueryParameters panic, because Read generates a legacy SQL statement without parameters. Generated-By: Claude Opus 5.5 --- sdks/go/pkg/beam/io/bigqueryio/bigquery.go | 54 +++++++- .../pkg/beam/io/bigqueryio/bigquery_test.go | 125 ++++++++++++++++++ 2 files changed, 173 insertions(+), 6 deletions(-) diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go index bbe951969b92..98818ccf1af1 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery.go @@ -34,6 +34,7 @@ import ( bq "google.golang.org/api/bigquery/v2" "google.golang.org/api/googleapi" "google.golang.org/api/iterator" + "google.golang.org/api/option" ) // writeSizeLimit is the maximum number of rows allowed by BQ in a write. @@ -88,14 +89,22 @@ func NewQualifiedTableName(s string) (QualifiedTableName, error) { // Read reads all rows from the given table. The table must have a schema // compatible with the given type, t, and Read returns a PCollection. If the // table has more rows than t, then Read is implicitly a projection. -func Read(s beam.Scope, project, table string, t reflect.Type) beam.PCollection { +// Of the options, only WithQuotaProject applies to Read; UseStandardSQL and +// WithQueryParameters panic, because the generated statement is legacy SQL and +// takes no parameters. +func Read(s beam.Scope, project, table string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection { mustParseTable(table) s = s.Scope("bigquery.Read") + qo := applyQueryOptions(options) + if qo.UseStandardSQL || len(qo.parameters) > 0 { + panic("bigqueryio.Read: UseStandardSQL and WithQueryParameters are only supported by Query") + } + stmt := constructSelectStatement(t, bigQueryTag, table) - return query(s, project, stmt, t) + return queryWithOptions(s, project, stmt, t, qo) } func constructSelectStatement(t reflect.Type, tagKey string, table string) string { @@ -114,6 +123,9 @@ func constructSelectStatement(t reflect.Type, tagKey string, table string) strin type QueryOptions struct { // UseStandardSQL enables BigQuery's Standard SQL dialect when executing a query. UseStandardSQL bool + // QuotaProject is the GCP project ID used for quota and billing attribution + // of the BigQuery API calls, if different from the project the data resides in. + QuotaProject string // Parameters are the query parameters for parameterized queries. // In the current implementation, user-defines types are not supported in Value field. // Use *bigquery.QueryParameterValue to build STRUCT/ARRAY parameters @@ -129,6 +141,17 @@ func UseStandardSQL() func(qo *QueryOptions) error { } } +// WithQuotaProject sets the GCP project ID used for quota and billing +// attribution of the BigQuery API calls, if different from the project the +// data resides in. The credentials used must have the +// serviceusage.services.use permission on that project. +func WithQuotaProject(project string) func(qo *QueryOptions) error { + return func(qo *QueryOptions) error { + qo.QuotaProject = project + return nil + } +} + // WithQueryParameters sets the query parameters for parameterized queries. // See QueryOptions.parameters for the list of supported Value types. func WithQueryParameters(params ...bigquery.QueryParameter) func(qo *QueryOptions) error { @@ -138,6 +161,15 @@ func WithQueryParameters(params ...bigquery.QueryParameter) func(qo *QueryOption } } +// clientOptions returns the BigQuery client options implied by qo. +func clientOptions(qo QueryOptions) []option.ClientOption { + var opts []option.ClientOption + if qo.QuotaProject != "" { + opts = append(opts, option.WithQuotaProject(qo.QuotaProject)) + } + return opts +} + // Query executes a query. The output must have a schema compatible with the given // type, t. It returns a PCollection. func Query(s beam.Scope, project, q string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection { @@ -145,15 +177,22 @@ func Query(s beam.Scope, project, q string, t reflect.Type, options ...func(*Que return query(s, project, q, t, options...) } -func query(s beam.Scope, project, query string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection { - mustInferSchema(t) - +func applyQueryOptions(options []func(*QueryOptions) error) QueryOptions { queryOptions := QueryOptions{} for _, opt := range options { if err := opt(&queryOptions); err != nil { panic(err) } } + return queryOptions +} + +func query(s beam.Scope, project, query string, t reflect.Type, options ...func(*QueryOptions) error) beam.PCollection { + return queryWithOptions(s, project, query, t, applyQueryOptions(options)) +} + +func queryWithOptions(s beam.Scope, project, query string, t reflect.Type, queryOptions QueryOptions) beam.PCollection { + mustInferSchema(t) imp := beam.Impulse(s) queryParameters, err := encodeQueryParameters(queryOptions.parameters) @@ -181,8 +220,11 @@ type queryFn struct { Options QueryOptions `json:"options"` } +// newClient creates the BigQuery client used by queryFn. Tests replace it. +var newClient = bigquery.NewClient + func (f *queryFn) ProcessElement(ctx context.Context, _ []byte, emit func(beam.X)) error { - client, err := bigquery.NewClient(ctx, f.Project) + client, err := newClient(ctx, f.Project, clientOptions(f.Options)...) if err != nil { return err } diff --git a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go index 9f83eeb847f7..4cf0c46b9f62 100644 --- a/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go +++ b/sdks/go/pkg/beam/io/bigqueryio/bigquery_test.go @@ -16,11 +16,16 @@ package bigqueryio import ( + "context" + "errors" + "fmt" "reflect" + "strings" "testing" "cloud.google.com/go/bigquery" "github.com/apache/beam/sdks/v2/go/pkg/beam" + "google.golang.org/api/option" ) func TestNewQualifiedTableName(t *testing.T) { @@ -149,6 +154,88 @@ func Test_mustInferSchema(t *testing.T) { } } +func TestWithQuotaProject(t *testing.T) { + qo := QueryOptions{} + if err := WithQuotaProject("quota-project")(&qo); err != nil { + t.Fatalf("WithQuotaProject() returned err: %v", err) + } + if got, want := qo.QuotaProject, "quota-project"; got != want { + t.Errorf("qo.QuotaProject = %q, want %q", got, want) + } +} + +func TestClientOptions(t *testing.T) { + if got := clientOptions(QueryOptions{}); len(got) != 0 { + t.Errorf("clientOptions(no quota) = %v, want none", got) + } + got := clientOptions(QueryOptions{QuotaProject: "quota-project"}) + if want := []option.ClientOption{option.WithQuotaProject("quota-project")}; !reflect.DeepEqual(got, want) { + t.Errorf("clientOptions(with quota) = %v, want %v", got, want) + } +} + +func TestQueryFn_passesQuotaProjectToClient(t *testing.T) { + var gotProject string + var gotOpts []option.ClientOption + errStop := errors.New("stop") + defer func(orig func(context.Context, string, ...option.ClientOption) (*bigquery.Client, error)) { + newClient = orig + }(newClient) + newClient = func(_ context.Context, project string, opts ...option.ClientOption) (*bigquery.Client, error) { + gotProject, gotOpts = project, opts + return nil, errStop + } + + fn := &queryFn{Project: "project", Options: QueryOptions{QuotaProject: "quota-project"}} + if err := fn.ProcessElement(context.Background(), nil, nil); !errors.Is(err, errStop) { + t.Fatalf("ProcessElement() err = %v, want %v", err, errStop) + } + if gotProject != "project" { + t.Errorf("client project = %q, want %q", gotProject, "project") + } + if want := []option.ClientOption{option.WithQuotaProject("quota-project")}; !reflect.DeepEqual(gotOpts, want) { + t.Errorf("client options = %v, want %v", gotOpts, want) + } +} + +// queryFnOf returns the queryFn in pipeline p. +func queryFnOf(t *testing.T, p *beam.Pipeline) *queryFn { + t.Helper() + edges, _, err := p.Build() + if err != nil { + t.Fatalf("Build() failed: %v", err) + } + for _, e := range edges { + if e.DoFn != nil { + if fn, ok := e.DoFn.Recv.(*queryFn); ok { + return fn + } + } + } + t.Fatal("no queryFn in pipeline") + return nil +} + +func TestReadAndQuery_passQuotaProjectToQueryFn(t *testing.T) { + type quotaSchema struct { + Name bigquery.NullString `bigquery:"name"` + } + beam.RegisterType(reflect.TypeOf((*quotaSchema)(nil)).Elem()) + typ := reflect.TypeOf(quotaSchema{}) + + p, s := beam.NewPipelineWithRoot() + Read(s, "project", "project:dataset.table", typ, WithQuotaProject("quota-project")) + if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" { + t.Errorf("Read: queryFn.Options.QuotaProject = %q, want %q", got, "quota-project") + } + + p, s = beam.NewPipelineWithRoot() + Query(s, "project", "SELECT 1", typ, WithQuotaProject("quota-project")) + if got := queryFnOf(t, p).Options.QuotaProject; got != "quota-project" { + t.Errorf("Query: queryFn.Options.QuotaProject = %q, want %q", got, "quota-project") + } +} + func TestWithQueryParameters(t *testing.T) { t.Run("WithQueryParameters sets parameters correctly", func(t *testing.T) { params := []bigquery.QueryParameter{ @@ -183,3 +270,41 @@ func TestWithQueryParameters(t *testing.T) { } }) } + +func TestRead_queryOnlyOptions(t *testing.T) { + type readSchema struct { + Name bigquery.NullString `bigquery:"name"` + } + beam.RegisterType(reflect.TypeOf((*readSchema)(nil)).Elem()) + typ := reflect.TypeOf(readSchema{}) + + mustPanic := func(t *testing.T, opt func(*QueryOptions) error) { + t.Helper() + defer func() { + r := recover() + if r == nil { + t.Errorf("Read() with query-only option does not panic") + } else if !strings.Contains(fmt.Sprint(r), "only supported by Query") { + t.Errorf("Read() panicked with unexpected value: %v", r) + } + }() + _, s := beam.NewPipelineWithRoot() + Read(s, "project", "project:dataset.table", typ, opt) + } + + t.Run("UseStandardSQL panics", func(t *testing.T) { + mustPanic(t, UseStandardSQL()) + }) + t.Run("WithQueryParameters panics", func(t *testing.T) { + mustPanic(t, WithQueryParameters(bigquery.QueryParameter{Name: "p", Value: 1})) + }) + t.Run("WithQuotaProject does not panic", func(t *testing.T) { + defer func() { + if r := recover(); r != nil { + t.Errorf("Read() with WithQuotaProject panicked: %v", r) + } + }() + _, s := beam.NewPipelineWithRoot() + Read(s, "project", "project:dataset.table", typ, WithQuotaProject("quota-project")) + }) +} From d5703b77c103fe0f0ee9e5f07226b772ca91f1d0 Mon Sep 17 00:00:00 2001 From: Shahar Epstein <60007259+shahar1@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:27:29 +0300 Subject: [PATCH 4/4] Add CHANGES.md entry for BigQuery quota project support Generated-By: Claude Opus 5.5 --- CHANGES.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGES.md b/CHANGES.md index 537c3a32447c..d7328eda2933 100644 --- a/CHANGES.md +++ b/CHANGES.md @@ -66,7 +66,7 @@ * Support for X source added (Java/Python) ([#X](https://github.com/apache/beam/issues/X)). * (Java) FirestoreIO: added `RpcQosOptions.Builder.withRampupThrottlingDisabled()`, mirroring the existing DatastoreIO option, to opt out of the 500/50/5 write ramp-up ([#27170](https://github.com/apache/beam/issues/27170)). - +* Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--quota_project_id` pipeline option, honored by the Python BigQuery IO transforms (`WriteToBigQuery` with `FILE_LOADS` or `STREAMING_INSERTS`; not yet `STORAGE_WRITE_API`, which is Java-backed), `--bigQueryQuotaProjectId` pipeline option (Java), and `bigqueryio.WithQuotaProject` read/query option (Go) ([#37431](https://github.com/apache/beam/issues/37431)). ## New Features / Improvements