gemini-code-assist[bot] commented on code in PR #39087:
URL: https://github.com/apache/beam/pull/39087#discussion_r3483048669
##########
sdks/python/apache_beam/io/gcp/big_query_query_to_table_it_test.py:
##########
@@ -153,10 +142,10 @@ def _setup_new_types_env(self):
}]
# the API Tools bigquery client expects byte values to be base-64 encoded
# TODO https://github.com/apache/beam/issues/19073: upgrade to
- # google-cloud-bigquery which does not require handling the encoding in
- # beam
- for row in table_data:
- row['bytes'] = base64.b64encode(row['bytes']).decode('utf-8')
+ # the new BigQuery client and check this behavior.
+ for r in table_data:
+ r['bytes'] = base64.b64encode(r['bytes'])
Review Comment:
With the migration to `google-cloud-bigquery`, the client natively handles
`bytes` objects and automatically base64-encodes them when sending to the
BigQuery API. Manually calling `base64.b64encode` on `r['bytes']` will result
in double-encoding if passed as raw bytes, or a `TypeError` during JSON
serialization because `base64.b64encode` returns a `bytes` object (which is not
JSON serializable).
You should remove the manual base64 encoding entirely.
```python
# No manual base64 encoding is needed with the new BigQuery client.
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -888,47 +844,35 @@ def get_or_create_dataset(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _is_table_empty(self, project_id, dataset_id, table_id):
- request = bigquery.BigqueryTabledataListRequest(
- projectId=project_id,
- datasetId=dataset_id,
- tableId=table_id,
- maxResults=1)
- response = self.client.tabledata.List(request)
- # The response is a bigquery.TableDataList instance.
- return response.totalRows == 0
+ rows = list(
+ self.client.list_rows(
+ f"{project_id}.{dataset_id}.{table_id}", max_results=1))
+ return len(rows) == 0
@retry.with_exponential_backoff(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _delete_table(self, project_id, dataset_id, table_id):
- request = bigquery.BigqueryTablesDeleteRequest(
- projectId=project_id, datasetId=dataset_id, tableId=table_id)
try:
- self.client.tables.Delete(request)
- except HttpError as exn:
- if exn.status_code == 404:
- _LOGGER.warning(
- 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id)
- return
- else:
- raise
+ self.client.delete_table(
+ f"{project_id}.{dataset_id}.{table_id}", not_found_ok=True)
+ except GoogleAPICallError as exn:
+ _LOGGER.warning(
+ 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id)
+ return
Review Comment:
By passing `not_found_ok=True` to `delete_table`, the client already handles
404 (Not Found) errors silently without raising an exception.
However, catching all `GoogleAPICallError` exceptions and swallowing them
with a warning that the table 'does not exist' is a regression. If a different
error occurs (such as a 403 Forbidden or 500 Internal Server Error), it will be
incorrectly logged as 'does not exist' and silently swallowed, which can hide
critical failures.
You should remove the `try...except` block entirely and let other exceptions
propagate.
```python
self.client.delete_table(
f"{project_id}.{dataset_id}.{table_id}", not_found_ok=True)
```
##########
sdks/python/apache_beam/io/gcp/bigquery_file_loads.py:
##########
@@ -796,6 +821,9 @@ def process(
source_format=self.source_format,
job_labels=self.bq_io_metadata.add_additional_bq_job_labels(),
load_job_project_id=self.load_job_project_id)
+
+ print("YIELDING JOB REFERENCE:", type(job_reference), job_reference)
Review Comment:
This is a leftover debug `print` statement and should be removed.
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -888,47 +844,35 @@ def get_or_create_dataset(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _is_table_empty(self, project_id, dataset_id, table_id):
- request = bigquery.BigqueryTabledataListRequest(
- projectId=project_id,
- datasetId=dataset_id,
- tableId=table_id,
- maxResults=1)
- response = self.client.tabledata.List(request)
- # The response is a bigquery.TableDataList instance.
- return response.totalRows == 0
+ rows = list(
+ self.client.list_rows(
+ f"{project_id}.{dataset_id}.{table_id}", max_results=1))
+ return len(rows) == 0
@retry.with_exponential_backoff(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _delete_table(self, project_id, dataset_id, table_id):
- request = bigquery.BigqueryTablesDeleteRequest(
- projectId=project_id, datasetId=dataset_id, tableId=table_id)
try:
- self.client.tables.Delete(request)
- except HttpError as exn:
- if exn.status_code == 404:
- _LOGGER.warning(
- 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id)
- return
- else:
- raise
+ self.client.delete_table(
+ f"{project_id}.{dataset_id}.{table_id}", not_found_ok=True)
+ except GoogleAPICallError as exn:
+ _LOGGER.warning(
+ 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id)
+ return
@retry.with_exponential_backoff(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def _delete_dataset(self, project_id, dataset_id, delete_contents=True):
- request = bigquery.BigqueryDatasetsDeleteRequest(
- projectId=project_id,
- datasetId=dataset_id,
- deleteContents=delete_contents)
try:
- self.client.datasets.Delete(request)
- except HttpError as exn:
- if exn.status_code == 404:
- _LOGGER.warning('Dataset %s:%s does not exist', project_id, dataset_id)
- return
- else:
- raise
+ self.client.delete_dataset(
+ f"{project_id}.{dataset_id}",
+ delete_contents=delete_contents,
+ not_found_ok=True)
+ except GoogleAPICallError as exn:
+ _LOGGER.warning('Dataset %s:%s does not exist', project_id, dataset_id)
+ return
Review Comment:
By passing `not_found_ok=True` to `delete_dataset`, the client already
handles 404 (Not Found) errors silently without raising an exception.
Catching all `GoogleAPICallError` exceptions and swallowing them with a
warning that the dataset 'does not exist' is a regression. If a different error
occurs (such as a 403 Forbidden or 500 Internal Server Error), it will be
incorrectly logged as 'does not exist' and silently swallowed, which can hide
critical failures.
You should remove the `try...except` block entirely and let other exceptions
propagate.
```python
self.client.delete_dataset(
f"{project_id}.{dataset_id}",
delete_contents=delete_contents,
not_found_ok=True)
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -1007,29 +949,26 @@ def _clean_up_beam_labelled_temporary_datasets(
filter_str = _build_filter_from_labels(labels)
if not self.is_user_configured_dataset() and labels is not None:
- response = (
- self.client.datasets.List(
- bigquery.BigqueryDatasetsListRequest(
- projectId=project_id, filter=filter_str)))
- for dataset in response.datasets:
+ datasets = self.client.list_datasets(
+ project=project_id, filter=filter_str)
+ for dataset in datasets:
try:
- dataset_id = dataset.datasetReference.datasetId
+ dataset_id = dataset.dataset_id
self._delete_dataset(project_id, dataset_id, True)
- except HttpError as exn:
- if exn.status_code == 403:
- _LOGGER.warning(
- 'Permission denied to delete temporary dataset %s:%s for '
- 'clean up.',
- project_id,
- dataset_id)
- return
- else:
- raise
+ except google_api_core_exceptions.Forbidden as exn:
+ _LOGGER.warning(
+ 'Permission denied to delete temporary dataset %s:%s for '
+ 'clean up.',
+ project_id,
+ dataset_id)
+ return
+ except Exception:
+ raise
Review Comment:
The `except Exception: raise` block is redundant and can be safely removed,
as unhandled exceptions will propagate naturally.
```suggestion
return
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -1103,27 +1044,36 @@ def perform_extract_job(
"""Starts a job to export data from BigQuery.
Returns:
- bigquery.JobReference with the information about the job that was
started.
+ gcp_bigquery.JobReference with the information about the job that was
started.
"""
- job_project = project or table_reference.projectId
- job_reference = bigquery.JobReference(jobId=job_id, projectId=job_project)
- request = bigquery.BigqueryJobsInsertRequest(
- projectId=job_project,
- job=bigquery.Job(
- configuration=bigquery.JobConfiguration(
- extract=bigquery.JobConfigurationExtract(
- destinationUris=destination,
- sourceTable=table_reference,
- printHeader=include_header,
- destinationFormat=destination_format,
- compression=compression,
- useAvroLogicalTypes=use_avro_logical_types,
- ),
- labels=_build_job_labels(job_labels),
- ),
- jobReference=job_reference,
- ))
- return self._start_job(request).jobReference
+ job_project = project or table_reference.project
+ job_config = gcp_bigquery.ExtractJobConfig(
+ print_header=include_header,
+ destination_format=destination_format,
+ compression=compression,
+ use_avro_logical_types=use_avro_logical_types,
+ labels=job_labels or {},
+ )
+ try:
+ job = self.client.extract_table(
+ source=table_reference,
+ destination_uris=destination,
+ job_id=job_id,
+ project=job_project,
+ job_config=job_config)
+ return job
+ except google_api_core_exceptions.Conflict as exn:
+ _LOGGER.info("BigQuery job %s already exists", job_id)
+ job_location = self._parse_location_from_exc(exn.message, job_id)
+
+ class MockJob:
+ pass
+
+ job = MockJob()
+ job.job_id = job_id
+ job.project = job_project
+ job.location = job_location
+ return job
Review Comment:
Instead of defining and returning an ad-hoc `MockJob` class when a
`Conflict` exception occurs, you can retrieve the actual existing job from the
BigQuery client using `self.get_job(job_project, job_id, job_location)`. This
ensures the returned object is a real `ExtractJob` instance, maintaining
consistency with the successful path.
```suggestion
except google_api_core_exceptions.Conflict as exn:
_LOGGER.info("BigQuery job %s already exists", job_id)
job_location = self._parse_location_from_exc(exn.message, job_id)
return self.get_job(job_project, job_id, job_location)
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -791,23 +765,12 @@ def _insert_all_rows(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_timeout_or_quota_issues_filter)
def get_table(self, project_id, dataset_id, table_id):
- """Lookup a table's metadata object.
-
- Args:
- client: bigquery.BigqueryV2 instance
- project_id: table lookup parameter
- dataset_id: table lookup parameter
- table_id: table lookup parameter
-
- Returns:
- bigquery.Table instance
- Raises:
- HttpError: if lookup failed.
- """
- request = bigquery.BigqueryTablesGetRequest(
- projectId=project_id, datasetId=dataset_id, tableId=table_id)
- response = self.client.tables.Get(request)
- return response
+ try:
+ return self.client.get_table(f"{project_id}.{dataset_id}.{table_id}")
+ except GoogleAPICallError as e:
+ if e.code == 404:
+ raise
+ raise
Review Comment:
The `try...except` block in `get_table` is redundant because both branches
of the `if/else` statement simply re-raise the caught `GoogleAPICallError`. You
can remove the `try...except` block and let the exception propagate naturally.
```python
return self.client.get_table(f"{project_id}.{dataset_id}.{table_id}")
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -1044,12 +983,13 @@ def _clean_up_beam_labelled_temporary_datasets(
num_retries=MAX_RETRIES,
retry_filter=retry.retry_on_server_errors_and_timeout_filter)
def get_job(self, project, job_id, location=None):
- request = bigquery.BigqueryJobsGetRequest()
- request.jobId = job_id
- request.projectId = project
- request.location = location
-
- return self.client.jobs.Get(request)
+ try:
+ job = self.client.get_job(job_id, project=project, location=location)
+ # Reload to get status
+ job.reload()
+ return job
+ except GoogleAPICallError as e:
+ raise
Review Comment:
The `try...except` block here is redundant because it catches
`GoogleAPICallError` and immediately re-raises it without any additional
handling or logging. You can simplify this by removing the `try...except` block
entirely.
```python
job = self.client.get_job(job_id, project=project, location=location)
# Reload to get status
job.reload()
return job
```
--
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]
