gemini-code-assist[bot] commented on code in PR #39203:
URL: https://github.com/apache/beam/pull/39203#discussion_r3510976832
##########
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:
The `try` block only catches `(auth_exceptions.DefaultCredentialsError,
AttributeError)`. However, `google.auth.default()` can raise other auth-related
exceptions (such as `RefreshError` or `TransportError`), and `import
google.auth` can raise `ModuleNotFoundError` if the dependencies are missing.
To ensure the fallback to the default client works robustly under any
credential loading failure, consider catching `Exception` instead. This also
allows simplifying the imports inside the `try` block.
```suggestion
try:
import google.auth
credentials, _ = google.auth.default()
credentials = auth.with_quota_project(credentials, quota_project_id)
return bq_storage.BigQueryReadClient(credentials=credentials)
except Exception as e:
```
##########
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:
If `_GOOGLE_AUTH_AVAILABLE` is `False`, `_ApitoolsCredentialsAdapter` is not
defined. To prevent potential `NameError` or static analysis warnings when
`google-auth` is not installed, consider adding a check for
`_GOOGLE_AUTH_AVAILABLE` at the beginning of `with_quota_project`.
```suggestion
if not _GOOGLE_AUTH_AVAILABLE or credentials is None or quota_project_id
is None:
return credentials
```
##########
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:
Similar to `_create_bq_storage_client`, catching only
`(auth_exceptions.DefaultCredentialsError, AttributeError)` can lead to
uncaught exceptions (like `RefreshError`, `TransportError`, or
`ModuleNotFoundError`) crashing the pipeline instead of falling back to the
default client. Consider catching `Exception` and simplifying the imports
inside the `try` block.
```suggestion
try:
import google.auth
credentials, _ = google.auth.default()
credentials = auth.with_quota_project(credentials, quota_project_id)
except Exception as e:
```
--
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]
