Move quota_project_id

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

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

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',

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(),

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,

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)
 

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')

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))

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 @@ 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)