This is an automated email from the ASF dual-hosted git repository.
claudevdm pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new cba73417cd2 Move quota_project_id (#40851)
cba73417cd2 is described below
commit cba73417cd23a9bde511bdfc1ed54b82f5978e48
Author: claudevdm <[email protected]>
AuthorDate: Tue Oct 27 21:43:33 2026 -0400
Move quota_project_id (#40851)
---
CHANGES.md | 2 +-
sdks/python/apache_beam/io/gcp/bigquery.py | 6 ++---
.../io/gcp/bigquery_change_history_test.py | 6 +++--
.../apache_beam/io/gcp/bigquery_file_loads_test.py | 3 ++-
sdks/python/apache_beam/io/gcp/bigquery_test.py | 29 +++++++++++---------
sdks/python/apache_beam/io/gcp/bigquery_tools.py | 7 +++--
.../apache_beam/io/gcp/bigquery_tools_test.py | 5 ++--
.../python/apache_beam/options/pipeline_options.py | 31 +++++++++++++---------
.../apache_beam/options/pipeline_options_test.py | 28 ++++++++++++++++++-
9 files changed, 79 insertions(+), 38 deletions(-)
diff --git a/CHANGES.md b/CHANGES.md
index 1fef9780020..59253c35bf8 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 [...]
+* 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://g [...]
* (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)).
* [IcebergIO] AddFiles is now available through the Managed API with
`Managed.write(Managed.ICEBERG_ADD_FILES)` (Java/Python)
([#40827](https://github.com/apache/beam/issues/40827)).
diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py
b/sdks/python/apache_beam/io/gcp/bigquery.py
index f3d64bfcfa6..2a40db1c78d 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 20d614c9362..f9fdda7b8d5 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 @@ class QuotaProjectTest(unittest.TestCase):
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 91ecb0d700d..73cd4df2aa4 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 @@ class
TestFileLoadsQuotaProject(_TestCaseWithTempDirCleanUp):
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 ae03104612c..c0ac8473288 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
BigqueryFullResultStreamin
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 @@ class TestReadFromBigQuery(unittest.TestCase):
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 @@ class TestReadFromBigQueryQuotaProject(unittest.TestCase):
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 @@ class TestReadFromBigQueryQuotaProject(unittest.TestCase):
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 @@ class TestReadFromBigQueryQuotaProject(unittest.TestCase):
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 @@ class TestReadFromBigQueryQuotaProject(unittest.TestCase):
@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 @@ class TestReadFromBigQueryQuotaProject(unittest.TestCase):
@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 @@ class TestWriteToBigQueryQuotaProject(unittest.TestCase):
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 ea3aa1bbbd4..0cfbf6a2a04 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 import monitoring_infos
from apache_beam.metrics.metric import Metrics
from apache_beam.options import value_provider
from apache_beam.options.pipeline_options import BigQueryOptions
-from apache_beam.options.pipeline_options import GoogleCloudOptions
from apache_beam.options.pipeline_options import PipelineOptions
from apache_beam.transforms import DoFn
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 @@ class BigQueryWrapper(object):
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 b65faac45b1..95e8969bbf2 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 @@ class
TestBigQueryWrapperQuotaProject(unittest.TestCase):
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 5b73b22476b..113dcb00e19 100644
--- a/sdks/python/apache_beam/options/pipeline_options.py
+++ b/sdks/python/apache_beam/options/pipeline_options.py
@@ -1197,17 +1197,6 @@ class GoogleCloudOptions(PipelineOptions):
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 @@ class BigQueryOptions(PipelineOptions):
'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 94083a55c6f..79ea0aa650f 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:<dest>:v1, and Java derives the property name from <dest>.
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 @@ class BigQueryOptionsTest(unittest.TestCase):
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 @@ class BigQueryOptionsTest(unittest.TestCase):
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)
