gemini-code-assist[bot] commented on code in PR #39087:
URL: https://github.com/apache/beam/pull/39087#discussion_r3579520403
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -646,53 +630,62 @@ def _start_query_job(
dry_run=False,
kms_key=None,
job_labels=None):
- reference = bigquery.JobReference(jobId=job_id, projectId=project_id)
- request = bigquery.BigqueryJobsInsertRequest(
- projectId=project_id,
- job=bigquery.Job(
- configuration=bigquery.JobConfiguration(
- dryRun=dry_run,
- query=bigquery.JobConfigurationQuery(
- query=query,
- useLegacySql=use_legacy_sql,
- allowLargeResults=not dry_run,
- destinationTable=self._get_temp_table(
- self._get_temp_table_project(project_id))
- if not dry_run else None,
- flattenResults=flatten_results,
- priority=priority,
- destinationEncryptionConfiguration=bigquery.
- EncryptionConfiguration(kmsKeyName=kms_key)),
- labels=_build_job_labels(job_labels),
- ),
- jobReference=reference))
-
- return self._start_job(request)
-
- def wait_for_bq_job(self, job_reference, sleep_duration_sec=5,
max_retries=0):
- """Poll job until it is DONE.
+ job_config = gcp_bigquery.QueryJobConfig(
+ use_legacy_sql=use_legacy_sql,
+ flatten_results=flatten_results,
+ priority=priority,
+ dry_run=dry_run,
+ labels=job_labels or {},
+ allow_large_results=not dry_run,
+ )
+ if kms_key:
+ job_config.destination_encryption_configuration =
gcp_bigquery.EncryptionConfiguration(
+ kms_key_name=kms_key)
+ if not dry_run:
+ job_config.destination = self._get_temp_table(
+ self._get_temp_table_project(project_id))
- Args:
- job_reference: bigquery.JobReference instance.
- sleep_duration_sec: Specifies the delay in seconds between retries.
- max_retries: The total number of times to retry. If equals to 0,
- the function waits forever.
+ try:
+ job = self.client.query(
+ query, job_config=job_config, job_id=job_id, project=project_id)
+ return BeamJobReference(
+ projectId=job.project,
+ jobId=job.job_id,
+ location=job.location,
+ totalBytesProcessed=getattr(job, 'total_bytes_processed', None))
+ 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 BeamJobReference(
+ projectId=project_id, jobId=job_id, location=job_location)
- Raises:
- `RuntimeError`: If the job is FAILED or the number of retries has been
- reached.
- """
+ def wait_for_bq_job(self, job_reference, sleep_duration_sec=5,
max_retries=0):
retry = 0
+ project_id = getattr(
+ job_reference, 'project', getattr(job_reference, 'projectId', None))
+ job_id = getattr(
+ job_reference, 'job_id', getattr(job_reference, 'jobId', None))
+ location = getattr(job_reference, 'location', None)
while True:
retry += 1
- job = self.get_job(
- job_reference.projectId, job_reference.jobId, job_reference.location)
- _LOGGER.info('Job %s status: %s', job.id, job.status.state)
- if job.status.state == 'DONE' and job.status.errorResult:
+ job = self.get_job(project_id, job_id, location)
+
+ if hasattr(job, 'status') and hasattr(job.status, 'state'):
+ state = job.status.state
+ error_result = job.status.errorResult
+ job_id = getattr(job, 'id', getattr(job_reference, 'jobId', ''))
+ else:
+ state = job.state
+ error_result = job.error_result
+ job_id = job.job_id if hasattr(job, 'job_id') else getattr(
+ job_reference, 'jobId', '')
+
+ _LOGGER.info('Job %s status: %s', job_id, state)
+ if state == 'DONE' and error_result:
raise RuntimeError(
'BigQuery job {} failed. Error Result: {}'.format(
- job_reference.jobId, job.status.errorResult))
- elif job.status.state == 'DONE':
+ job_reference.jobId, error_result))
Review Comment:
The `job_reference` object can be a `google.cloud.bigquery.job` instance
(such as `LoadJob` or `QueryJob`) which has `job_id` but not `jobId`. Accessing
`job_reference.jobId` will raise an `AttributeError` at runtime. Use the local
variable `job_id` instead, which was safely extracted using `getattr` above.
```suggestion
if state == 'DONE' and error_result:
raise RuntimeError(
'BigQuery job {} failed. Error Result: {}'.format(
job_id, error_result))
```
##########
sdks/python/apache_beam/io/gcp/bigquery_read_internal.py:
##########
@@ -444,15 +469,19 @@ def _to_bytes(value):
@classmethod
def _convert_to_tuple(cls, table_field_schemas):
- """Recursively converts the list of TableFieldSchema instances to the
+ """Recursively converts the list of SchemaField instances to the
list of tuples to prevent errors when pickling and unpickling
- TableFieldSchema instances.
+ SchemaField instances.
"""
if not table_field_schemas:
return []
return [
- FieldSchema(cls._convert_to_tuple(x.fields), x.mode, x.name, x.type)
+ FieldSchema(
+ cls._convert_to_tuple(x.fields),
+ x.mode,
+ x.name,
+ getattr(x, "field_type", getattr(x, "type", None)))
for x in table_field_schemas
]
Review Comment:
When `table_schema` is a dictionary, `table_field_schemas` is a list of
dictionaries. Attempting to access attributes like `x.fields`, `x.mode`,
`x.name` on `x` will raise an `AttributeError` because `x` is a dictionary. We
should robustly handle both `SchemaField` objects and `dict` objects.
```suggestion
res = []
for x in table_field_schemas:
if isinstance(x, dict):
fields = x.get('fields', [])
mode = x.get('mode', 'NULLABLE')
name = x.get('name')
field_type = x.get('type')
else:
fields = getattr(x, 'fields', [])
mode = getattr(x, 'mode', 'NULLABLE')
name = getattr(x, 'name')
field_type = getattr(x, 'field_type', getattr(x, 'type', None))
res.append(FieldSchema(cls._convert_to_tuple(fields), mode, name,
field_type))
return res
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -1103,27 +1063,28 @@ 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)
+ return self.get_job(job_project, job_id, job_location)
Review Comment:
To maintain consistency with other job-starting methods in `BigQueryWrapper`
(such as `_insert_load_job`, `_insert_copy_job`, and `_start_query_job`),
`perform_extract_job` should return a `BeamJobReference` instead of an
`ExtractJob` or `Job` object. Callers expecting a job reference with
`projectId` and `jobId` attributes will fail with `AttributeError` if they
receive an `ExtractJob` or `Job` object (which only have `project` and
`job_id`).
```python
try:
job = self.client.extract_table(
source=table_reference,
destination_uris=destination,
job_id=job_id,
project=job_project,
job_config=job_config)
return BeamJobReference(
projectId=job.project, jobId=job.job_id, location=job.location)
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 BeamJobReference(
projectId=job_project, jobId=job_id, location=job_location)
```
##########
sdks/python/apache_beam/io/gcp/bigquery_file_loads.py:
##########
@@ -397,17 +397,22 @@ def process(self, element, schema_mod_job_name_prefix):
return
table_reference = bigquery_tools.parse_table_reference(destination)
- if table_reference.projectId is None:
- table_reference.projectId = vp.RuntimeValueProvider.get_value(
+ if table_reference.project is None or table_reference.project ==
bigquery_tools.FALLBACK_PROJECT:
+ from google.cloud.bigquery import DatasetReference
+ from google.cloud.bigquery import TableReference
+ new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
+ table_reference = TableReference(
+ DatasetReference(new_project, table_reference.dataset_id),
+ table_reference.table_id)
Review Comment:
Avoid inline imports of `DatasetReference` and `TableReference` from
`google.cloud.bigquery` inside `process`. This can raise `ImportError` if
`google-cloud-bigquery` is not installed, and is redundant because
`bigquery_tools` already exports `TableReference` and `DatasetReference` safely
with fallback dummy classes.
```python
if table_reference.project is None or table_reference.project ==
bigquery_tools.FALLBACK_PROJECT:
new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
table_reference = bigquery_tools.TableReference(
bigquery_tools.DatasetReference(new_project,
table_reference.dataset_id),
table_reference.table_id)
```
##########
sdks/python/apache_beam/io/gcp/bigquery_file_loads.py:
##########
@@ -537,15 +542,26 @@ def process_one(self, element, job_name_prefix):
destination, job_reference = element
copy_to_reference = bigquery_tools.parse_table_reference(destination)
- if copy_to_reference.projectId is None:
- copy_to_reference.projectId = vp.RuntimeValueProvider.get_value(
+ if copy_to_reference.project is None or copy_to_reference.project ==
bigquery_tools.FALLBACK_PROJECT:
+ from google.cloud.bigquery import DatasetReference
+ from google.cloud.bigquery import TableReference
+ new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
+ copy_to_reference = TableReference(
+ DatasetReference(new_project, copy_to_reference.dataset_id),
+ copy_to_reference.table_id)
copy_from_reference = bigquery_tools.parse_table_reference(destination)
- copy_from_reference.tableId = job_reference.jobId
- if copy_from_reference.projectId is None:
- copy_from_reference.projectId = vp.RuntimeValueProvider.get_value(
+ new_table_id = job_reference.jobId
+ new_project = copy_from_reference.project
+ if new_project is None or new_project == bigquery_tools.FALLBACK_PROJECT:
+ new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
+ from google.cloud.bigquery import DatasetReference
+ from google.cloud.bigquery import TableReference
+ copy_from_reference = TableReference(
+ DatasetReference(new_project, copy_from_reference.dataset_id),
+ new_table_id)
Review Comment:
Avoid inline imports of `DatasetReference` and `TableReference` from
`google.cloud.bigquery` inside `process_one`. This can raise `ImportError` if
`google-cloud-bigquery` is not installed, and is redundant because
`bigquery_tools` already exports `TableReference` and `DatasetReference` safely
with fallback dummy classes.
```suggestion
if copy_to_reference.project is None or copy_to_reference.project ==
bigquery_tools.FALLBACK_PROJECT:
new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
copy_to_reference = bigquery_tools.TableReference(
bigquery_tools.DatasetReference(new_project,
copy_to_reference.dataset_id),
copy_to_reference.table_id)
copy_from_reference = bigquery_tools.parse_table_reference(destination)
new_table_id = job_reference.jobId
new_project = copy_from_reference.project
if new_project is None or new_project == bigquery_tools.FALLBACK_PROJECT:
new_project = vp.RuntimeValueProvider.get_value(
'project', str, '') or self.project
copy_from_reference = bigquery_tools.TableReference(
bigquery_tools.DatasetReference(new_project,
copy_from_reference.dataset_id),
new_table_id)
```
##########
sdks/python/apache_beam/dataframe/io_it_test.py:
##########
@@ -34,8 +33,11 @@
try:
from google.api_core.exceptions import GoogleAPICallError
+
+ import apache_beam.io.gcp.bigquery
except ImportError:
GoogleAPICallError = None
+ bigquery = None
Review Comment:
In the `try` block, `import apache_beam.io.gcp.bigquery` is called, but
`bigquery` is not bound to the local scope (only `apache_beam` is). In the
`except` block, `bigquery = None` is set. If `bigquery` is accessed directly
later in the file, it will raise a `NameError` when `ImportError` is not
raised. Import it as `import apache_beam.io.gcp.bigquery as bigquery` instead.
```suggestion
from google.api_core.exceptions import GoogleAPICallError
import apache_beam.io.gcp.bigquery as bigquery
except ImportError:
GoogleAPICallError = None
bigquery = None
```
##########
sdks/python/apache_beam/io/gcp/bigquery_tools.py:
##########
@@ -425,13 +443,16 @@ def _get_temp_table_project(self, fallback_project_id):
Otherwise, returns the fallback_project_id.
"""
if self.temp_table_ref:
- return self.temp_table_ref.projectId
+ return getattr(
+ self.temp_table_ref,
+ 'projectId',
+ getattr(self.temp_table_ref, 'project', None))
else:
return fallback_project_id
def _get_temp_dataset(self):
if self.temp_table_ref:
- return self.temp_table_ref.datasetId
+ return self.temp_table_ref.dataset_id
return BigQueryWrapper.TEMP_DATASET + self._temporary_table_suffix
Review Comment:
If `self.temp_table_ref` is an older `apitools` `TableReference` or a mock
object that only has `datasetId` (camelCase), accessing
`self.temp_table_ref.dataset_id` will raise an `AttributeError`. Since
`self.temp_dataset_id` is already safely initialized in `__init__` using
`getattr` to handle both `datasetId` and `dataset_id`, `_get_temp_dataset`
should simply return `self.temp_dataset_id`.
```suggestion
def _get_temp_dataset(self):
if self.temp_dataset_id:
return self.temp_dataset_id
return BigQueryWrapper.TEMP_DATASET + self._temporary_table_suffix
```
--
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]
