[Gemini] Migrate Python BigQuery off of apitools client by jrmccluskey · Pull Request #39087 · apache/beam

jrmccluskey · 2026-06-24T12:08:36Z

Migrates the BQ client to the supported Python library and off of the deprecated apitools client.

DO NOT MERGE - needs internal testing

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

jrmccluskey · 2026-06-26T17:26:26Z

@gemini-code-assist review

gemini-code-assist

Code Review

This pull request migrates the Apache Beam Python SDK's BigQuery integration from the legacy apitools client to the official google-cloud-bigquery library. The feedback highlights several critical issues and regressions introduced during this migration. Specifically, manual base64 encoding of bytes should be removed as the new client handles this natively, and broad exception-swallowing during table and dataset deletions should be avoided to prevent hiding critical errors. Additionally, several redundant try-except blocks, a leftover debug print statement, and an ad-hoc MockJob implementation on conflict should be cleaned up.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

jrmccluskey · 2026-07-14T13:45:12Z

@gemini-code-assist review

gemini-code-assist

Code Review

This pull request migrates the Apache Beam Python SDK's BigQuery integration from the legacy apitools client to the modern google-cloud-bigquery library. The review feedback highlights several critical runtime issues introduced by this migration, including potential AttributeErrors due to mixed usage of camelCase and snake_case attributes on job and table references, potential ImportErrors from inline imports of google.cloud.bigquery classes where fallback classes should be used, and a NameError in tests due to an unbound import. Addressing these issues by standardizing on google-cloud-bigquery types and robustly handling dictionary schemas will ensure a stable migration.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

gemini-code-assist · 2026-07-14T13:48:09Z

    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
    ]


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.

Suggested change

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

]

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

gemini-code-assist · 2026-07-14T13:48:09Z

+      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))


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.

Suggested change

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))

if state == 'DONE' and error_result:

raise RuntimeError(

'BigQuery job {} failed. Error Result: {}'.format(

job_id, error_result))

gemini-code-assist · 2026-07-14T13:48:09Z

+    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)


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).

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)

gemini-code-assist · 2026-07-14T13:48:09Z

  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


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.

Suggested change

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

def _get_temp_dataset(self):

if self.temp_dataset_id:

return self.temp_dataset_id

return BigQueryWrapper.TEMP_DATASET + self._temporary_table_suffix

gemini-code-assist · 2026-07-14T13:48:09Z

+    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)


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.

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)

gemini-code-assist · 2026-07-14T13:48:09Z

+    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)


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.

Suggested change

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)

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)

gemini-code-assist · 2026-07-14T13:48:09Z

  from google.api_core.exceptions import GoogleAPICallError
+
+  import apache_beam.io.gcp.bigquery
 except ImportError:
  GoogleAPICallError = None
+  bigquery = None


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.

Suggested change

from google.api_core.exceptions import GoogleAPICallError

import apache_beam.io.gcp.bigquery

except ImportError:

GoogleAPICallError = None

bigquery = None

from google.api_core.exceptions import GoogleAPICallError

import apache_beam.io.gcp.bigquery as bigquery

except ImportError:

GoogleAPICallError = None

bigquery = None

github-actions Bot added python examples io gcp yaml labels Jun 24, 2026

gemini-code-assist Bot reviewed Jun 26, 2026

View reviewed changes

derrickaw added the Next Action: Author label Jul 8, 2026

Abacn mentioned this pull request Jul 8, 2026

Update Bigquery Client #38950

Merged

jrmccluskey force-pushed the bqMigration branch from df61084 to f1c6c63 Compare July 9, 2026 14:28

jrmccluskey and others added 20 commits July 9, 2026 15:12

[Gemini] Migrate Python BigQuery off of apitools client

19294ca

Fix tests that break with import changes

98d4e61

isort

404ee81

yapf

a1f5c09

Fix dataframes tests

a313e5c

BQ IT import

9815ab4

More formatting

5006650

Linting, imports, formatting cleanup

0d22ea9

import issues

551b32b

more fixes

b149665

yapf

3602f75

tools import

51f5ede

please be the last yapf change

0ad534a

docs precommit fix

f32c6a1

code review suggestions

81ef1b2

Update sdks/python/apache_beam/io/gcp/bigquery_tools.py

42c3ac2

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Trigger postcommit suites

dbde319

Fix broken unit test

93e37bf

fix camelcase issues

2ca00a9

exception mocks

e6cceaf

jrmccluskey added 16 commits July 9, 2026 15:12

projectID failure

502e54f

re-add encoding

73f7edd

IT test fixes

73a8453

attribute fix

926c3b2

api repr fix

94f2e7f

re-trigger postcommit

fbda22d

matcher update

fe021ef

more matcher adjustments

2cd3eba

yapf

4fdc3cf

fix imports

72e3d45

fix error propagation

e0a6aad

try to fix prop again

9557b0d

precommit fixes

c21c140

api representation update

61c34fc

query response update

dc92c77

plumb pipeline options

783f3fb

jrmccluskey force-pushed the bqMigration branch from f1c6c63 to 783f3fb Compare July 9, 2026 15:12

gemini-code-assist Bot reviewed Jul 14, 2026

View reviewed changes

-    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
-    ]
+    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

Uh oh!

Conversation

jrmccluskey commented Jun 24, 2026

GitHub Actions Tests Status (on master branch)

Uh oh!

jrmccluskey commented Jun 26, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

jrmccluskey commented Jul 14, 2026

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist Bot Jul 14, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants