shahar1 commented on code in PR #39203:
URL: https://github.com/apache/beam/pull/39203#discussion_r3511039964
##########
sdks/python/apache_beam/internal/gcp/auth.py:
##########
@@ -82,6 +82,47 @@ def get_service_credentials(pipeline_options):
return _Credentials.get_service_credentials(pipeline_options)
+def with_quota_project(credentials, quota_project_id):
+ """For internal use only; no backwards-compatibility guarantees.
+
+ Apply a quota project to credentials if supported.
+
+ The quota project is used to bill API requests to a specific GCP project,
+ separate from the project that owns the service account or data.
+
+ Args:
+ credentials: The credentials object (either _ApitoolsCredentialsAdapter
+ or a google.auth credentials object).
+ quota_project_id: The GCP project ID to use for quota and billing.
+
+ Returns:
+ Credentials with the quota project applied, or the original credentials
+ if quota project is not supported or credentials is None.
+ """
+ if credentials is None or quota_project_id is None:
+ return credentials
Review Comment:
Done in 0f949cce66e.
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -929,6 +946,31 @@ def _export_files(self, bq):
return table.schema, metadata_list
+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.
+
+ Returns:
+ A BigQueryReadClient instance.
+ """
+ if quota_project_id:
+ try:
+ import google.auth
+ from google.auth import exceptions as auth_exceptions
+ credentials, _ = google.auth.default()
+ credentials = auth.with_quota_project(credentials, quota_project_id)
+ return bq_storage.BigQueryReadClient(credentials=credentials)
+ except (auth_exceptions.DefaultCredentialsError, AttributeError) as e:
Review Comment:
Done in 0f949cce66e.
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -1406,19 +1413,69 @@ def convert_row_to_dict(self, row, schema):
@staticmethod
def from_pipeline_options(pipeline_options: PipelineOptions):
+ """Create a BigQueryWrapper from pipeline options.
+
+ Args:
+ pipeline_options: Pipeline options containing GCP configuration.
+ The quota_project_id is read from GoogleCloudOptions if set.
+ """
+ quota_project_id = None
+ if pipeline_options is not None:
+ quota_project_id = pipeline_options.view_as(
+ GoogleCloudOptions).quota_project_id
return BigQueryWrapper(
- client=BigQueryWrapper._bigquery_client(pipeline_options))
+ client=BigQueryWrapper._bigquery_client(pipeline_options),
+ quota_project_id=quota_project_id)
@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
+ if quota_project_id is None and pipeline_options is not None:
+ quota_project_id = pipeline_options.view_as(
+ GoogleCloudOptions).quota_project_id
+ if quota_project_id:
+ credentials = auth.with_quota_project(credentials, quota_project_id)
return bigquery.BigqueryV2(
http=get_new_http(),
- credentials=auth.get_service_credentials(pipeline_options),
+ 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."""
+ credentials = None
+
+ if quota_project_id:
+ # Get default credentials and apply quota project
+ try:
+ import google.auth
+ from google.auth import exceptions as auth_exceptions
+ credentials, _ = google.auth.default()
+ credentials = auth.with_quota_project(credentials, quota_project_id)
+ except (auth_exceptions.DefaultCredentialsError, AttributeError) as e:
Review Comment:
Done in 0f949cce66e.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
