claudevdm commented on code in PR #39203:
URL: https://github.com/apache/beam/pull/39203#discussion_r3703993746
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -358,11 +359,17 @@ class BigQueryWrapper(object):
HISTOGRAM_METRIC_LOGGER = MetricLogger()
- def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None):
- self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions())
- self.gcp_bq_client = client or gcp_bigquery.Client(
- client_info=ClientInfo(
- user_agent="apache-beam-%s" % apache_beam.__version__))
+ def __init__(
+ self,
+ client=None,
Review Comment:
Passing a client that overrides self.client and self.gcp_bq_client seems
error prone. Can we reserve this for testing only? And should client and
gcp_bq_client have separate overrides? (apitools BigqueryV2 vs
gcp_bigquery.Client)
Maybe we can do someting like
```
def __init__(self, client=None, gcp_client=None, temp_dataset_id=None,
temp_table_ref=None, pipeline_options=None,
quota_project_id=None):
effective = quota_project_id or _quota_from(pipeline_options)
self.client = client or self._bigquery_client(
pipeline_options or PipelineOptions(),
quota_project_id=effective)
self.gcp_bq_client = gcp_client or self._gcp_bigquery_client(
quota_project_id=effective)
```
Then in `_CustomBigQuerySource`, `_CustomBigQueryStorageSource` add a utility
```
def _create_bq_wrapper(self, **kwargs):
"""Creates a BigQueryWrapper for the API calls made by this source.
Every BigQuery client this source uses is built here, so that quota
attribution cannot be forgotten at an individual call site.
"""
return bigquery_tools.BigQueryWrapper(
pipeline_options=self.options,
quota_project_id=self.quota_project_id,
**kwargs)
```
call that from split etc
```
bq = self._create_bq_wrapper(
temp_dataset_id=(
self.temp_dataset.datasetId if self.temp_dataset else None))
```
##########
sdks/python/apache_beam/io/gcp/bigquery.py:
##########
@@ -929,6 +947,30 @@ 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
+ credentials, _ = google.auth.default()
+ credentials = auth.with_quota_project(credentials, quota_project_id)
+ return bq_storage.BigQueryReadClient(credentials=credentials)
+ except Exception as e:
Review Comment:
Should we rather fail than falling back to default if user explicitly asked
to use quota project?
--
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]
