From ee1c57077aaa2bfcfad06ba3db1868400b419ae2 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 16:04:28 -0400 Subject: [PATCH] Move quota_project_id --- CHANGES.md | 2 +- sdks/python/apache_beam/io/gcp/bigquery.py | 6 ++-- .../io/gcp/bigquery_change_history_test.py | 6 ++-- .../io/gcp/bigquery_file_loads_test.py | 3 +- .../apache_beam/io/gcp/bigquery_test.py | 29 ++++++++++------- .../apache_beam/io/gcp/bigquery_tools.py | 7 ++--- .../apache_beam/io/gcp/bigquery_tools_test.py | 5 +-- .../apache_beam/options/pipeline_options.py | 31 ++++++++++++------- .../options/pipeline_options_test.py | 28 ++++++++++++++++- 9 files changed, 79 insertions(+), 38 deletions(-) diff --git a/CHANGES.md b/CHANGES.md index 04ac961fcd93..69484124319a 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)). +* Added support for attributing BigQuery API quota and billing to a specific GCP project (quota project): `quota_project_id` parameter in `ReadFromBigQuery` or `--bigquery_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)). * (Java/Python/Go) DebeziumIO: `readAsJson()` can now include `op` and the remaining Debezium `source` fields (e.g. MySQL `file`, `pos`, `row`, `ts_ms`; PostgreSQL `lsn`, `txId`, `ts_ms`) in `metadata` via `withFullMetadata(true)` / `full_metadata=True` / `debeziumio.FullMetadata()`. Off by default ([#40701](https://github.com/apache/beam/issues/40701)). ## New Features / Improvements diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index f3d64bfcfa6c..2a40db1c78d2 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -3103,9 +3103,9 @@ class ReadFromBigQuery(PTransform): 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. + ``--bigquery_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 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 20d614c93622..f9fdda7b8d59 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 @@ -298,8 +298,10 @@ def test_read_sdf_creates_storage_client_with_quota_project(self): 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']) + options = PipelineOptions([ + '--bigquery_quota_project_id=' + self._QUOTA, + '--no_pipeline_type_check' + ]) classes = [ '_PollChangeHistoryFn', '_ExecuteQueryFn', 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 91ecb0d700de..73cd4df2aa4d 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 @@ -1357,7 +1357,8 @@ def test_start_bundle_passes_quota_project(self, dofn, quota_attr): 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']) + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) transform = bqfl.BigQueryBatchFileLoads( 'project:dataset.table', custom_gcs_temp_location=self._new_tempdir(), diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index ae03104612c9..c0ac84732880 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -66,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 BigQueryOptions from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.pipeline_options import StandardOptions @@ -892,16 +893,17 @@ def test_expand_output_type_uses_query_schema(self): 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') + """Test that quota_project_id is read from BigQueryOptions.""" + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) + bq_options = options.view_as(BigQueryOptions) + self.assertEqual(bq_options.bigquery_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) + bq_options = options.view_as(BigQueryOptions) + self.assertIsNone(bq_options.bigquery_quota_project_id) def test_export_source_explicit_quota_project(self): """Test that explicit quota_project_id is stored in @@ -916,7 +918,8 @@ def test_export_source_explicit_quota_project(self): 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']) + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) source = beam_bq._CustomBigQuerySource( method=ReadFromBigQuery.Method.EXPORT, table='project:dataset.table', @@ -925,7 +928,7 @@ def test_export_source_gets_quota_from_options(self): def test_export_source_explicit_overrides_options(self): """Test that explicit quota_project_id overrides options.""" - options = PipelineOptions(['--quota_project_id=options-project']) + options = PipelineOptions(['--bigquery_quota_project_id=options-project']) source = beam_bq._CustomBigQuerySource( method=ReadFromBigQuery.Method.EXPORT, table='project:dataset.table', @@ -945,7 +948,8 @@ def test_storage_source_explicit_quota_project(self): def test_storage_source_gets_quota_from_options(self): """Test that _CustomBigQueryStorageSource falls back to options.""" - options = PipelineOptions(['--quota_project_id=my-billing-project']) + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) source = beam_bq._CustomBigQueryStorageSource( method=ReadFromBigQuery.Method.DIRECT_READ, table='project:dataset.table', @@ -955,7 +959,7 @@ def test_storage_source_gets_quota_from_options(self): @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.""" + client, not only the --bigquery_quota_project_id pipeline option.""" mock_bq_client.side_effect = RuntimeError('stop') source = beam_bq._CustomBigQuerySource( method=ReadFromBigQuery.Method.EXPORT, @@ -972,7 +976,7 @@ def test_export_source_split_passes_quota_to_client(self, mock_bq_client): @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.""" + client, not only the --bigquery_quota_project_id pipeline option.""" mock_bq_client.side_effect = RuntimeError('stop') source = beam_bq._CustomBigQueryStorageSource( method=ReadFromBigQuery.Method.DIRECT_READ, @@ -1204,7 +1208,8 @@ def test_write_fn_start_bundle_passes_quota_project(self, mock_wrapper): 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']) + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) p = TestPipeline(options=options) with mock.patch.object( beam_bq.BigQueryWriteFn, diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index ea3aa1bbbd40..0cfbf6a2a048 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -62,7 +62,6 @@ 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 from apache_beam.typehints.row_type import RowTypeConstraint @@ -500,10 +499,10 @@ def _build_dataset_encryption_config(kms_key): def _quota_project_id_from_options(pipeline_options): - """Returns --quota_project_id from pipeline options, if any.""" + """Returns --bigquery_quota_project_id from pipeline options, if any.""" if pipeline_options is None: return None - return pipeline_options.view_as(GoogleCloudOptions).quota_project_id + return pipeline_options.view_as(BigQueryOptions).bigquery_quota_project_id def create_bq_storage_client(quota_project_id=None): @@ -1654,7 +1653,7 @@ def from_pipeline_options(pipeline_options: PipelineOptions): Args: pipeline_options: Pipeline options containing GCP configuration. - The quota_project_id is read from GoogleCloudOptions if set. + The quota project is read from --bigquery_quota_project_id if set. """ return BigQueryWrapper(pipeline_options=pipeline_options) 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 b65faac45b19..95e8969bbf2d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -1655,13 +1655,14 @@ def test_quota_project_id_stored(self, mock_gcp_client, mock_bq_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.""" + BigQueryOptions.""" 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']) + options = PipelineOptions( + ['--bigquery_quota_project_id=my-billing-project']) wrapper = BigQueryWrapper.from_pipeline_options(options) self.assertEqual(wrapper.quota_project_id, 'my-billing-project') diff --git a/sdks/python/apache_beam/options/pipeline_options.py b/sdks/python/apache_beam/options/pipeline_options.py index 5b73b22476b4..113dcb00e190 100644 --- a/sdks/python/apache_beam/options/pipeline_options.py +++ b/sdks/python/apache_beam/options/pipeline_options.py @@ -1197,17 +1197,6 @@ 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, @@ -1397,7 +1386,10 @@ class BigQueryOptions(PipelineOptions): """ # Alternative attribute names, mapped to the option names they stand for. - _ATTRIBUTE_ALIASES = {'bigquery_endpoint': 'big_query_endpoint'} + _ATTRIBUTE_ALIASES = { + 'bigquery_endpoint': 'big_query_endpoint', + 'bigquery_quota_project_id': 'big_query_quota_project_id', + } @classmethod def _add_argparse_args(cls, parser): @@ -1413,6 +1405,21 @@ def _add_argparse_args(cls, parser): 'variable when present, and otherwise the default public URL is ' 'used. Also passed to Java cross-language BigQuery transforms as ' 'bigQueryEndpoint.')) + parser.add_argument( + '--bigquery_quota_project_id', + '--big_query_quota_project_id', + dest='big_query_quota_project_id', + default=None, + help=( + '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. Honored by ' + 'ReadFromBigQuery, ReadAllFromBigQuery, WriteToBigQuery with ' + 'FILE_LOADS or STREAMING_INSERTS, and ReadBigQueryChangeHistory; ' + 'STORAGE_WRITE_API (Java-backed) does not use it yet. The ' + 'credentials used must have the serviceusage.services.use ' + 'permission on this project. Also passed to Java cross-language ' + 'BigQuery transforms as bigQueryQuotaProjectId.')) def __getattr__(self, name): return super().__getattr__(self._ATTRIBUTE_ALIASES.get(name, name)) diff --git a/sdks/python/apache_beam/options/pipeline_options_test.py b/sdks/python/apache_beam/options/pipeline_options_test.py index 94083a55c6f4..79ea0aa650f8 100644 --- a/sdks/python/apache_beam/options/pipeline_options_test.py +++ b/sdks/python/apache_beam/options/pipeline_options_test.py @@ -1284,6 +1284,7 @@ class BigQueryOptionsTest(unittest.TestCase): # beam:option::v1, and Java derives the property name from . JAVA_OPTION_NAMES = { 'big_query_endpoint': 'bigQueryEndpoint', + 'big_query_quota_project_id': 'bigQueryQuotaProjectId', } def test_option_names_map_to_java_option_names(self): @@ -1297,12 +1298,19 @@ def test_option_names_map_to_java_option_names(self): 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']) + options = PipelineOptions([ + '--bigquery_endpoint=https://bq.example.com', + '--bigquery_quota_project_id=billing-project', + ]) 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) + self.assertEqual( + fields['beam:option:big_query_quota_project_id:v1'].string_value, + 'billing-project') + self.assertNotIn('beam:option:bigquery_quota_project_id:v1', fields) @parameterized.expand([ ('--bigquery_endpoint', ), @@ -1317,6 +1325,24 @@ def test_endpoint_flags(self, flag): self.assertEqual(bq_options.big_query_endpoint, 'https://bq.example.com') self.assertEqual(bq_options.bigquery_endpoint, 'https://bq.example.com') + @parameterized.expand([ + ('--bigquery_quota_project_id', ), + ('--big_query_quota_project_id', ), + ]) + def test_quota_project_id_flags(self, flag): + options = PipelineOptions([flag, 'billing-project']) + self.assertEqual( + options.get_all_options()['big_query_quota_project_id'], + 'billing-project') + bq_options = options.view_as(BigQueryOptions) + self.assertEqual(bq_options.big_query_quota_project_id, 'billing-project') + self.assertEqual(bq_options.bigquery_quota_project_id, 'billing-project') + + def test_quota_project_id_default(self): + bq_options = PipelineOptions([]).view_as(BigQueryOptions) + self.assertIsNone(bq_options.big_query_quota_project_id) + self.assertIsNone(bq_options.bigquery_quota_project_id) + def test_endpoint_default(self): bq_options = PipelineOptions([]).view_as(BigQueryOptions) self.assertIsNone(bq_options.big_query_endpoint)