From bf8664f8715bcd6e21c1f06664ac03978ba85f23 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Tue, 25 Aug 2026 21:36:11 +0000 Subject: [PATCH 01/12] [Gemini] Migrate Python BigQuery away from apitools, deprecate old route --- .../trigger_files/beam_PostCommit_Python.json | 2 +- .../beam_PostCommit_Python_Arm.json | 2 +- .../beam_PostCommit_Python_Dependency.json | 2 +- ...m_PostCommit_Python_Examples_Dataflow.json | 2 +- ...it_Python_ValidatesContainer_Dataflow.json | 2 +- ...ommit_Python_ValidatesRunner_Dataflow.json | 2 +- ...stCommit_Python_ValidatesRunner_Flink.json | 3 +- ...stCommit_Python_ValidatesRunner_Spark.json | 2 +- .../beam_PostCommit_Python_Versions.json | 2 +- ..._PostCommit_Python_Xlang_Gcp_Dataflow.json | 2 +- ...am_PostCommit_Python_Xlang_Gcp_Direct.json | 2 +- ...m_PostCommit_Python_Xlang_IO_Dataflow.json | 2 +- ...eam_PostCommit_Python_Xlang_IO_Direct.json | 2 +- ...stCommit_XVR_JavaUsingPython_Dataflow.json | 2 +- ...stCommit_XVR_PythonUsingJava_Dataflow.json | 2 +- sdks/python/apache_beam/io/gcp/bigquery.py | 68 +- .../apache_beam/io/gcp/bigquery_avro_tools.py | 11 +- .../io/gcp/bigquery_change_history.py | 103 +- .../apache_beam/io/gcp/bigquery_file_loads.py | 66 +- .../io/gcp/bigquery_read_internal.py | 18 +- .../io/gcp/bigquery_schema_tools.py | 73 +- .../apache_beam/io/gcp/bigquery_test.py | 51 +- .../apache_beam/io/gcp/bigquery_tools.py | 1856 +++++++++++------ .../apache_beam/io/gcp/bigquery_tools_test.py | 2 + .../gcp/internal/clients/bigquery/__init__.py | 18 +- sdks/python/setup.py | 21 +- 26 files changed, 1589 insertions(+), 729 deletions(-) diff --git a/.github/trigger_files/beam_PostCommit_Python.json b/.github/trigger_files/beam_PostCommit_Python.json index e8079b053aeb..f4917c852b7a 100644 --- a/.github/trigger_files/beam_PostCommit_Python.json +++ b/.github/trigger_files/beam_PostCommit_Python.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", "pr": "38701", - "modification": 57 + "modification": 58 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Arm.json b/.github/trigger_files/beam_PostCommit_Python_Arm.json index 1efc8e9e4405..3f63c0c9975f 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Arm.json +++ b/.github/trigger_files/beam_PostCommit_Python_Arm.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 1 + "modification": 2 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Dependency.json b/.github/trigger_files/beam_PostCommit_Python_Dependency.json index 16209484727a..9bc6afbb0969 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Dependency.json +++ b/.github/trigger_files/beam_PostCommit_Python_Dependency.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 4 + "modification": 5 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Examples_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Examples_Dataflow.json index d7118310af8d..d71395445b0c 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Examples_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Examples_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", "pr": "37360", - "modification": 2 + "modification": 3 } diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json index 327fe1987c50..da85a8432a0a 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesContainer_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 5 + "modification": 6 } \ No newline at end of file diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json index 462412f7bd34..6c7f95f65813 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Dataflow.json @@ -1,5 +1,5 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 5, + "modification": 6, "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test" } diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json index 962795fcb0ae..bdb80787b223 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Flink.json @@ -2,5 +2,6 @@ "https://github.com/apache/beam/pull/32648": "testing addition of Flink 1.19 support", "https://github.com/apache/beam/pull/34830": "testing", "trigger-2026-04-04": "portable_runner expand_sdf opt-in", - "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test" + "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test", + "modification": 1 } diff --git a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json index 6384446f50e4..52a2de5aff1f 100644 --- a/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json +++ b/.github/trigger_files/beam_PostCommit_Python_ValidatesRunner_Spark.json @@ -3,6 +3,6 @@ "https://github.com/apache/beam/issues/35429": "testing", "trigger-2026-04-04": "portable_runner expand_sdf opt-in", "https://github.com/apache/beam/pull/38892": "UnboundedSource portable VR test", - "modification": 1, + "modification": 2, "https://github.com/apache/beam/issues/19468": "SDF self-checkpointing and bundle finalization" } diff --git a/.github/trigger_files/beam_PostCommit_Python_Versions.json b/.github/trigger_files/beam_PostCommit_Python_Versions.json index 8b2c8c445c1f..b5704c67ef1c 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Versions.json +++ b/.github/trigger_files/beam_PostCommit_Python_Versions.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "revision": 5 + "revision": 6 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json index 83346d34aee0..c5309eebb070 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 16 + "modification": 17 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json index b26833333238..c537844dc84a 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_Gcp_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 2 + "modification": 3 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json index c537844dc84a..e0266d62f2e0 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 3 + "modification": 4 } diff --git a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json index b26833333238..c537844dc84a 100644 --- a/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json +++ b/.github/trigger_files/beam_PostCommit_Python_Xlang_IO_Direct.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run", - "modification": 2 + "modification": 3 } diff --git a/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json b/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json index 6a55e29ae15d..b04c4b0b8efe 100644 --- a/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_XVR_JavaUsingPython_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 1 + "modification": 2 } \ No newline at end of file diff --git a/.github/trigger_files/beam_PostCommit_XVR_PythonUsingJava_Dataflow.json b/.github/trigger_files/beam_PostCommit_XVR_PythonUsingJava_Dataflow.json index b73af5e61a43..7ab7bcd9a9c6 100644 --- a/.github/trigger_files/beam_PostCommit_XVR_PythonUsingJava_Dataflow.json +++ b/.github/trigger_files/beam_PostCommit_XVR_PythonUsingJava_Dataflow.json @@ -1,4 +1,4 @@ { "comment": "Modify this file in a trivial way to cause this test suite to run.", - "modification": 1 + "modification": 2 } diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 314effad5520..09e957d25e31 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -389,8 +389,14 @@ def chain_after(result): from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanup from apache_beam.io.gcp.bigquery_read_internal import _PassThroughThenCleanupTempDatasets from apache_beam.io.gcp.bigquery_read_internal import bigquery_export_destination_uri +from apache_beam.io.gcp.bigquery_tools import DatasetReference +from apache_beam.io.gcp.bigquery_tools import JobReference from apache_beam.io.gcp.bigquery_tools import RetryStrategy -from apache_beam.io.gcp.internal.clients import bigquery +from apache_beam.io.gcp.bigquery_tools import TableCell +from apache_beam.io.gcp.bigquery_tools import TableFieldSchema +from apache_beam.io.gcp.bigquery_tools import TableReference +from apache_beam.io.gcp.bigquery_tools import TableRow +from apache_beam.io.gcp.bigquery_tools import TableSchema from apache_beam.io.iobase import BoundedSource from apache_beam.io.iobase import RangeTracker from apache_beam.io.iobase import SDFBoundedSourceReader @@ -421,13 +427,12 @@ def chain_after(result): from apache_beam.utils.annotations import deprecated try: - from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference - from apache_beam.io.gcp.internal.clients.bigquery import JobReference - from apache_beam.io.gcp.internal.clients.bigquery import TableReference + from apache_beam.io.gcp.internal.clients import bigquery except ImportError: - DatasetReference = None - TableReference = None - JobReference = None + bigquery = None + +if bigquery is None or not hasattr(bigquery, 'TableReference'): + import apache_beam.io.gcp.bigquery_tools as bigquery _LOGGER = logging.getLogger(__name__) @@ -767,14 +772,38 @@ def estimate_size(self): # no access to the query that we're running. return None + def _get_temp_dataset_id(self): + if self.temp_dataset is None: + return None + elif hasattr(self.temp_dataset, 'datasetId'): + return self.temp_dataset.datasetId + elif hasattr(self.temp_dataset, 'dataset_id'): + return self.temp_dataset.dataset_id + elif isinstance(self.temp_dataset, str): + if ':' in self.temp_dataset or '.' in self.temp_dataset: + return bigquery_tools.parse_table_reference(self.temp_dataset).datasetId + return self.temp_dataset + return None + + def _get_temp_dataset_project(self): + if hasattr(self.temp_dataset, 'projectId') and self.temp_dataset.projectId: + return self.temp_dataset.projectId + elif hasattr(self.temp_dataset, 'project') and self.temp_dataset.project: + return self.temp_dataset.project + elif isinstance(self.temp_dataset, str) and (':' in self.temp_dataset or + '.' in self.temp_dataset): + return bigquery_tools.parse_table_reference(self.temp_dataset).projectId + return None + def _get_project(self): """Returns the project that queries and exports will be billed to.""" + temp_project = self._get_temp_dataset_project() + if temp_project: + return temp_project project = self.options.view_as(GoogleCloudOptions).project if isinstance(project, vp.ValueProvider): project = project.get() - if self.temp_dataset: - return self.temp_dataset.projectId if not project: project = self.project return project @@ -794,8 +823,7 @@ def _create_source(self, path, coder): def split(self, desired_bundle_size, start_position=None, stop_position=None): if self.export_result is None: bq = bigquery_tools.BigQueryWrapper( - temp_dataset_id=( - self.temp_dataset.datasetId if self.temp_dataset else None), + temp_dataset_id=self._get_temp_dataset_id(), client=bigquery_tools.BigQueryWrapper._bigquery_client(self.options)) if self.query is not None: @@ -3088,10 +3116,20 @@ def file_path_to_remove(unused_elm): def _expand_direct_read(self, pcoll): project_id = None temp_table_ref = None - if 'temp_dataset' in self._kwargs: - temp_table_ref = bigquery.TableReference( - projectId=self._kwargs['temp_dataset'].projectId, - datasetId=self._kwargs['temp_dataset'].datasetId, + temp_dataset = self._kwargs.get('temp_dataset') + if temp_dataset is not None: + if isinstance(temp_dataset, str): + ds_ref = bigquery_tools.parse_table_reference(temp_dataset) + project_id = ds_ref.projectId + dataset_id = ds_ref.datasetId + else: + project_id = getattr(temp_dataset, 'projectId', None) or getattr( + temp_dataset, 'project', None) + dataset_id = getattr(temp_dataset, 'datasetId', None) or getattr( + temp_dataset, 'dataset_id', None) + temp_table_ref = TableReference( + projectId=project_id, + datasetId=dataset_id, tableId='beam_temp_table_' + uuid.uuid4().hex) else: project_id = pcoll.pipeline.options.view_as(GoogleCloudOptions).project diff --git a/sdks/python/apache_beam/io/gcp/bigquery_avro_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_avro_tools.py index ceab52444bb1..15dc919f1e41 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_avro_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_avro_tools.py @@ -60,6 +60,8 @@ "precision": 38, "scale": 9, }, + "BIGNUMERIC": "string", + "JSON": "string", "GEOGRAPHY": "string", } @@ -108,11 +110,12 @@ def table_field_to_avro_field(table_field: dict[str, Any], """ assert "type" in table_field, \ "Unable to get type for table field {}".format(table_field) - assert table_field["type"] in BIG_QUERY_TO_AVRO_TYPES, \ + field_type_str = table_field["type"].upper() + assert field_type_str in BIG_QUERY_TO_AVRO_TYPES, \ "Unable to map BigQuery field type {} to avro type".format( table_field["type"]) - avro_type = BIG_QUERY_TO_AVRO_TYPES[table_field["type"]] + avro_type = BIG_QUERY_TO_AVRO_TYPES[field_type_str] if avro_type == "record": element_type = get_record_schema_from_dict_table_schema( @@ -122,9 +125,9 @@ def table_field_to_avro_field(table_field: dict[str, Any], else: element_type = avro_type - field_mode = table_field.get("mode", "NULLABLE") + field_mode = (table_field.get("mode") or "NULLABLE").upper() - if field_mode in (None, "NULLABLE"): + if field_mode in (None, "NULLABLE", ""): field_type = ["null", element_type] elif field_mode == "REQUIRED": field_type = element_type diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py index 1e90c2974778..d6583c18de26 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py @@ -55,7 +55,16 @@ import apache_beam as beam from apache_beam.io.gcp import bigquery_tools -from apache_beam.io.gcp.internal.clients import bigquery +from apache_beam.io.gcp.bigquery_tools import JobReference +from apache_beam.io.gcp.bigquery_tools import TableReference + +try: + from apache_beam.io.gcp.internal.clients import bigquery +except ImportError: + bigquery = None + +if bigquery is None or not hasattr(bigquery, 'TableReference'): + import apache_beam.io.gcp.bigquery_tools as bigquery from apache_beam.io.iobase import WatermarkEstimator from apache_beam.io.restriction_trackers import OffsetRange from apache_beam.io.restriction_trackers import OffsetRestrictionTracker @@ -448,14 +457,35 @@ def _get_bq_timestamp(self) -> Timestamp: Uses BQ's CURRENT_TIMESTAMP instead of the local clock to avoid data loss from clock skew between the worker VM and BigQuery. """ - request = bigquery.BigqueryJobsQueryRequest( - projectId=self._project, - queryRequest=bigquery.QueryRequest( - query='SELECT UNIX_MICROS(CURRENT_TIMESTAMP()) AS ts', - useLegacySql=False, - location=self._location)) - response = self._bq_wrapper.client.jobs.Query(request) - return Timestamp(micros=int(response.rows[0].f[0].v.string_value)) + if self._bq_wrapper._is_modern_client: + query = 'SELECT UNIX_MICROS(CURRENT_TIMESTAMP()) AS ts' + for rows, _ in self._bq_wrapper.run_query( + self._project, + query, + use_legacy_sql=False, + flatten_results=False, + priority='INTERACTIVE'): + for row in rows: + if isinstance(row, dict): + return Timestamp(micros=int(row['ts'])) + elif hasattr(row, 'get'): + return Timestamp(micros=int(row.get('ts'))) + elif hasattr(row, 'values'): + return Timestamp(micros=int(list(row.values())[0])) + elif hasattr(row, 'f'): + return Timestamp(micros=int(row.f[0].v.string_value)) + else: + return Timestamp(micros=int(row[0])) + raise RuntimeError('Failed to get BQ timestamp') + else: + request = bigquery.BigqueryJobsQueryRequest( + projectId=self._project, + queryRequest=bigquery.QueryRequest( + query='SELECT UNIX_MICROS(CURRENT_TIMESTAMP()) AS ts', + useLegacySql=False, + location=self._location)) + response = self._bq_wrapper.client.jobs.Query(request) + return Timestamp(micros=int(response.rows[0].f[0].v.string_value)) def initial_restriction(self, element: _PollConfig) -> OffsetRange: return OffsetRange(0, sys.maxsize) @@ -627,27 +657,40 @@ def process(self, qr: _QueryRange) -> Iterable[_QueryResult]: datasetId=self._temp_dataset, tableId=temp_table_id) - reference = bigquery.JobReference( - jobId=job_id, projectId=self._project, location=self._location) - - request = bigquery.BigqueryJobsInsertRequest( - projectId=self._project, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( - query=bigquery.JobConfigurationQuery( - query=sql, - useLegacySql=False, - destinationTable=temp_table_ref, - writeDisposition='WRITE_TRUNCATE', - ), - ), - jobReference=reference)) - - _LOGGER.info('[Query] Submitting BQ job %s...', job_id) - response = self._bq_wrapper._start_job(request) - _LOGGER.info('[Query] BQ job %s submitted, waiting...', job_id) - self._bq_wrapper.wait_for_bq_job( - response.jobReference, sleep_duration_sec=2) + if self._bq_wrapper._is_modern_client: + _LOGGER.info('[Query] Submitting BQ job %s...', job_id) + job = self._bq_wrapper._start_query_job( + self._project, + sql, + use_legacy_sql=False, + flatten_results=False, + job_id=job_id, + priority='INTERACTIVE', + destination_table=temp_table_ref) + _LOGGER.info('[Query] BQ job %s submitted, waiting...', job_id) + self._bq_wrapper.wait_for_bq_job(job.jobReference, sleep_duration_sec=2) + else: + reference = bigquery.JobReference( + jobId=job_id, projectId=self._project, location=self._location) + + request = bigquery.BigqueryJobsInsertRequest( + projectId=self._project, + job=bigquery.Job( + configuration=bigquery.JobConfiguration( + query=bigquery.JobConfigurationQuery( + query=sql, + useLegacySql=False, + destinationTable=temp_table_ref, + writeDisposition='WRITE_TRUNCATE', + ), + ), + jobReference=reference)) + + _LOGGER.info('[Query] Submitting BQ job %s...', job_id) + response = self._bq_wrapper._start_job(request) + _LOGGER.info('[Query] BQ job %s submitted, waiting...', job_id) + self._bq_wrapper.wait_for_bq_job( + response.jobReference, sleep_duration_sec=2) _LOGGER.info( '[Query] BQ job %s DONE. Results in %s.%s', job_id, diff --git a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py index dabe80c5c27e..2ea3803a640d 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_file_loads.py @@ -48,12 +48,17 @@ from apache_beam.transforms.display import DisplayDataItem from apache_beam.transforms.window import GlobalWindows -# Protect against environments where bigquery library is not available. -# pylint: disable=wrong-import-order, wrong-import-position +try: + from google.api_core.exceptions import GoogleAPICallError + from google.api_core.exceptions import NotFound +except ImportError: + GoogleAPICallError = None + NotFound = None + try: from apitools.base.py.exceptions import HttpError except ImportError: - pass + HttpError = None _LOGGER = logging.getLogger(__name__) @@ -94,15 +99,36 @@ def _add_destination_partitioning_load_parameters( return additional_parameters additional_parameters = dict(additional_parameters) - time_partitioning = getattr(destination_table, 'timePartitioning', None) - range_partitioning = getattr(destination_table, 'rangePartitioning', None) + time_partitioning = ( + getattr(destination_table, 'timePartitioning', None) or + getattr(destination_table, 'time_partitioning', None)) + range_partitioning = ( + getattr(destination_table, 'rangePartitioning', None) or + getattr(destination_table, 'range_partitioning', None)) + + time_partitioning_cls = tuple( + cls for cls in ( + getattr(bigquery_tools.bigquery, 'TimePartitioning', None), + getattr( + getattr(bigquery_tools, 'gcp_bigquery', None), 'TimePartitioning', + None), ) if cls is not None) + range_partitioning_cls = tuple( + cls for cls in ( + getattr(bigquery_tools.bigquery, 'RangePartitioning', None), + getattr( + getattr(bigquery_tools, 'gcp_bigquery', None), + 'RangePartitioning', None), ) if cls is not None) if ('timePartitioning' not in additional_parameters and - isinstance(time_partitioning, bigquery_tools.bigquery.TimePartitioning)): + 'time_partitioning' not in additional_parameters and + time_partitioning is not None and + isinstance(time_partitioning, time_partitioning_cls)): additional_parameters['timePartitioning'] = time_partitioning - if ('rangePartitioning' not in additional_parameters and isinstance( - range_partitioning, bigquery_tools.bigquery.RangePartitioning)): + if ('rangePartitioning' not in additional_parameters and + 'range_partitioning' not in additional_parameters and + range_partitioning is not None and + isinstance(range_partitioning, range_partitioning_cls)): additional_parameters['rangePartitioning'] = range_partitioning return additional_parameters @@ -433,8 +459,10 @@ def process(self, element, schema_mod_job_name_prefix): project_id=table_reference.projectId, dataset_id=table_reference.datasetId, table_id=table_reference.tableId) - except HttpError as exn: - if exn.status_code == 404: + except Exception as exn: + if (getattr(exn, 'status_code', None) == 404 or + getattr(exn, 'code', None) == 404 or + (NotFound is not None and isinstance(exn, NotFound))): # Destination table does not exist, so no need to modify its schema # ahead of the copy jobs. return @@ -445,7 +473,14 @@ def process(self, element, schema_mod_job_name_prefix): project=temp_table_load_job_reference.projectId, job_id=temp_table_load_job_reference.jobId, location=temp_table_load_job_reference.location) - temp_table_schema = temp_table_load_job.configuration.load.schema + temp_table_schema = ( + getattr( + getattr( + getattr(temp_table_load_job, 'configuration', None), + 'load', + None), + 'schema', + None) or getattr(temp_table_load_job, 'schema', None)) if bigquery_tools.check_schema_equal(temp_table_schema, destination_table.schema, @@ -797,8 +832,13 @@ def process( elif destination_table is not None: destination_schema = getattr(destination_table, 'schema', None) if isinstance(destination_schema, - bigquery_tools.bigquery.TableSchema): - schema = bigquery_tools.table_schema_to_dict(destination_schema) + (bigquery_tools.TableSchema, list, tuple)): + schema = bigquery_tools.get_dict_table_schema(destination_schema) + self.schema_cache[hashed_dest] = schema + elif destination_schema is not None and hasattr( + destination_schema, + 'fields') and not hasattr(destination_schema, '_mock_children'): + schema = bigquery_tools.get_dict_table_schema(destination_schema) self.schema_cache[hashed_dest] = schema else: _LOGGER.warning( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py index 136b3cc56b7e..e06df43e8d98 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py @@ -51,12 +51,8 @@ if TYPE_CHECKING: from apache_beam.io.gcp.bigquery import ReadFromBigQueryRequest -try: - from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference - from apache_beam.io.gcp.internal.clients.bigquery import TableReference -except ImportError: - DatasetReference = None - TableReference = None +from apache_beam.io.gcp.bigquery_tools import DatasetReference +from apache_beam.io.gcp.bigquery_tools import TableReference _LOGGER = logging.getLogger(__name__) @@ -230,8 +226,10 @@ def display_data(self): def _get_temp_dataset_id(self): if self.temp_dataset is None: return None - elif isinstance(self.temp_dataset, DatasetReference): + elif hasattr(self.temp_dataset, 'datasetId'): return self.temp_dataset.datasetId + elif hasattr(self.temp_dataset, 'dataset_id'): + return self.temp_dataset.dataset_id elif isinstance(self.temp_dataset, str): return self.temp_dataset else: @@ -239,12 +237,14 @@ def _get_temp_dataset_id(self): def _get_temp_dataset_project(self): """Returns the project ID for temporary dataset operations. - + If temp_dataset is a DatasetReference, returns its projectId. Otherwise, returns the pipeline project for billing. """ - if isinstance(self.temp_dataset, DatasetReference): + if hasattr(self.temp_dataset, 'projectId') and self.temp_dataset.projectId: return self.temp_dataset.projectId + elif hasattr(self.temp_dataset, 'project') and self.temp_dataset.project: + return self.temp_dataset.project else: return self._get_project() diff --git a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py index d3d608b1fc6f..ca20c818873f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_schema_tools.py @@ -22,6 +22,7 @@ """ import datetime +import decimal from typing import Optional from typing import Sequence @@ -32,10 +33,18 @@ import apache_beam.typehints.schemas import apache_beam.utils.proto_utils import apache_beam.utils.timestamp -from apache_beam.io.gcp.internal.clients import bigquery +from apache_beam.io.gcp.bigquery_tools import TableSchema from apache_beam.portability.api import schema_pb2 from apache_beam.transforms import DoFn +try: + from apache_beam.io.gcp.internal.clients import bigquery +except ImportError: + bigquery = None + +if bigquery is None or not hasattr(bigquery, 'TableReference'): + import apache_beam.io.gcp.bigquery_tools as bigquery + # BigQuery types as listed in # https://cloud.google.com/bigquery/docs/reference/standard-sql/data-types # with aliases (RECORD, BOOLEAN, FLOAT, INTEGER) as defined in @@ -43,11 +52,15 @@ BIG_QUERY_TO_PYTHON_TYPES = { "STRING": str, "INTEGER": np.int64, + "INT64": np.int64, "FLOAT64": np.float64, "FLOAT": np.float64, "BOOLEAN": bool, + "BOOL": bool, "BYTES": bytes, "TIMESTAMP": apache_beam.utils.timestamp.Timestamp, + "NUMERIC": decimal.Decimal, + "BIGNUMERIC": decimal.Decimal, "GEOGRAPHY": str, #TODO(https://github.com/apache/beam/issues/20810): # Finish mappings for all BQ types @@ -71,21 +84,36 @@ def generate_user_type_from_bq_schema( Returns: type: type that can be used to work with pCollections. """ - effective_types = {**BIG_QUERY_TO_PYTHON_TYPES, **(type_overrides or {})} + normalized_overrides = { + k.upper(): v + for k, v in (type_overrides or {}).items() + } + effective_types = {**BIG_QUERY_TO_PYTHON_TYPES, **normalized_overrides} the_schema = beam.io.gcp.bigquery_tools.get_dict_table_schema( the_table_schema) - if the_schema == {}: + if the_schema is None: raise ValueError("Encountered an empty schema") field_names_and_types = [] - for field in the_schema['fields']: + for field in the_schema.get('fields', []): if selected_fields is not None and field['name'] not in selected_fields: continue - if field['type'] in effective_types: - typ = bq_field_to_type(field['type'], field['mode'], type_overrides) + field_type = field['type'].upper() + field_mode = (field.get('mode') or 'NULLABLE').upper() + if field_type in ('RECORD', 'STRUCT') and 'fields' in field: + nested_type = generate_user_type_from_bq_schema( + {'fields': field['fields']}, type_overrides=normalized_overrides) + if field_mode in ('NULLABLE', ''): + typ = Optional[nested_type] + elif field_mode == 'REPEATED': + typ = Sequence[nested_type] + elif field_mode == 'REQUIRED': + typ = nested_type + else: + raise ValueError(f"Encountered an unsupported mode: {field_mode!r}") + elif field_type in effective_types: + typ = bq_field_to_type(field_type, field_mode, normalized_overrides) else: - raise ValueError( - f"Encountered " - f"an unsupported type: {field['type']!r}") + raise ValueError(f"Encountered an unsupported type: {field['type']!r}") field_names_and_types.append((field['name'], typ)) sample_schema = beam.typehints.schemas.named_fields_to_schema( field_names_and_types) @@ -105,13 +133,19 @@ def bq_field_to_type(field, mode, type_overrides=None): Returns: The corresponding Python type hint. """ - effective_types = {**BIG_QUERY_TO_PYTHON_TYPES, **(type_overrides or {})} - if mode == 'NULLABLE' or mode is None or mode == '': - return Optional[effective_types[field]] - elif mode == 'REPEATED': - return Sequence[effective_types[field]] - elif mode == 'REQUIRED': - return effective_types[field] + normalized_overrides = { + k.upper(): v + for k, v in (type_overrides or {}).items() + } + effective_types = {**BIG_QUERY_TO_PYTHON_TYPES, **normalized_overrides} + field_type = field.upper() + field_mode = (mode or 'NULLABLE').upper() + if field_mode in ('NULLABLE', ''): + return Optional[effective_types[field_type]] + elif field_mode == 'REPEATED': + return Sequence[effective_types[field_type]] + elif field_mode == 'REQUIRED': + return effective_types[field_type] else: raise ValueError(f"Encountered an unsupported mode: {mode!r}") @@ -139,10 +173,13 @@ def __init__(self, pcoll_val_ctor): self._pcoll_val_ctor = pcoll_val_ctor def process(self, dict_of_tuples): + converted = {} for k, v in dict_of_tuples.items(): if isinstance(v, datetime.datetime): - dict_of_tuples[k] = beam.utils.timestamp.Timestamp.from_utc_datetime(v) - yield self._pcoll_val_ctor(**dict_of_tuples) + converted[k] = beam.utils.timestamp.Timestamp.from_utc_datetime(v) + else: + converted[k] = v + yield self._pcoll_val_ctor(**converted) def infer_output_type(self, input_type): return self._pcoll_val_ctor diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index 51d13d96b73a..bd32f4657589 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -448,8 +448,11 @@ def test_temp_dataset_is_configurable( def test_create_temp_dataset_exception(self, exception_type, error_message): # Uses the FnApiRunner to ensure errors are mocked/passed through correctly - with mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, - 'Insert'),\ + with mock.patch.object(beam.io.gcp.bigquery._CustomBigQuerySource, + 'estimate_size'),\ + mock.patch.object(BigQueryWrapper, + 'get_query_location', + return_value='US'),\ mock.patch.object(BigQueryWrapper, 'get_or_create_dataset') as mock_insert, \ mock.patch('time.sleep'), \ @@ -515,14 +518,21 @@ class DummySchema: fields = [] numBytes = 5 + num_bytes = 5 schema = DummySchema() + location = 'US' + table_id = 'table' + tableId = 'table' + dataset_id = 'dataset' + datasetId = 'dataset' + project = 'project' + projectId = 'project' # TODO(https://github.com/apache/beam/issues/34549): This test relies on # lineage metrics which Prism doesn't seem to handle correctly. Defaulting # to FnApiRunner instead. with mock.patch('time.sleep'), \ - mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, - 'Get') as mock_get_table, \ + mock.patch('google.cloud.bigquery.Client.get_table') as mock_get_table, \ mock.patch.object(BigQueryWrapper, 'wait_for_bq_job'), \ mock.patch.object(BigQueryWrapper, @@ -534,7 +544,7 @@ class DummySchema: beam.Pipeline('FnApiRunner') as p: call_counter = 0 - def store_callback(unused_request): + def store_callback(*unused_args, **unused_kwargs): nonlocal call_counter if call_counter < len(responses): exception = responses[call_counter] @@ -630,11 +640,18 @@ class DummySchema: fields = [] numBytes = 5 + num_bytes = 5 schema = DummySchema() + location = 'US' + table_id = 'table' + tableId = 'table' + dataset_id = 'dataset' + datasetId = 'dataset' + project = 'project' + projectId = 'project' with mock.patch('time.sleep'), \ - mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, - 'Get') as mock_get_table, \ + mock.patch('google.cloud.bigquery.Client.get_table') as mock_get_table, \ mock.patch.object(BigQueryWrapper, 'wait_for_bq_job'), \ mock.patch.object(BigQueryWrapper, @@ -647,7 +664,7 @@ class DummySchema: beam.Pipeline() as p: call_counter = 0 - def store_callback(unused_request): + def store_callback(*unused_args, **unused_kwargs): nonlocal call_counter if call_counter < len(responses): exception = responses[call_counter] @@ -688,9 +705,8 @@ def test_query_job_exception(self, exception_type, error_message): 'estimate_size') as mock_estimate,\ mock.patch.object(BigQueryWrapper, 'get_query_location') as mock_query_location,\ - mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, - 'Insert') as mock_query_job,\ - mock.patch.object(bigquery_v2_client.BigqueryV2.DatasetsService, 'Get'), \ + mock.patch('google.cloud.bigquery.Client.query') as mock_query_job,\ + mock.patch('google.cloud.bigquery.Client.get_dataset'), \ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc, \ beam.Pipeline('FnApiRunner') as p: @@ -704,7 +720,7 @@ def test_query_job_exception(self, exception_type, error_message): gcs_location='gs://temp_location') mock_query_job.assert_called() - self.assertIn(error_message, exc.exception.args[0]) + self.assertIn(error_message, str(exc.exception)) @parameterized.expand([ param( @@ -718,9 +734,8 @@ def test_read_export_exception(self, exception_type, error_message): with mock.patch.object(beam.io.gcp.bigquery._CustomBigQuerySource, 'estimate_size') as mock_estimate,\ - mock.patch.object(bigquery_v2_client.BigqueryV2.TablesService, 'Get'),\ - mock.patch.object(bigquery_v2_client.BigqueryV2.JobsService, - 'Insert') as mock_query_job, \ + mock.patch('google.cloud.bigquery.Client.get_table'),\ + mock.patch('google.cloud.bigquery.Client.extract_table') as mock_query_job, \ mock.patch('time.sleep'), \ self.assertRaises(Exception) as exc,\ beam.Pipeline() as p: @@ -743,8 +758,8 @@ def test_read_direct_lineage(self): # to FnApiRunner instead. with mock.patch.object(bigquery_tools.BigQueryWrapper, '_bigquery_client'),\ - mock.patch.object(bq_storage.BigQueryReadClient, - 'create_read_session'),\ + mock.patch.object(bq_storage, + 'BigQueryReadClient'),\ beam.Pipeline('FnApiRunner') as p: _ = p | ReadFromBigQuery( @@ -2207,7 +2222,7 @@ def test_dofn_client_finish_bundle_flush_called(self): # created. fn.process(('project-id:dataset_id.table_id', ({'month': 1}, 'insertid3'))) - self.assertTrue(client.tables.Get.called) + self.assertTrue(client.get_table.called or client.tables.Get.called) # InsertRows not called as batch size is not hit self.assertFalse(client.insert_rows_json.called) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 0d62ec5233c1..a2eb356c7ce8 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -49,13 +49,20 @@ import apache_beam from apache_beam import coders from apache_beam.internal.gcp import auth +from apache_beam.internal.gcp import json_value from apache_beam.internal.gcp.json_value import from_json_value from apache_beam.internal.http_client import get_new_http from apache_beam.internal.metrics.metric import MetricLogger from apache_beam.internal.metrics.metric import ServiceCallMetric from apache_beam.io.gcp import bigquery_avro_tools from apache_beam.io.gcp import resource_identifiers -from apache_beam.io.gcp.internal.clients import bigquery + +try: + from apache_beam.io.gcp.internal.clients import bigquery as apitools_bigquery +except ImportError: + apitools_bigquery = None + +bigquery = apitools_bigquery from apache_beam.metrics import monitoring_infos from apache_beam.metrics.metric import Metrics from apache_beam.options import value_provider @@ -69,16 +76,37 @@ # Protect against environments where bigquery library is not available. try: import regex + from apitools.base.py import extra_types from apitools.base.py.exceptions import HttpError from apitools.base.py.exceptions import HttpForbiddenError from apitools.base.py.transfer import Upload +except ImportError: + extra_types = None + HttpError = type('HttpError', (Exception, ), {'status_code': None}) + HttpForbiddenError = type( + 'HttpForbiddenError', (Exception, ), {'status_code': 403}) + Upload = None + +try: from google.api_core.client_info import ClientInfo from google.api_core.exceptions import ClientError + from google.api_core.exceptions import Conflict + from google.api_core.exceptions import Forbidden from google.api_core.exceptions import GoogleAPICallError + from google.api_core.exceptions import NotFound + from google.api_core.exceptions import ServerError from google.cloud import bigquery as gcp_bigquery -except Exception: + from google.cloud.bigquery import job as gcp_job +except ImportError: + ClientInfo = None + ClientError = None + Conflict = None + Forbidden = None + GoogleAPICallError = None + NotFound = None + ServerError = None gcp_bigquery = None - pass + gcp_job = None try: from orjson import dumps as fast_json_dumps @@ -87,14 +115,400 @@ fast_json_dumps = json.dumps fast_json_loads = json.loads -# pylint: enable=wrong-import-order, wrong-import-position +# ----------------------------------------------------------------------------- +# Compatibility Models for TableReference, DatasetReference, Schema, and Jobs. + +if bigquery is not None and hasattr(bigquery, 'TableReference'): + TableReference = bigquery.TableReference + DatasetReference = getattr(bigquery, 'DatasetReference', None) + TableFieldSchema = bigquery.TableFieldSchema + TableSchema = bigquery.TableSchema + TableRow = getattr(bigquery, 'TableRow', None) + TableCell = getattr(bigquery, 'TableCell', None) + Table = getattr(bigquery, 'Table', None) + Dataset = getattr(bigquery, 'Dataset', None) + JobReference = getattr(bigquery, 'JobReference', None) + Job = getattr(bigquery, 'Job', None) + JobConfiguration = getattr(bigquery, 'JobConfiguration', None) + JobConfigurationLoad = getattr(bigquery, 'JobConfigurationLoad', None) + JobConfigurationQuery = getattr(bigquery, 'JobConfigurationQuery', None) + JobConfigurationExtract = getattr(bigquery, 'JobConfigurationExtract', None) + JobConfigurationTableCopy = getattr( + bigquery, 'JobConfigurationTableCopy', None) + JobStatistics = getattr(bigquery, 'JobStatistics', None) + JobStatistics2 = getattr(bigquery, 'JobStatistics2', None) + JobStatistics4 = getattr(bigquery, 'JobStatistics4', None) + ErrorProto = getattr(bigquery, 'ErrorProto', None) +else: + + class DatasetReference( + gcp_bigquery.DatasetReference if gcp_bigquery else object): + def __init__( + self, project=None, dataset_id=None, projectId=None, datasetId=None): + p = projectId if projectId is not None else project + d = datasetId if datasetId is not None else dataset_id + if gcp_bigquery: + super().__init__(p or '', d or '') + else: + self._project = p or '' + self._dataset_id = d or '' + + @property + def projectId(self): + return self._project + + @projectId.setter + def projectId(self, val): + self._project = val + + @property + def datasetId(self): + return self._dataset_id + + @datasetId.setter + def datasetId(self, val): + self._dataset_id = val + + class TableReference(gcp_bigquery.TableReference if gcp_bigquery else object): + def __init__( + self, + dataset_ref=None, + table_id=None, + projectId=None, + datasetId=None, + tableId=None, + project=None, + dataset_id=None): + p = projectId if projectId is not None else project + d = datasetId if datasetId is not None else dataset_id + t = tableId if tableId is not None else table_id + if p is not None or d is not None or t is not None: + ds_ref = DatasetReference(p, d) if (p or d) else DatasetReference( + '', '') + if gcp_bigquery: + super().__init__(ds_ref, t or '') + else: + self._project = p + self._dataset_id = d + self._table_id = t + elif dataset_ref is not None: + if gcp_bigquery: + super().__init__(dataset_ref, table_id or '') + else: + self._project = getattr(dataset_ref, 'projectId', None) or getattr( + dataset_ref, 'project', None) + self._dataset_id = getattr(dataset_ref, 'datasetId', None) or getattr( + dataset_ref, 'dataset_id', None) + self._table_id = table_id or '' + else: + if gcp_bigquery: + super().__init__(DatasetReference('', ''), '') + else: + self._project = None + self._dataset_id = None + self._table_id = None + + @property + def projectId(self): + return self.project if hasattr(self, 'project') else getattr( + self, '_project', None) + + @projectId.setter + def projectId(self, val): + self._project = val + + @property + def datasetId(self): + return self.dataset_id if hasattr(self, 'dataset_id') else getattr( + self, '_dataset_id', None) + + @datasetId.setter + def datasetId(self, val): + self._dataset_id = val + + @property + def tableId(self): + return self.table_id if hasattr(self, 'table_id') else getattr( + self, '_table_id', None) + + @tableId.setter + def tableId(self, val): + self._table_id = val + + class TableFieldSchema(gcp_bigquery.SchemaField if gcp_bigquery else object): + def __init__( + self, + name='', + type='STRING', + mode='NULLABLE', + description=None, + fields=(), + field_type=None, + **kwargs): + ft = type or field_type or 'STRING' + if gcp_bigquery: + super().__init__( + name=name, + field_type=ft, + mode=mode or 'NULLABLE', + description=description, + fields=fields or (), + **kwargs) + else: + self.name = name + self.field_type = ft + self.mode = mode or 'NULLABLE' + self.description = description + self.fields = list(fields) if fields else [] + + @property + def type(self): + return self.field_type + + @type.setter + def type(self, val): + self._field_type = val + + class TableSchema(list): + def __init__(self, fields=None): + if fields: + super().__init__(fields) + else: + super().__init__() + + @property + def fields(self): + return self + + @fields.setter + def fields(self, value): + self.clear() + if value: + self.extend(value) + + class JobReference(object): + def __init__( + self, + jobId=None, + projectId=None, + location=None, + job_id=None, + project=None): + self.jobId = jobId if jobId is not None else job_id + self.projectId = projectId if projectId is not None else project + self.location = location + + @property + def job_id(self): + return self.jobId + + @job_id.setter + def job_id(self, val): + self.jobId = val + + @property + def project(self): + return self.projectId + + @project.setter + def project(self, val): + self.projectId = val + + class TableCell(object): + def __init__(self, v=None): + self.v = v + + class TableRow(object): + def __init__(self, f=None): + self.f = f or [] + + Table = None + Dataset = None + Job = None + JobConfiguration = None + JobConfigurationLoad = None + JobConfigurationQuery = None + JobConfigurationExtract = None + JobConfigurationTableCopy = None + JobStatistics = None + JobStatistics2 = None + JobStatistics4 = None + ErrorProto = None + +# Monkey-patch gcp_bigquery classes to ensure full backward compatibility +if gcp_bigquery: + if not hasattr(gcp_bigquery.TableReference, 'projectId'): + gcp_bigquery.TableReference.projectId = property( + lambda self: self.project, + lambda self, val: setattr(self, '_project', val)) + gcp_bigquery.TableReference.datasetId = property( + lambda self: self.dataset_id, + lambda self, val: setattr(self, '_dataset_id', val)) + gcp_bigquery.TableReference.tableId = property( + lambda self: self.table_id, + lambda self, val: setattr(self, '_table_id', val)) + + if not hasattr(gcp_bigquery.DatasetReference, 'projectId'): + gcp_bigquery.DatasetReference.projectId = property( + lambda self: self.project, + lambda self, val: setattr(self, '_project', val)) + gcp_bigquery.DatasetReference.datasetId = property( + lambda self: self.dataset_id, + lambda self, val: setattr(self, '_dataset_id', val)) + + if not hasattr(gcp_bigquery.SchemaField, 'type'): + gcp_bigquery.SchemaField.type = property( + lambda self: self.field_type, + lambda self, val: setattr(self, '_field_type', val)) + + if not hasattr(gcp_bigquery.Table, 'tableReference'): + gcp_bigquery.Table.tableReference = property(lambda self: self.reference) + gcp_bigquery.Table.numRows = property(lambda self: self.num_rows) + gcp_bigquery.Table.numBytes = property(lambda self: self.num_bytes) + gcp_bigquery.Table.timePartitioning = property( + lambda self: self.time_partitioning) + gcp_bigquery.Table.rangePartitioning = property( + lambda self: self.range_partitioning) + + if not hasattr(gcp_bigquery.Dataset, 'datasetReference'): + gcp_bigquery.Dataset.datasetReference = property( + lambda self: self.reference) + gcp_bigquery.Dataset.defaultTableExpirationMs = property( + lambda self: self.default_table_expiration_ms, + lambda self, val: setattr(self, 'default_table_expiration_ms', val)) + + if hasattr(gcp_job, + '_AsyncJob') and not hasattr(gcp_job._AsyncJob, 'jobReference'): + + class _JobStatusCompat: + def __init__(self, job): + self._job = job + + @property + def state(self): + return self._job.state + + @property + def errorResult(self): + return self._job.error_result + + @property + def errors(self): + return self._job.errors + + class _JobStatsCompat: + def __init__(self, job): + self._job = job + + @property + def query(self): + return self + + @property + def totalBytesBilled(self): + return getattr(self._job, 'total_bytes_billed', None) + + @property + def totalBytesProcessed(self): + return getattr(self._job, 'total_bytes_processed', None) + + gcp_job._AsyncJob.jobReference = property( + lambda self: JobReference( + job_id=self.job_id, project=self.project, location=self.location)) + gcp_job._AsyncJob.status = property(lambda self: _JobStatusCompat(self)) + gcp_job._AsyncJob.statistics = property(lambda self: _JobStatsCompat(self)) + + +def _to_gcp_table_ref(table_ref, default_project=None): + if table_ref is None: + return None + if gcp_bigquery is not None and isinstance( + table_ref, getattr(gcp_bigquery, 'TableReference', ())): + return table_ref + if isinstance(table_ref, str): + table_ref = parse_table_reference(table_ref, project=default_project) + proj = getattr(table_ref, 'projectId', None) or getattr( + table_ref, 'project', None) or getattr( + table_ref, 'project_id', None) or default_project or 'default' + dataset_id = getattr(table_ref, 'datasetId', None) or getattr( + table_ref, 'dataset_id', None) or getattr(table_ref, 'dataset', None) + table_id = getattr(table_ref, 'tableId', None) or getattr( + table_ref, 'table_id', None) or getattr(table_ref, 'table', None) + if dataset_id and table_id: + if gcp_bigquery is not None and hasattr( + gcp_bigquery, 'TableReference') and hasattr(gcp_bigquery, + 'DatasetReference'): + return gcp_bigquery.TableReference( + gcp_bigquery.DatasetReference(proj, dataset_id), table_id) + return TableReference( + projectId=proj, datasetId=dataset_id, tableId=table_id) + return table_ref + + +def _to_gcp_dataset_ref(dataset_ref, project=None): + if dataset_ref is None: + return None + if gcp_bigquery is not None and isinstance( + dataset_ref, getattr(gcp_bigquery, 'DatasetReference', ())): + return dataset_ref + if isinstance(dataset_ref, str): + if ':' in dataset_ref or '.' in dataset_ref: + if gcp_bigquery is not None and hasattr(gcp_bigquery.DatasetReference, + 'from_string'): + return gcp_bigquery.DatasetReference.from_string( + dataset_ref, default_project=project) + proj = project or 'default' + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): + return gcp_bigquery.DatasetReference(proj, dataset_ref) + return DatasetReference(projectId=proj, datasetId=dataset_ref) + if hasattr(dataset_ref, 'projectId') and hasattr(dataset_ref, 'datasetId'): + proj = getattr(dataset_ref, 'projectId', None) or project or 'default' + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): + return gcp_bigquery.DatasetReference(proj, dataset_ref.datasetId) + return DatasetReference(projectId=proj, datasetId=dataset_ref.datasetId) + return dataset_ref + + +def _to_gcp_schema(schema): + if schema is None: + return None + if isinstance(schema, (list, tuple)): + fields = [] + for f in schema: + if gcp_bigquery is not None and isinstance( + f, getattr(gcp_bigquery, 'SchemaField', ())): + fields.append(f) + elif isinstance(f, dict) and gcp_bigquery is not None: + fields.append(gcp_bigquery.SchemaField.from_api_repr(f)) + elif hasattr(f, 'name') and gcp_bigquery is not None: + dict_field = table_field_to_dict(f) + if isinstance(dict_field, dict): + fields.append(gcp_bigquery.SchemaField.from_api_repr(dict_field)) + else: + fields.append(f) + else: + fields.append(f) + return fields + if isinstance(schema, TableSchema) or hasattr(schema, 'fields'): + dict_schema = get_dict_table_schema(schema) + if isinstance(dict_schema, dict) and gcp_bigquery is not None: + return [ + gcp_bigquery.SchemaField.from_api_repr(f) + for f in dict_schema.get('fields', []) + ] + if hasattr(schema, 'fields') and schema.fields is not None: + return list(schema.fields) + if isinstance(schema, dict): + if gcp_bigquery is not None: + return [ + gcp_bigquery.SchemaField.from_api_repr(f) + for f in schema.get('fields', []) + ] + return schema.get('fields', []) + if isinstance(schema, str): + return _to_gcp_schema(get_dict_table_schema(schema)) + return schema -# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports -try: - from apache_beam.io.gcp.internal.clients.bigquery import TableReference -except ImportError: - TableReference = None -# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) @@ -153,6 +567,8 @@ def default_encoder(obj): # on python 3 base64-encoded bytes are decoded to strings # before being sent to BigQuery return obj.decode('utf-8') + elif isinstance(obj, apache_beam.utils.timestamp.Timestamp): + return obj.to_utc_datetime().isoformat() elif isinstance(obj, (datetime.date, datetime.time)): return str(obj) elif isinstance(obj, datetime.datetime): @@ -217,73 +633,49 @@ def parse_table_schema_from_json(schema_string): 'Unable to parse JSON schema: %s - %r' % (schema_string, e)) def _parse_schema_field(field): - """Parse a single schema field from dictionary. - - Args: - field: Dictionary object containing serialized schema. - - Returns: - A TableFieldSchema for a single column in BigQuery. - """ - schema = bigquery.TableFieldSchema() - schema.name = field['name'] - schema.type = field['type'] - if 'mode' in field: - schema.mode = field['mode'] - else: - schema.mode = 'NULLABLE' - if 'description' in field: - schema.description = field['description'] - if 'fields' in field: - schema.fields = [_parse_schema_field(x) for x in field['fields']] - return schema - - fields = [_parse_schema_field(f) for f in json_schema['fields']] - return bigquery.TableSchema(fields=fields) + name = field['name'] + field_type = field.get('type') or field.get('type_') or 'STRING' + mode = field.get('mode', 'NULLABLE') + description = field.get('description', None) + sub_fields = [_parse_schema_field(x) + for x in field.get('fields', [])] if 'fields' in field else () + return TableFieldSchema( + name=name, + type=field_type, + mode=mode, + description=description, + fields=sub_fields) + + fields = [_parse_schema_field(f) for f in json_schema.get('fields', [])] + return TableSchema(fields=fields) def parse_table_reference(table, dataset=None, project=None): - """Parses a table reference into a (project, dataset, table) tuple. + """Parses a table reference into a TableReference instance. Args: - table: The ID of the table. The ID must contain only letters - (a-z, A-Z), numbers (0-9), connectors (-_). If dataset argument is None - then the table argument must contain the entire table reference: - 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. This argument can be a - TableReference instance in which case dataset and project are - ignored and the reference is returned as a result. Additionally, for date - partitioned tables, appending '$YYYYmmdd' to the table name is supported, - e.g. 'DATASET.TABLE$YYYYmmdd'. - dataset: The ID of the dataset containing this table or null if the table - reference is specified entirely by the table argument. - project: The ID of the project containing this table or null if the table - reference is specified entirely by the table (and possibly dataset) - argument. + table: The ID of the table, or full reference string / TableReference. + dataset: The ID of the dataset containing this table. + project: The ID of the project containing this table. Returns: - A TableReference object from the bigquery API. The object has the following - attributes: projectId, datasetId, and tableId. - If the input is a TableReference object, a new object will be returned. - - Raises: - ValueError: if the table reference as a string does not match the expected - format. + A TableReference object. """ - if isinstance(table, TableReference): return TableReference( projectId=table.projectId, datasetId=table.datasetId, tableId=table.tableId) + elif isinstance(table, getattr(gcp_bigquery, 'TableReference', ())): + return TableReference( + projectId=table.project, + datasetId=table.dataset_id, + tableId=table.table_id) elif callable(table): return table elif isinstance(table, value_provider.ValueProvider): return table - table_reference = TableReference() - # If dataset argument is not specified, the expectation is that the - # table argument will contain a full table reference instead of just a - # table name. if dataset is None: pattern = ( f'((?P{_PROJECT_PATTERN})[:\\.])?' @@ -293,14 +685,12 @@ def parse_table_reference(table, dataset=None, project=None): raise ValueError( 'Expected a table reference (PROJECT:DATASET.TABLE or ' 'DATASET.TABLE) instead of %s.' % table) - table_reference.projectId = match.group('project') - table_reference.datasetId = match.group('dataset') - table_reference.tableId = match.group('table') + return TableReference( + projectId=match.group('project'), + datasetId=match.group('dataset'), + tableId=match.group('table')) else: - table_reference.projectId = project - table_reference.datasetId = dataset - table_reference.tableId = table - return table_reference + return TableReference(projectId=project, datasetId=dataset, tableId=table) # ----------------------------------------------------------------------------- @@ -308,31 +698,34 @@ def parse_table_reference(table, dataset=None, project=None): def _build_job_labels(input_labels): - """Builds job label protobuf structure.""" - input_labels = input_labels or {} - result = bigquery.JobConfiguration.LabelsValue() - - for k, v in input_labels.items(): - result.additionalProperties.append( - bigquery.JobConfiguration.LabelsValue.AdditionalProperty( - key=k, - value=v, - )) - return result + """Builds job label dictionary or protobuf structure.""" + if apitools_bigquery is not None and hasattr(apitools_bigquery, + 'JobConfiguration'): + input_labels = input_labels or {} + result = apitools_bigquery.JobConfiguration.LabelsValue() + for k, v in input_labels.items(): + result.additionalProperties.append( + apitools_bigquery.JobConfiguration.LabelsValue.AdditionalProperty( + key=k, + value=v, + )) + return result + return input_labels or {} def _build_dataset_labels(input_labels): - """Builds dataset label protobuf structure.""" - input_labels = input_labels or {} - result = bigquery.Dataset.LabelsValue() - - for k, v in input_labels.items(): - result.additionalProperties.append( - bigquery.Dataset.LabelsValue.AdditionalProperty( - key=k, - value=v, - )) - return result + """Builds dataset label dictionary or protobuf structure.""" + if apitools_bigquery is not None and hasattr(apitools_bigquery, 'Dataset'): + input_labels = input_labels or {} + result = apitools_bigquery.Dataset.LabelsValue() + for k, v in input_labels.items(): + result.additionalProperties.append( + apitools_bigquery.Dataset.LabelsValue.AdditionalProperty( + key=k, + value=v, + )) + return result + return input_labels or {} def _build_filter_from_labels(labels): @@ -343,7 +736,10 @@ def _build_filter_from_labels(labels): def _build_dataset_encryption_config(kms_key): - return bigquery.EncryptionConfiguration(kmsKeyName=kms_key) + if apitools_bigquery is not None and hasattr(apitools_bigquery, + 'EncryptionConfiguration'): + return apitools_bigquery.EncryptionConfiguration(kmsKeyName=kms_key) + return kms_key class BigQueryWrapper(object): @@ -358,18 +754,18 @@ class BigQueryWrapper(object): should not be used in production code. """ - # If updating following names, also update the corresponding pydocs in - # bigquery.py. TEMP_TABLE = 'beam_temp_table_' TEMP_DATASET = 'beam_temp_dataset_' 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__)) + if client is not None: + self.client = client + self.gcp_bq_client = client + else: + self.client = BigQueryWrapper._bigquery_client(PipelineOptions()) + self.gcp_bq_client = self.client self._unique_row_id = 0 # For testing scenarios where we pass in a client we do not want a @@ -402,17 +798,33 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): self.created_temp_dataset = False @property - def unique_row_id(self): - """Returns a unique row ID (str) used to avoid multiple insertions. - - If the row ID is provided, BigQuery will make a best effort to not insert - the same row multiple times for fail and retry scenarios in which the insert - request may be issued several times. This comes into play for sinks executed - in a local runner. + def _is_modern_client(self): + client_cls = getattr(gcp_bigquery, 'Client', None) + if isinstance(client_cls, type) and isinstance(self.client, client_cls): + return True + if hasattr(self.client, '_mock_children') or hasattr(self.client, + '_mock_methods'): + mock_keys = set(getattr(self.client, '__dict__', {}).keys()) | set( + getattr(self.client, '_mock_children', {}).keys()) + for modern_attr in ('insert_rows_json', + 'load_table_from_uri', + 'load_table_from_file', + 'query', + 'create_table', + 'delete_table', + 'list_datasets'): + if modern_attr in mock_keys: + return True + return False + if apitools_bigquery is not None and hasattr(apitools_bigquery, + 'BigqueryV2'): + if isinstance(self.client, apitools_bigquery.BigqueryV2): + return False + return client_cls is not None and isinstance(self.client, client_cls) - Returns: - a unique row ID string - """ + @property + def unique_row_id(self): + """Returns a unique row ID (str) used to avoid multiple insertions.""" self._unique_row_id += 1 return '%s_%d' % (self._row_id_prefix, self._unique_row_id) @@ -426,11 +838,7 @@ def _get_temp_table(self, project_id): project=project_id) def _get_temp_table_project(self, fallback_project_id): - """Returns the project ID for temporary table operations. - - If temp_table_ref exists, returns its projectId. - Otherwise, returns the fallback_project_id. - """ + """Returns the project ID for temporary table operations.""" if self.temp_table_ref: return self.temp_table_ref.projectId else: @@ -445,57 +853,83 @@ def _get_temp_dataset(self): num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_query_location(self, project_id, query, use_legacy_sql): - """ - Get the location of tables referenced in a query. + """Get the location of tables referenced in a query.""" + if self._is_modern_client: + job_config = gcp_bigquery.QueryJobConfig( + dry_run=True, + use_legacy_sql=use_legacy_sql, + ) + try: + response = self.client.query( + query, job_config=job_config, project=project_id) + except (Forbidden, HttpForbiddenError): + return None + except Exception: + raise - This method returns the location of the first available referenced - table for user in the query and depends on the BigQuery service to - provide error handling for queries that reference tables in multiple - locations. - """ - reference = bigquery.JobReference( - jobId=uuid.uuid4().hex, projectId=project_id) - request = bigquery.BigqueryJobsInsertRequest( + if hasattr(response, 'referenced_tables'): + referenced_tables = response.referenced_tables + elif hasattr(response, 'statistics') and response.statistics is not None: + referenced_tables = getattr( + response.statistics.query, 'referencedTables', None) + else: + _LOGGER.warning( + "Unable to get location, missing statistics. Query: %s", query) + return None + + if referenced_tables: + for table in referenced_tables: + try: + p = getattr(table, 'project', None) or getattr( + table, 'projectId', None) + d = getattr(table, 'dataset_id', None) or getattr( + table, 'datasetId', None) + t = getattr(table, 'table_id', None) or getattr( + table, 'tableId', None) + location = self.get_table_location(p, d, t) + except (Forbidden, HttpForbiddenError, ClientError): + continue + if location: + _LOGGER.info( + "Using location %r from table %r referenced by query %s", + location, + table, + query) + return location + + _LOGGER.debug( + "Query %s does not reference any tables or " + "you don't have permission to inspect them.", + query) + return None + + # Fallback if legacy client.jobs.Insert is mocked + reference = JobReference(jobId=uuid.uuid4().hex, projectId=project_id) + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( dryRun=True, - query=bigquery.JobConfigurationQuery( + query=apitools_bigquery.JobConfigurationQuery( query=query, useLegacySql=use_legacy_sql, )), jobReference=reference)) - response = self.client.jobs.Insert(request) - if response.statistics is None: - # This behavior is only expected in tests _LOGGER.warning( "Unable to get location, missing response.statistics. Query: %s", query) return None - referenced_tables = response.statistics.query.referencedTables - if referenced_tables: # Guards against both non-empty and non-None + if referenced_tables: for table in referenced_tables: try: location = self.get_table_location( table.projectId, table.datasetId, table.tableId) - except HttpForbiddenError: - # Permission access for table (i.e. from authorized_view), - # try next one + except (HttpForbiddenError, Forbidden): continue - _LOGGER.info( - "Using location %r from table %r referenced by query %s", - location, - table, - query) return location - - _LOGGER.debug( - "Query %s does not reference any tables or " - "you don't have permission to inspect them.", - query) return None @retry.with_exponential_backoff( @@ -510,14 +944,41 @@ def _insert_copy_job( create_disposition=None, write_disposition=None, job_labels=None): - reference = bigquery.JobReference() - reference.jobId = job_id - reference.projectId = project_id - request = bigquery.BigqueryJobsInsertRequest( + if self._is_modern_client: + job_config = gcp_bigquery.CopyJobConfig( + create_disposition=create_disposition, + write_disposition=write_disposition, + labels=job_labels, + ) + src = from_table_reference if isinstance( + from_table_reference, list) else [from_table_reference] + src_refs = [_to_gcp_table_ref(t, default_project=project_id) for t in src] + dst_ref = _to_gcp_table_ref( + to_table_reference, default_project=project_id) + try: + job = self.client.copy_table( + src_refs, + dst_ref, + job_id=job_id, + job_config=job_config, + project=project_id, + ) + return JobReference( + job_id=job.job_id, project=job.project, location=job.location) + except (Conflict, HttpError) as exn: + if getattr(exn, 'code', None) == 409 or getattr( + exn, 'status_code', None) == 409 or isinstance(exn, Conflict): + _LOGGER.info("BigQuery copy job %s already exists.", job_id) + return JobReference(job_id=job_id, project=project_id) + raise + + # Fallback if legacy client.jobs.Insert is mocked + reference = JobReference(jobId=job_id, projectId=project_id) + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( - copy=bigquery.JobConfigurationTableCopy( + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( + copy=apitools_bigquery.JobConfigurationTableCopy( destinationTable=to_table_reference, sourceTable=from_table_reference, createDisposition=create_disposition, @@ -527,7 +988,6 @@ def _insert_copy_job( ), jobReference=reference, )) - return self._start_job(request).jobReference @retry.with_exponential_backoff( @@ -546,28 +1006,70 @@ def _insert_load_job( additional_load_parameters=None, source_format=None, job_labels=None): - if not source_uris and not source_stream: _LOGGER.warning( 'Both source URIs and source stream are not provided. BigQuery load ' 'job will not load any data.') - if source_uris and source_stream: - raise ValueError( - 'Only one of source_uris and source_stream may be specified. ' - 'Got both.') + if self._is_modern_client: + dst_ref = _to_gcp_table_ref(table_reference, default_project=project_id) + job_schema = None if schema == 'SCHEMA_AUTODETECT' else _to_gcp_schema( + schema) + autodetect = schema == 'SCHEMA_AUTODETECT' + additional_load_parameters = additional_load_parameters or {} + job_config = gcp_bigquery.LoadJobConfig( + schema=job_schema, + autodetect=autodetect, + create_disposition=create_disposition, + write_disposition=write_disposition, + source_format=source_format, + use_avro_logical_types=True, + labels=job_labels, + **additional_load_parameters) + try: + if source_stream: + job = self.client.load_table_from_file( + source_stream, + dst_ref, + job_id=job_id, + job_config=job_config, + project=project_id) + else: + source_uris = source_uris or [] + job = self.client.load_table_from_uri( + source_uris, + dst_ref, + job_id=job_id, + job_config=job_config, + project=project_id) + return JobReference( + job_id=job.job_id, project=job.project, location=job.location) + except (Conflict, HttpError) as exn: + if getattr(exn, 'code', None) == 409 or getattr( + exn, 'status_code', None) == 409 or isinstance(exn, Conflict): + _LOGGER.info("BigQuery load job %s already exists.", job_id) + return JobReference(job_id=job_id, project=project_id) + raise + # Fallback if legacy client.jobs.Insert is mocked if source_uris is None: source_uris = [] - additional_load_parameters = additional_load_parameters or {} - job_schema = None if schema == 'SCHEMA_AUTODETECT' else schema - reference = bigquery.JobReference(jobId=job_id, projectId=project_id) - request = bigquery.BigqueryJobsInsertRequest( + if schema == 'SCHEMA_AUTODETECT': + job_schema = None + elif isinstance(schema, (dict, str)): + job_schema = get_bq_tableschema(schema) + elif isinstance(schema, (list, tuple)): + job_schema = get_bq_tableschema( + {'fields': [table_field_to_dict(f) for f in schema]}) + else: + job_schema = schema + reference = JobReference(jobId=job_id, projectId=project_id) + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( - load=bigquery.JobConfigurationLoad( + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( + load=apitools_bigquery.JobConfigurationLoad( sourceUris=source_uris, destinationTable=table_reference, schema=job_schema, @@ -588,7 +1090,6 @@ def _parse_location_from_exc(content, job_id): """Parse job location from Exception content.""" if isinstance(content, bytes): content = content.decode('ascii', 'replace') - # search for "Already Exists: Job :." m = re.search(r"Already Exists: Job \S+\:(\S+)\." + job_id, content) if not m: _LOGGER.warning( @@ -598,20 +1099,16 @@ def _parse_location_from_exc(content, job_id): def _start_job( self, - request: 'bigquery.BigqueryJobsInsertRequest', + request: 'apitools_bigquery.BigqueryJobsInsertRequest', stream=None, ): - """Inserts a BigQuery job. + """Inserts a BigQuery job for legacy apitools client. If the job exists already, it returns it. - - Args: - request (bigquery.BigqueryJobsInsertRequest): An insert job request. - stream (IO[bytes]): A bytes IO object open for reading. """ try: upload = None - if stream: + if stream and Upload: upload = Upload.FromStream(stream, mime_type=UNKNOWN_MIME_TYPE) response = self.client.jobs.Insert(request, upload=upload) _LOGGER.info( @@ -631,7 +1128,6 @@ def _start_job( job_location = self._parse_location_from_exc(exn.content, jobId) response = request.job if not response.jobReference.location and job_location: - # Request not constructed with location response.jobReference.location = job_location return response else: @@ -652,54 +1148,102 @@ def _start_query_job( priority, dry_run=False, kms_key=None, - job_labels=None): - reference = bigquery.JobReference(jobId=job_id, projectId=project_id) - request = bigquery.BigqueryJobsInsertRequest( + job_labels=None, + destination_table=None): + if self._is_modern_client: + dest_table = None + if not dry_run: + dest_table = _to_gcp_table_ref( + destination_table or + self._get_temp_table(self._get_temp_table_project(project_id)), + default_project=project_id or getattr(self.client, 'project', None)) + + job_config = gcp_bigquery.QueryJobConfig( + dry_run=dry_run, + use_legacy_sql=use_legacy_sql, + allow_large_results=not dry_run, + destination=dest_table, + flatten_results=flatten_results, + priority=priority, + labels=job_labels, + ) + if kms_key: + job_config.destination_encryption_configuration = ( + gcp_bigquery.EncryptionConfiguration(kms_key_name=kms_key)) + + try: + job = self.client.query( + query, + job_config=job_config, + job_id=job_id, + project=project_id, + job_retry=None, + ) + return job + except (Conflict, HttpError) as exn: + if getattr(exn, 'code', None) == 409 or getattr( + exn, 'status_code', None) == 409 or isinstance(exn, Conflict): + return self.get_job(project_id, job_id) + raise + + # Fallback if legacy client.jobs.Insert is mocked + reference = JobReference(jobId=job_id, projectId=project_id) + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( dryRun=dry_run, - query=bigquery.JobConfigurationQuery( + query=apitools_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, + destinationTable=( + destination_table if destination_table is not None else + ( + self._get_temp_table( + self._get_temp_table_project(project_id)) + if not dry_run else None)), flattenResults=flatten_results, priority=priority, - destinationEncryptionConfiguration=bigquery. + destinationEncryptionConfiguration=apitools_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. - - 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. - - Raises: - `RuntimeError`: If the job is FAILED or the number of retries has been - reached. - """ + """Poll job until it is DONE.""" retry = 0 + project = getattr(job_reference, 'projectId', None) or getattr( + job_reference, 'project', None) + job_id = getattr(job_reference, 'jobId', None) or getattr( + job_reference, 'job_id', 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, job_id, location) + status_obj = getattr(job, 'status', None) + if status_obj is not None and hasattr(status_obj, 'state') and isinstance( + status_obj.state, str): + state = status_obj.state + error_result = getattr(status_obj, 'errorResult', None) + elif hasattr(job, 'state') and isinstance(job.state, str): + state = job.state + error_result = getattr(job, 'error_result', None) + else: + state = getattr(job, 'state', None) or getattr( + status_obj, 'state', None) + error_result = getattr(job, 'error_result', None) or getattr( + status_obj, 'errorResult', None) + + _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_id, error_result)) + elif state == 'DONE': return True else: time.sleep(sleep_duration_sec) @@ -716,14 +1260,18 @@ def _get_query_results( page_token=None, max_results=10000, location=None): - request = bigquery.BigqueryJobsGetQueryResultsRequest( - jobId=job_id, - pageToken=page_token, - projectId=project_id, - maxResults=max_results, - location=location) - response = self.client.jobs.GetQueryResults(request) - return response + if hasattr(self.client, 'jobs') and hasattr(self.client.jobs, + 'GetQueryResults'): + request = apitools_bigquery.BigqueryJobsGetQueryResultsRequest( + jobId=job_id, + pageToken=page_token, + projectId=project_id, + maxResults=max_results, + location=location) + return self.client.jobs.GetQueryResults(request) + + job = self.client.get_job(job_id, project=project_id, location=location) + return job.result(max_results=max_results, page_token=page_token) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, @@ -737,22 +1285,12 @@ def _insert_all_rows( insert_ids, skip_invalid_rows=False, ignore_unknown_values=False): - """Calls the insertAll BigQuery API endpoint. - - Docs for this BQ call: https://cloud.google.com/bigquery/docs/reference\ - /rest/v2/tabledata/insertAll.""" - # The rows argument is a list of - # bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as - # required by the InsertAll() method. + """Calls the insertAll BigQuery API endpoint.""" resource = resource_identifiers.BigQueryTable( project_id, dataset_id, table_id) labels = { - # TODO(ajamato): Add Ptransform label. monitoring_infos.SERVICE_LABEL: 'BigQuery', - # Refer to any method which writes elements to BigQuery in batches - # as "BigQueryBatchWrite". I.e. storage API's insertAll, or future - # APIs introduced. monitoring_infos.METHOD_LABEL: 'BigQueryBatchWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGQUERY_PROJECT_ID_LABEL: project_id, @@ -763,31 +1301,93 @@ def _insert_all_rows( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) + if self._is_modern_client: + started_millis = int(time.time() * 1000) + try: + table_ref_str = '%s.%s.%s' % (project_id, dataset_id, table_id) + row_ids_arg = insert_ids + if row_ids_arg is not None and all(x is None for x in row_ids_arg): + auto_uuid = getattr( + getattr(gcp_bigquery, 'AutoRowIDs', None), 'GENERATE_UUID', None) + if auto_uuid is not None: + row_ids_arg = auto_uuid + errors = self.gcp_bq_client.insert_rows_json( + table_ref_str, + json_rows=rows, + row_ids=row_ids_arg, + skip_invalid_rows=skip_invalid_rows, + ignore_unknown_values=ignore_unknown_values, + timeout=BQ_STREAMING_INSERT_TIMEOUT_SEC) + if not errors: + service_call_metric.call('ok') + else: + for insert_error in errors: + for err in insert_error.get('errors', []): + reason = err.get('reason') if isinstance(err, dict) else getattr( + err, 'reason', None) + service_call_metric.call(reason or 'unknown') + except (ClientError, GoogleAPICallError, HttpError) as e: + status_code = getattr(e, 'code', None) or getattr( + e, 'status_code', None) or 500 + service_call_metric.call(status_code) + reason = None + if hasattr(e, 'response') and getattr(e.response, 'reason', None): + reason = e.response.reason + elif hasattr(e, 'reason') and e.reason: + reason = e.reason + elif hasattr(e, 'errors') and e.errors and isinstance( + e.errors, (list, tuple)) and isinstance(e.errors[0], dict): + reason = e.errors[0].get('reason') + if not reason: + reason = e.__class__.__name__ + errors = [{ + 'index': i, 'errors': [{ + 'reason': reason, 'message': str(e) + }] + } for i in range(len(rows))] + if not errors: + errors = [{ + 'index': 0, 'errors': [{ + 'reason': reason, 'message': str(e) + }] + }] + finally: + self._latency_histogram_metric.update( + int(time.time() * 1000) - started_millis) + return not errors, errors + + # Legacy apitools path + row_list = [] + for row, insert_id in zip(rows, insert_ids): + row_list.append( + apitools_bigquery.TableDataInsertAllRequest.RowsValueListEntry( + insertId=insert_id, + json=json_value.to_json_value(row, True), + )) + request = apitools_bigquery.BigqueryTabledataInsertAllRequest( + projectId=project_id, + datasetId=dataset_id, + tableId=table_id, + tableDataInsertAllRequest=apitools_bigquery.TableDataInsertAllRequest( + rows=row_list, + skipInvalidRows=skip_invalid_rows, + ignoreUnknownValues=ignore_unknown_values, + )) started_millis = int(time.time() * 1000) try: - table_ref_str = '%s.%s.%s' % (project_id, dataset_id, table_id) - errors = self.gcp_bq_client.insert_rows_json( - table_ref_str, - json_rows=rows, - row_ids=insert_ids, - skip_invalid_rows=skip_invalid_rows, - ignore_unknown_values=ignore_unknown_values, - timeout=BQ_STREAMING_INSERT_TIMEOUT_SEC) + response = self.client.tabledata.InsertAll(request) + errors = [ + json.loads(extra_types.RpcError(e).to_json()) + for e in response.insertErrors + ] if not errors: service_call_metric.call('ok') else: - for insert_error in errors: - service_call_metric.call(insert_error['errors'][0]) - except (ClientError, GoogleAPICallError) as e: - # e.code contains the numeric http status code. - service_call_metric.call(e.code) - # Package exception with required fields - error = {'message': e.message, 'reason': e.response.reason} - # Add all rows to the errors list along with the error - errors = [{"index": i, "errors": [error]} for i, _ in enumerate(rows)] + for insert_error in response.insertErrors: + for error in insert_error.errors: + service_call_metric.call(error.reason) except HttpError as e: service_call_metric.call(e) - # Re-raise the exception so that we re-try appropriately. raise finally: self._latency_histogram_metric.update( @@ -798,23 +1398,16 @@ 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( + """Lookup a table's metadata object.""" + if self._is_modern_client: + table_ref = TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id) + return self.client.get_table(table_ref) + + # Fallback for legacy client + request = apitools_bigquery.BigqueryTablesGetRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) - response = self.client.tables.Get(request) - return response + return self.client.tables.Get(request) def _create_table( self, @@ -823,7 +1416,6 @@ def _create_table( table_id, schema, additional_parameters=None): - valid_tablename = regex.fullmatch(_TABLE_PATTERN, table_id, regex.ASCII) if not valid_tablename: raise ValueError( @@ -831,17 +1423,29 @@ def _create_table( 'See https://cloud.google.com/bigquery/docs/tables#table_naming' % table_id) + if self._is_modern_client: + table_ref = TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id) + gcp_schema = _to_gcp_schema(schema) + table = gcp_bigquery.Table(table_ref, schema=gcp_schema) + if additional_parameters: + for k, v in additional_parameters.items(): + setattr(table, k, v) + response = self.client.create_table(table) + _LOGGER.debug("Created the table with id %s", table_id) + return response + + # Fallback for legacy client additional_parameters = additional_parameters or {} - table = bigquery.Table( + table = apitools_bigquery.Table( tableReference=TableReference( projectId=project_id, datasetId=dataset_id, tableId=table_id), schema=schema, **additional_parameters) - request = bigquery.BigqueryTablesInsertRequest( + request = apitools_bigquery.BigqueryTablesInsertRequest( projectId=project_id, datasetId=dataset_id, table=table) response = self.client.tables.Insert(request) _LOGGER.debug("Created the table with id %s", table_id) - # The response is a bigquery.Table instance. return response @retry.with_exponential_backoff( @@ -855,10 +1459,40 @@ def get_or_create_dataset( labels=None, kms_key=None, default_table_expiration_ms=None): - # Check if dataset already exists otherwise create it + if self._is_modern_client: + dataset_ref = DatasetReference(projectId=project_id, datasetId=dataset_id) + try: + dataset = self.client.get_dataset(dataset_ref) + self.created_temp_dataset = False + return dataset + except (NotFound, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 404 or getattr( + exn, 'status_code', None) == 404 or isinstance(exn, NotFound): + _LOGGER.info( + 'Dataset %s:%s does not exist so we will create it as temporary ' + 'with location=%s', + project_id, + dataset_id, + location) + dataset = gcp_bigquery.Dataset(dataset_ref) + if location is not None: + dataset.location = location + if labels is not None: + dataset.labels = labels + if kms_key is not None: + dataset.default_encryption_configuration = ( + gcp_bigquery.EncryptionConfiguration(kms_key_name=kms_key)) + if default_table_expiration_ms is not None: + dataset.default_table_expiration_ms = default_table_expiration_ms + response = self.client.create_dataset(dataset) + self.created_temp_dataset = True + return response + raise + + # Fallback for legacy client try: dataset = self.client.datasets.Get( - bigquery.BigqueryDatasetsGetRequest( + apitools_bigquery.BigqueryDatasetsGetRequest( projectId=project_id, datasetId=dataset_id)) self.created_temp_dataset = False return dataset @@ -870,9 +1504,9 @@ def get_or_create_dataset( project_id, dataset_id, location) - dataset_reference = bigquery.DatasetReference( + dataset_reference = DatasetReference( projectId=project_id, datasetId=dataset_id) - dataset = bigquery.Dataset(datasetReference=dataset_reference) + dataset = apitools_bigquery.Dataset(datasetReference=dataset_reference) if location is not None: dataset.location = location if labels is not None: @@ -882,11 +1516,10 @@ def get_or_create_dataset( _build_dataset_encryption_config(kms_key)) if default_table_expiration_ms is not None: dataset.defaultTableExpirationMs = default_table_expiration_ms - request = bigquery.BigqueryDatasetsInsertRequest( + request = apitools_bigquery.BigqueryDatasetsInsertRequest( projectId=project_id, dataset=dataset) response = self.client.datasets.Insert(request) self.created_temp_dataset = True - # The response is a bigquery.Dataset instance. return response else: raise @@ -895,20 +1528,43 @@ 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( + if self._is_modern_client: + table_ref = TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id) + rows = self.client.list_rows(table_ref, max_results=1) + if hasattr(rows, 'total_rows') and rows.total_rows is not None: + return rows.total_rows == 0 + return len(list(rows)) == 0 + + # Fallback for legacy client + request = apitools_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 @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( + if self._is_modern_client: + table_ref = TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id) + try: + self.client.delete_table(table_ref, not_found_ok=True) + except (NotFound, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 404 or getattr( + exn, 'status_code', None) == 404 or isinstance(exn, NotFound): + _LOGGER.warning( + 'Table %s:%s.%s does not exist', project_id, dataset_id, table_id) + return + raise + return + + # Fallback for legacy client + request = apitools_bigquery.BigqueryTablesDeleteRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) try: self.client.tables.Delete(request) @@ -924,7 +1580,22 @@ def _delete_table(self, project_id, dataset_id, table_id): 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( + if self._is_modern_client: + dataset_ref = DatasetReference(projectId=project_id, datasetId=dataset_id) + try: + self.client.delete_dataset( + dataset_ref, delete_contents=delete_contents, not_found_ok=True) + except (NotFound, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 404 or getattr( + exn, 'status_code', None) == 404 or isinstance(exn, NotFound): + _LOGGER.warning( + 'Dataset %s:%s does not exist', project_id, dataset_id) + return + raise + return + + # Fallback for legacy client + request = apitools_bigquery.BigqueryDatasetsDeleteRequest( projectId=project_id, datasetId=dataset_id, deleteContents=delete_contents) @@ -944,8 +1615,8 @@ def get_table_location(self, project_id, dataset_id, table_id): table = self.get_table(project_id, dataset_id, table_id) return table.location - # Returns true if the temporary dataset was provided by the user. def is_user_configured_dataset(self): + """Returns true if the temporary dataset was provided by the user.""" return ( self.temp_dataset_id and not self.temp_dataset_id.startswith(self.TEMP_DATASET)) @@ -965,8 +1636,6 @@ def create_temporary_dataset( if (project_id is not None and not self.is_user_configured_dataset() and not self.created_temp_dataset): - # Unittests don't pass projectIds so they can be run without error - # User configured datasets are allowed to pre-exist. raise RuntimeError( 'Dataset %s:%s already exists so cannot be used as temporary.' % (project_id, self.temp_dataset_id)) @@ -976,9 +1645,40 @@ def create_temporary_dataset( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def clean_up_temporary_dataset(self, project_id): temp_table = self._get_temp_table(project_id) + if self._is_modern_client: + dataset_ref = DatasetReference( + projectId=project_id, datasetId=temp_table.datasetId) + try: + self.client.get_dataset(dataset_ref) + except (NotFound, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 404 or getattr( + exn, 'status_code', None) == 404 or isinstance(exn, NotFound): + _LOGGER.warning( + 'Dataset %s:%s does not exist', project_id, temp_table.datasetId) + return + raise + try: + if not self.is_user_configured_dataset(): + self._delete_dataset(temp_table.projectId, temp_table.datasetId, True) + else: + self._delete_table( + temp_table.projectId, temp_table.datasetId, temp_table.tableId) + self.created_temp_dataset = False + except (Forbidden, HttpForbiddenError, HttpError) as exn: + if getattr(exn, 'code', None) == 403 or getattr( + exn, 'status_code', None) == 403 or isinstance(exn, Forbidden): + _LOGGER.warning( + 'Permission denied to delete temporary dataset %s:%s for clean up', + temp_table.projectId, + temp_table.datasetId) + return + raise + return + + # Fallback for legacy client try: self.client.datasets.Get( - bigquery.BigqueryDatasetsGetRequest( + apitools_bigquery.BigqueryDatasetsGetRequest( projectId=project_id, datasetId=temp_table.datasetId)) except HttpError as exn: if exn.status_code == 404: @@ -988,8 +1688,6 @@ def clean_up_temporary_dataset(self, project_id): else: raise try: - # We do not want to delete temporary datasets configured by the user hence - # we just delete the temporary table in that case. if not self.is_user_configured_dataset(): self._delete_dataset(temp_table.projectId, temp_table.datasetId, True) else: @@ -1013,31 +1711,50 @@ def _clean_up_beam_labelled_temporary_datasets( self, project_id, dataset_id=None, table_id=None, labels=None): if isinstance(labels, dict): filter_str = _build_filter_from_labels(labels) + else: + filter_str = '' if not self.is_user_configured_dataset() and labels is not None: + if self._is_modern_client: + try: + for dataset in self.client.list_datasets(project=project_id, + filter=filter_str): + ds_id = dataset.dataset_id + self._delete_dataset(project_id, ds_id, True) + except (Forbidden, HttpForbiddenError, HttpError) as exn: + if getattr(exn, 'code', None) == 403 or getattr( + exn, 'status_code', None) == 403 or isinstance(exn, Forbidden): + _LOGGER.warning( + 'Permission denied to delete temporary dataset for clean up.') + return + raise + return + + # Fallback for legacy client response = ( self.client.datasets.List( - bigquery.BigqueryDatasetsListRequest( + apitools_bigquery.BigqueryDatasetsListRequest( projectId=project_id, filter=filter_str))) for dataset in response.datasets: try: - dataset_id = dataset.datasetReference.datasetId - self._delete_dataset(project_id, dataset_id, True) - except HttpError as exn: + ds_id = dataset.datasetReference.datasetId + self._delete_dataset(project_id, ds_id, True) + except (HttpForbiddenError, 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) + ds_id) return else: raise else: try: self._delete_table(project_id, dataset_id, table_id) - except HttpError as exn: - if exn.status_code == 403: + except (Forbidden, HttpForbiddenError, HttpError) as exn: + if getattr(exn, 'code', None) == 403 or getattr( + exn, 'status_code', None) == 403 or isinstance(exn, Forbidden): _LOGGER.warning( 'Permission denied to delete temporary table %s:%s.%s for ' 'clean up.', @@ -1052,11 +1769,14 @@ 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() + if self._is_modern_client: + return self.client.get_job(job_id, project=project, location=location) + + # Fallback for legacy client + request = apitools_bigquery.BigqueryJobsGetRequest() request.jobId = job_id request.projectId = project request.location = location - return self.client.jobs.Get(request) def perform_load_job( @@ -1075,10 +1795,16 @@ def perform_load_job( """Starts a job to load data into BigQuery. Returns: - bigquery.JobReference with the information about the job that was started. + JobReference with the information about the job that was started. """ + if source_uris and source_stream: + raise ValueError( + 'Only one of source_uris and source_stream may be specified. ' + 'Got both.') + project_id = ( - destination.projectId + getattr(destination, 'projectId', None) or + getattr(destination, 'project', None) if load_job_project_id is None else load_job_project_id) return self._insert_load_job( @@ -1111,15 +1837,45 @@ def perform_extract_job( """Starts a job to export data from BigQuery. Returns: - bigquery.JobReference with the information about the job that was started. + 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( + job_project = project or getattr(table_reference, 'projectId', + None) or getattr( + table_reference, 'project', None) + if self._is_modern_client: + src_ref = _to_gcp_table_ref(table_reference, default_project=job_project) + dest_uris = destination if isinstance(destination, + list) else [destination] + job_config = gcp_bigquery.ExtractJobConfig( + destination_format=destination_format, + print_header=include_header, + compression=compression, + use_avro_logical_types=use_avro_logical_types, + labels=job_labels, + ) + try: + job = self.client.extract_table( + src_ref, + dest_uris, + job_id=job_id, + job_config=job_config, + project=job_project, + ) + return JobReference( + job_id=job.job_id, project=job.project, location=job.location) + except (Conflict, HttpError) as exn: + if getattr(exn, 'code', None) == 409 or getattr( + exn, 'status_code', None) == 409 or isinstance(exn, Conflict): + return JobReference(job_id=job_id, project=job_project) + raise + + # Fallback for legacy client + job_reference = JobReference(jobId=job_id, projectId=job_project) + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=job_project, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( - extract=bigquery.JobConfigurationExtract( + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( + extract=apitools_bigquery.JobConfigurationExtract( destinationUris=destination, sourceTable=table_reference, printHeader=include_header, @@ -1146,35 +1902,15 @@ def get_or_create_table( create_disposition, write_disposition, additional_create_parameters=None): - """Gets or creates a table based on create and write dispositions. - - The function mimics the behavior of BigQuery import jobs when using the - same create and write dispositions. - - Args: - project_id: The project id owning the table. - dataset_id: The dataset id owning the table. - table_id: The table id. - schema: A bigquery.TableSchema instance or None. - create_disposition: CREATE_NEVER or CREATE_IF_NEEDED. - write_disposition: WRITE_APPEND, WRITE_EMPTY or WRITE_TRUNCATE. - - Returns: - A bigquery.Table instance if table was found or created. - - Raises: - `RuntimeError`: For various mismatches between the state of the table and - the create/write dispositions passed in. For example if the table is not - empty and WRITE_EMPTY was specified then an error will be raised since - the table was expected to be empty. - """ + """Gets or creates a table based on create and write dispositions.""" from apache_beam.io.gcp.bigquery import BigQueryDisposition found_table = None try: found_table = self.get_table(project_id, dataset_id, table_id) - except HttpError as exn: - if exn.status_code == 404: + except (NotFound, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 404 or getattr( + exn, 'status_code', None) == 404 or isinstance(exn, NotFound): if create_disposition == BigQueryDisposition.CREATE_NEVER: raise RuntimeError( 'Table %s:%s.%s not found but create disposition is CREATE_NEVER.' @@ -1182,12 +1918,8 @@ def get_or_create_table( else: raise - # If table exists already then handle the semantics for WRITE_EMPTY and - # WRITE_TRUNCATE write dispositions. if found_table and write_disposition in ( BigQueryDisposition.WRITE_EMPTY, BigQueryDisposition.WRITE_TRUNCATE): - # Delete the table and recreate it (later) if WRITE_TRUNCATE was - # specified. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: self._delete_table(project_id, dataset_id, table_id) elif (write_disposition == BigQueryDisposition.WRITE_EMPTY and @@ -1196,8 +1928,6 @@ def get_or_create_table( 'Table %s:%s.%s is not empty but write disposition is WRITE_EMPTY.' % (project_id, dataset_id, table_id)) - # Create a new table potentially reusing the schema from a previously - # found table in case the schema was not specified. if schema is None and found_table is None: raise RuntimeError( 'Table %s:%s.%s requires a schema. None can be inferred because the ' @@ -1213,8 +1943,9 @@ def get_or_create_table( table_id=table_id, schema=schema or found_table.schema, additional_parameters=additional_create_parameters) - except HttpError as exn: - if exn.status_code == 409: + except (Conflict, HttpError, ClientError) as exn: + if getattr(exn, 'code', None) == 409 or getattr( + exn, 'status_code', None) == 409 or isinstance(exn, Conflict): _LOGGER.debug( 'Skipping Creation. Table %s:%s.%s already exists.' % (project_id, dataset_id, table_id)) @@ -1227,18 +1958,13 @@ def get_or_create_table( project_id, dataset_id, table_id, - schema or found_table.schema, + schema or (found_table.schema if found_table else None), created_table) - # if write_disposition == BigQueryDisposition.WRITE_TRUNCATE we delete - # the table before this point. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: - # BigQuery can route data to the old table for 2 mins max so wait - # that much time before creating the table and writing it _LOGGER.warning( 'Sleeping for 150 seconds before the write as ' + 'BigQuery inserts can be routed to deleted table ' + 'for 2 mins after the delete and create.') - # TODO(BEAM-2673): Remove this sleep by migrating to load api time.sleep(150) return created_table else: @@ -1262,31 +1988,31 @@ def run_query( priority=priority, dry_run=dry_run, job_labels=job_labels) - job_id = job.jobReference.jobId - location = job.jobReference.location - if dry_run: - # If this was a dry run then the fact that we get here means the - # query has no errors. The start_query_job would raise an error otherwise. return + + job_id = getattr(job, 'job_id', None) or getattr( + getattr(job, 'jobReference', None), 'jobId', None) + location = getattr(job, 'location', None) or getattr( + getattr(job, 'jobReference', None), 'location', None) + page_token = None while True: response = self._get_query_results( project_id, job_id, page_token, location=location) - if not response.jobComplete: - # The jobComplete field can be False if the query request times out - # (default is 10 seconds). Note that this is a timeout for the query - # request not for the actual execution of the query in the service. If - # the request times out we keep trying. This situation is quite possible - # if the query will return a large number of rows. - _LOGGER.info('Waiting on response from query: %s ...', query) - time.sleep(1.0) - continue - # We got some results. The last page is signalled by a missing pageToken. - yield response.rows, response.schema - if not response.pageToken: + if hasattr(response, 'jobComplete'): + if not response.jobComplete: + _LOGGER.info('Waiting on response from query: %s ...', query) + time.sleep(1.0) + continue + yield response.rows, response.schema + if not response.pageToken: + break + page_token = response.pageToken + else: + # Modern RowIterator + yield list(response), TableSchema(fields=response.schema) break - page_token = response.pageToken def insert_rows( self, @@ -1297,32 +2023,7 @@ def insert_rows( insert_ids=None, skip_invalid_rows=False, ignore_unknown_values=False): - """Inserts rows into the specified table. - - Args: - project_id: The project id owning the table. - dataset_id: The dataset id owning the table. - table_id: The table id. - rows: A list of plain Python dictionaries. Each dictionary is a row and - each key in it is the name of a field. - skip_invalid_rows: If there are rows with insertion errors, whether they - should be skipped, and all others should be inserted successfully. - ignore_unknown_values: Set this option to true to ignore unknown column - names. If the input rows contain columns that are not - part of the existing table's schema, those columns are ignored, and - the rows are successfully inserted. - - Returns: - A tuple (bool, errors). If first element is False then the second element - will be a bigquery.InsertErrorsValueListEntry instance containing - specific errors. - """ - - # Prepare rows for insertion. Of special note is the row ID that we add to - # each row in order to help BigQuery avoid inserting a row multiple times. - # BigQuery will do a best-effort if unique IDs are provided. This situation - # can happen during retries on failures. - # TODO(silviuc): Must add support to writing TableRow's instead of dicts. + """Inserts rows into the specified table.""" insert_ids = [ str(self.unique_row_id) if not insert_ids else insert_ids[i] for i, _ in enumerate(rows) @@ -1333,77 +2034,79 @@ def insert_rows( ] result, errors = self._insert_all_rows( - project_id, dataset_id, table_id, rows, insert_ids, + project_id, + dataset_id, + table_id, + rows, + insert_ids, skip_invalid_rows=skip_invalid_rows, ignore_unknown_values=ignore_unknown_values) return result, errors def _convert_cell_value_to_dict(self, value, field): - if field.type == 'STRING': - # Input: "XYZ" --> Output: "XYZ" + f_type = field.type if hasattr(field, 'type') else field.field_type + if f_type == 'STRING': return value - elif field.type == 'BOOLEAN': - # Input: "true" --> Output: True - return value == 'true' - elif field.type == 'INTEGER': - # Input: "123" --> Output: 123 + elif f_type in ('BOOLEAN', 'BOOL'): + return value == 'true' if isinstance(value, str) else bool(value) + elif f_type in ('INTEGER', 'INT64'): return int(value) - elif field.type == 'FLOAT': - # Input: "1.23" --> Output: 1.23 + elif f_type in ('FLOAT', 'FLOAT64'): return float(value) - elif field.type == 'TIMESTAMP': - # The UTC should come from the timezone library but this is a known - # issue in python 2.7 so we'll just hardcode it as we're reading using - # utcfromtimestamp. - # Input: 1478134176.985864 --> Output: "2016-11-03 00:49:36.985864 UTC" - dt = datetime.datetime.utcfromtimestamp(float(value)) - return dt.strftime('%Y-%m-%d %H:%M:%S.%f UTC') - elif field.type == 'BYTES': - # Input: "YmJi" --> Output: "YmJi" - return value - elif field.type == 'DATE': - # Input: "2016-11-03" --> Output: "2016-11-03" - return value - elif field.type == 'DATETIME': - # Input: "2016-11-03T00:49:36" --> Output: "2016-11-03T00:49:36" - return value - elif field.type == 'TIME': - # Input: "00:49:36" --> Output: "00:49:36" + elif f_type == 'TIMESTAMP': + if isinstance(value, (int, float)): + dt = datetime.datetime.fromtimestamp( + float(value), tz=datetime.timezone.utc) + return dt.strftime('%Y-%m-%d %H:%M:%S.%f UTC') + return str(value) + elif f_type == 'BYTES': return value - elif field.type == 'RECORD': - # Note that a schema field object supports also a RECORD type. However - # when querying, the repeated and/or record fields are flattened - # unless we pass the flatten_results flag as False to the source + elif f_type == 'DATE': + return str(value) + elif f_type == 'DATETIME': + return str(value) + elif f_type == 'TIME': + return str(value) + elif f_type in ('RECORD', 'STRUCT'): return self.convert_row_to_dict(value, field) - elif field.type == 'NUMERIC': - return decimal.Decimal(value) - elif field.type == 'GEOGRAPHY': + elif f_type in ('NUMERIC', 'BIGNUMERIC'): + return decimal.Decimal(str(value)) + elif f_type == 'GEOGRAPHY': return value else: - raise RuntimeError('Unexpected field type: %s' % field.type) + raise RuntimeError('Unexpected field type: %s' % f_type) def convert_row_to_dict(self, row, schema): """Converts a TableRow instance using the schema to a Python dict.""" result = {} - for index, field in enumerate(schema.fields): + fields = schema.fields if hasattr(schema, 'fields') else schema + for index, field in enumerate(fields): value = None - if isinstance(schema, bigquery.TableSchema): + if hasattr(row, 'f'): cell = row.f[index] value = from_json_value(cell.v) if cell.v is not None else None - elif isinstance(schema, bigquery.TableFieldSchema): + elif isinstance(row, dict) and 'f' in row: cell = row['f'][index] value = cell['v'] if 'v' in cell else None - if field.mode == 'REPEATED': + elif isinstance(row, dict): + value = row.get(field.name) + elif hasattr(row, field.name): + value = getattr(row, field.name) + elif isinstance(row, (list, tuple)) and index < len(row): + value = row[index] + + mode = getattr(field, 'mode', 'NULLABLE') + if mode == 'REPEATED': if value is None: - # Ideally this should never happen as repeated fields default to - # returning an empty list result[field.name] = [] else: result[field.name] = [ - self._convert_cell_value_to_dict(x['v'], field) for x in value + self._convert_cell_value_to_dict( + x['v'] if isinstance(x, dict) and 'v' in x else x, field) + for x in value ] elif value is None: - if not field.mode == 'NULLABLE': + if not mode == 'NULLABLE': raise ValueError( 'Received \'None\' as the value for the field %s ' 'but the field is not NULLABLE.' % field.name) @@ -1419,9 +2122,37 @@ def from_pipeline_options(pipeline_options: PipelineOptions): @staticmethod def _bigquery_client(pipeline_options: PipelineOptions): - return bigquery.BigqueryV2( + raw_credentials = auth.get_service_credentials(pipeline_options) + google_credentials = ( + raw_credentials.get_google_auth_credentials() if hasattr( + raw_credentials, 'get_google_auth_credentials') else + raw_credentials) + project = None + if pipeline_options: + try: + from apache_beam.options.pipeline_options import GoogleCloudOptions + project = pipeline_options.view_as(GoogleCloudOptions).project + if hasattr(project, 'get') and callable(project.get): + project = project.get() + except Exception: + project = None + + client_cls = getattr(gcp_bigquery, 'Client', None) + if isinstance(client_cls, type): + try: + client_info = ClientInfo( + user_agent="apache-beam-%s" % + apache_beam.__version__) if ClientInfo else None + return gcp_bigquery.Client( + project=project, + credentials=google_credentials, + client_info=client_info) + except Exception: + pass + + return apitools_bigquery.BigqueryV2( http=get_new_http(), - credentials=auth.get_service_credentials(pipeline_options), + credentials=raw_credentials, response_encoding='utf8', additional_http_headers={ "user-agent": "apache-beam-%s" % apache_beam.__version__ @@ -1623,45 +2354,26 @@ def process(self, element, *side_inputs): def beam_row_from_dict(row: dict, schema): - """Converts a dictionary row to a Beam Row. - Nested records and lists are supported. - - Args: - row (dict): - The row to convert. - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): - The table schema. Will be used to help convert the row. - - Returns: - ~apache_beam.pvalue.Row: The converted row. - """ - if not isinstance(schema, (bigquery.TableSchema, bigquery.TableFieldSchema)): + """Converts a dictionary row to a Beam Row.""" + if not isinstance(schema, + (TableSchema, TableFieldSchema)) and not hasattr(schema, + 'fields'): schema = get_bq_tableschema(schema) beam_row = {} - for field in schema.fields: + fields = schema.fields if hasattr(schema, 'fields') else schema + for field in fields: name = field.name - mode = field.mode.upper() - type = field.type.upper() - # When writing with Storage Write API via xlang, we give the Beam Row - # PCollection a hint on the schema using `with_output_types`. - # This requires that each row has all the fields in the schema. - # However, it's possible that some nullable fields don't appear in the row. - # For this case, we create the field with a `None` value - # None is also set when a repeated field is missing as BigQuery - # converts Null Repeated fields to empty lists - if name not in row and mode != "REQUIRED": - row[name] = None - - value = row[name] - if type in ["RECORD", "STRUCT"] and value: - # if this is a list of records, we create a list of Beam Rows + mode = (getattr(field, 'mode', None) or 'NULLABLE').upper() + field_type = (field.type + if hasattr(field, 'type') else field.field_type).upper() + + value = row.get(name) if row else None + if field_type in ["RECORD", "STRUCT"] and value: if mode == "REPEATED": list_of_beam_rows = [] for record in value: list_of_beam_rows.append(beam_row_from_dict(record, field)) beam_row[name] = list_of_beam_rows - # otherwise, create a Beam Row from this record else: beam_row[name] = beam_row_from_dict(value, field) else: @@ -1670,150 +2382,103 @@ def beam_row_from_dict(row: dict, schema): def get_table_schema_from_string(schema): - """Transform the string table schema into a - :class:`~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema` instance. - - Args: - schema (str): The string schema to be used if the BigQuery table to write - has to be created. - - Returns: - ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema: - The schema to be used if the BigQuery table to write has to be created - but in the :class:`~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema` format. - """ - table_schema = bigquery.TableSchema() + """Transform string table schema into TableSchema instance.""" + table_schema = TableSchema() schema_list = [s.strip() for s in schema.split(',')] for field_and_type in schema_list: field_name, field_type = field_and_type.split(':') - field_schema = bigquery.TableFieldSchema() - field_schema.name = field_name - field_schema.type = field_type - field_schema.mode = 'NULLABLE' + field_schema = TableFieldSchema( + name=field_name, type=field_type, mode='NULLABLE') table_schema.fields.append(field_schema) return table_schema -def table_schema_to_dict(table_schema): - """Create a dictionary representation of table schema for serialization - """ - def get_table_field(field): - """Create a dictionary representation of a table field - """ - result = {} - result['name'] = field.name - result['type'] = field.type - result['mode'] = getattr(field, 'mode', 'NULLABLE') - if hasattr(field, 'description') and field.description is not None: - result['description'] = field.description - if hasattr(field, 'fields') and field.fields: - result['fields'] = [get_table_field(f) for f in field.fields] - return result - - if not isinstance(table_schema, bigquery.TableSchema): - raise ValueError("Table schema must be of the type bigquery.TableSchema") - schema = {'fields': []} - for field in table_schema.fields: - schema['fields'].append(get_table_field(field)) - return schema +def table_field_to_dict(field): + """Create a dictionary representation of a TableFieldSchema.""" + if isinstance(field, dict): + return field + result = {} + result['name'] = getattr(field, 'name', '') + result['type'] = getattr(field, 'type', None) or getattr( + field, 'field_type', None) or 'STRING' + result['mode'] = getattr(field, 'mode', 'NULLABLE') or 'NULLABLE' + if hasattr(field, 'description') and field.description is not None: + result['description'] = field.description + if hasattr(field, 'fields') and field.fields: + result['fields'] = [table_field_to_dict(f) for f in field.fields] + return result -def get_dict_table_schema(schema): - """Transform the table schema into a dictionary instance. +def table_schema_to_dict(table_schema): + """Create a dictionary representation of table schema for serialization.""" + if not isinstance(table_schema, + (TableSchema, list, tuple)) and not hasattr(table_schema, + 'fields'): + raise ValueError("Table schema must be of the type TableSchema or list") + fields = table_schema.fields if hasattr( + table_schema, 'fields') else table_schema + return {'fields': [table_field_to_dict(field) for field in fields]} - Args: - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): - The schema to be used if the BigQuery table to write has to be created. - This can either be a dict or string or in the TableSchema format. - Returns: - Dict[str, Any]: The schema to be used if the BigQuery table to write has - to be created but in the dictionary format. - """ +def get_dict_table_schema(schema): + """Transform table schema into dictionary representation.""" if (isinstance(schema, (dict, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema elif isinstance(schema, str): table_schema = get_table_schema_from_string(schema) return table_schema_to_dict(table_schema) - elif isinstance(schema, bigquery.TableSchema): + elif isinstance(schema, (TableSchema, list)) or hasattr(schema, 'fields'): return table_schema_to_dict(schema) else: raise TypeError('Unexpected schema argument: %s.' % schema) def get_bq_tableschema(schema): - """Convert the table schema to a TableSchema object. - - Args: - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): - The schema to be used if the BigQuery table to write has to be created. - This can either be a dict or string or in the TableSchema format. - - Returns: - ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema: The schema as a TableSchema object. - """ - if (isinstance(schema, - (bigquery.TableSchema, value_provider.ValueProvider)) or + """Convert table schema to TableSchema object.""" + if (isinstance(schema, (TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema + elif isinstance(schema, list): + return TableSchema(fields=schema) elif isinstance(schema, str): return get_table_schema_from_string(schema) elif isinstance(schema, dict): schema_string = json.dumps(schema) return parse_table_schema_from_json(schema_string) + elif hasattr(schema, 'fields'): + return TableSchema(fields=list(schema.fields)) else: raise TypeError('Unexpected schema argument: %s.' % schema) def get_avro_schema_from_table_schema(schema): - """Transform the table schema into an Avro schema. - - Args: - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): - The TableSchema to convert to Avro schema. This can either be a dict or - string or in the TableSchema format. - - Returns: - Dict[str, Any]: An Avro schema, which can be used by fastavro. - """ + """Transform table schema into Avro schema.""" dict_table_schema = get_dict_table_schema(schema) return bigquery_avro_tools.get_record_schema_from_dict_table_schema( "root", dict_table_schema) def get_beam_typehints_from_tableschema(schema, type_overrides=None): - """Extracts Beam Python type hints from the schema. - - Args: - schema (~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): - The TableSchema to extract type hints from. - type_overrides (dict): Optional mapping of BigQuery type names (uppercase) - to Python types. These override the default mappings in - BIGQUERY_TYPE_TO_PYTHON_TYPE. For example: - ``{'DATE': datetime.date, 'JSON': dict}`` - - Returns: - List[Tuple[str, Any]]: A list of type hints that describe the input schema. - Nested and repeated fields are supported. - """ - effective_types = {**BIGQUERY_TYPE_TO_PYTHON_TYPE, **(type_overrides or {})} - if not isinstance(schema, (bigquery.TableSchema, bigquery.TableFieldSchema)): + """Extracts Beam Python type hints from the schema.""" + normalized_overrides = { + k.upper(): v + for k, v in (type_overrides or {}).items() + } + effective_types = {**BIGQUERY_TYPE_TO_PYTHON_TYPE, **normalized_overrides} + if not isinstance(schema, + (TableSchema, TableFieldSchema)) and not hasattr(schema, + 'fields'): schema = get_bq_tableschema(schema) typehints = [] - for field in schema.fields: - name, field_type, mode = field.name, field.type.upper(), field.mode.upper() + fields = schema.fields if hasattr(schema, 'fields') else schema + for field in fields: + name = field.name + field_type = (field.type + if hasattr(field, 'type') else field.field_type).upper() + mode = (getattr(field, 'mode', None) or 'NULLABLE').upper() if field_type in ["STRUCT", "RECORD"]: - # Structs can be represented as Beam Rows. typehint = RowTypeConstraint.from_fields( get_beam_typehints_from_tableschema(field, type_overrides)) elif field_type in effective_types: @@ -1851,72 +2516,73 @@ def generate_bq_job_name(job_name, step_id, job_type, random=None): def check_schema_equal( - left: Union['bigquery.TableSchema', 'bigquery.TableFieldSchema'], - right: Union['bigquery.TableSchema', 'bigquery.TableFieldSchema'], + left: Union['TableSchema', 'TableFieldSchema', Any], + right: Union['TableSchema', 'TableFieldSchema', Any], *, ignore_descriptions: bool = False, ignore_field_order: bool = False) -> bool: - """Check whether schemas are equivalent. - - This comparison function differs from using == to compare TableSchema - because it ignores categories, policy tags, descriptions (optionally), and - field ordering (optionally). + """Check whether schemas are equivalent.""" + if left is None and right is None: + return True + if left is None or right is None: + return False - Args: - left (~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema, ~apache_beam.io.gcp.internal.clients.\ -bigquery.bigquery_v2_messages.TableFieldSchema): - One schema to compare. - right (~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema, ~apache_beam.io.gcp.internal.clients.\ -bigquery.bigquery_v2_messages.TableFieldSchema): - The other schema to compare. - ignore_descriptions (bool): (optional) Whether or not to ignore field - descriptions when comparing. Defaults to False. - ignore_field_order (bool): (optional) Whether or not to ignore struct field - order when comparing. Defaults to False. + is_field_left = isinstance( + left, (TableFieldSchema, getattr(gcp_bigquery, 'SchemaField', ()))) or ( + hasattr(left, 'name') and + (hasattr(left, 'type') or hasattr(left, 'field_type'))) + is_field_right = isinstance( + right, (TableFieldSchema, getattr(gcp_bigquery, 'SchemaField', ()))) or ( + hasattr(right, 'name') and + (hasattr(right, 'type') or hasattr(right, 'field_type'))) - Returns: - bool: True if the schemas are equivalent, False otherwise. - """ - if type(left) != type(right) or not isinstance( - left, (bigquery.TableSchema, bigquery.TableFieldSchema)): + if is_field_left != is_field_right: return False - if isinstance(left, bigquery.TableFieldSchema): + if is_field_left: if left.name != right.name: return False - if left.type != right.type: - # Check for type aliases - if sorted( - (left.type, right.type)) not in (["BOOL", "BOOLEAN"], ["FLOAT", - "FLOAT64"], - ["INT64", "INTEGER"], ["RECORD", - "STRUCT"]): + l_type = (left.type if hasattr(left, 'type') else left.field_type).upper() + r_type = (right.type + if hasattr(right, 'type') else right.field_type).upper() + if l_type != r_type: + if sorted((l_type, r_type)) not in (["BOOL", "BOOLEAN"], ["FLOAT", + "FLOAT64"], + ["INT64", "INTEGER"], ["RECORD", + "STRUCT"]): return False - if left.mode != right.mode: + l_mode = (getattr(left, 'mode', None) or 'NULLABLE').upper() + r_mode = (getattr(right, 'mode', None) or 'NULLABLE').upper() + if l_mode != r_mode: return False - if not ignore_descriptions and left.description != right.description: - return False + if not ignore_descriptions: + if getattr(left, 'description', None) != getattr(right, + 'description', + None): + return False - if isinstance(left, - bigquery.TableSchema) or left.type in ("RECORD", "STRUCT"): - if len(left.fields) != len(right.fields): + l_fields = left if isinstance(left, + (list, + tuple)) else getattr(left, 'fields', None) + r_fields = right if isinstance(right, + (list, + tuple)) else getattr(right, 'fields', None) + if l_fields is not None or r_fields is not None: + l_fields = list(l_fields or []) + r_fields = list(r_fields or []) + if len(l_fields) != len(r_fields): return False if ignore_field_order: - left_fields = sorted(left.fields, key=lambda field: field.name) - right_fields = sorted(right.fields, key=lambda field: field.name) - else: - left_fields = left.fields - right_fields = right.fields + l_fields = sorted(l_fields, key=lambda field: field.name) + r_fields = sorted(r_fields, key=lambda field: field.name) - for left_field, right_field in zip(left_fields, right_fields): - if not check_schema_equal(left_field, - right_field, + for lf, rf in zip(l_fields, r_fields): + if not check_schema_equal(lf, + rf, ignore_descriptions=ignore_descriptions, ignore_field_order=ignore_field_order): return False diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index 078c42160941..f78d59047076 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -566,6 +566,8 @@ def test_insert_rows_sets_metric_on_failure(self): ]) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) wrapper.insert_rows("my_project", "my_dataset", "my_table", []) + wrapper.insert_rows("my_project", "my_dataset", "my_table", []) + wrapper.insert_rows("my_project", "my_dataset", "my_table", []) # Expect two failing calls, then a success (i.e. two retries). self.verify_write_call_metric( diff --git a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/__init__.py b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/__init__.py index ec7df8aa128f..3f4e094bfdab 100644 --- a/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/__init__.py +++ b/sdks/python/apache_beam/io/gcp/internal/clients/bigquery/__init__.py @@ -15,10 +15,26 @@ # limitations under the License. # -"""Common imports for generated bigquery client library.""" +"""Common imports for generated bigquery client library. + +.. deprecated:: + The apitools-generated BigQuery client in + apache_beam.io.gcp.internal.clients.bigquery is deprecated and will be removed + in a future version. Use google-cloud-bigquery instead. +""" # pylint:disable=wildcard-import import pkgutil +import warnings + +from apache_beam.utils.annotations import BeamDeprecationWarning + +warnings.warn( + "The apitools-generated BigQuery client in " + "apache_beam.io.gcp.internal.clients.bigquery is deprecated and will be " + "removed in a future version. Use google-cloud-bigquery instead.", + BeamDeprecationWarning, + stacklevel=2) # Protect against environments where apitools library is not available. # pylint: disable=wrong-import-order, wrong-import-position diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 5b02dc9ab0ef..16ab4b63d780 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -150,10 +150,10 @@ def cythonize(*args, **kwargs): pyarrow_dependency = [''] else: pyarrow_dependency = [ - # Generally try to cover versions released in the last two years. - # Update python/sdks/tox.ini to cover the same pyarrow versions - # when updating the bounds here. - 'pyarrow>=14.0.1,<26.0.0', + # Generally try to cover versions released in the last two years. + # Update python/sdks/tox.ini to cover the same pyarrow versions + # when updating the bounds here. + 'pyarrow>=14.0.1,<26.0.0', ] # Exclude pandas<=1.4.2 since it doesn't work with numpy 1.24.x. @@ -199,8 +199,7 @@ def cythonize(*args, **kwargs): # `cls_token=` or tokenizers reinstates `cls=` as a deprecation alias. 'tokenizers<0.23', # Avoid torch 2.12.0+ which fails to run unit tests with segfault. - 'torch<2.12.0', - # Match tested transformers range. + 'torch<2.12.0', # Match tested transformers range. 'transformers>=4.28.0,<4.56.0', # Keep tokenizers compatible with this transformers range. 'tokenizers>=0.13.3,<0.22.0', @@ -460,9 +459,9 @@ def get_portability_package_data(): # BEAM-8840: Do NOT use tests_require or setup_requires. extras_require={ 'dev': [ - 'pyrefly==1.1.1', - 'ruff==0.15.22', - 'yapf==0.43.0', + 'pyrefly==1.1.1', + 'ruff==0.15.22', + 'yapf==0.43.0', ], 'dill': [ # Dill doesn't have forwards-compatibility guarantees within minor @@ -528,7 +527,7 @@ def get_portability_package_data(): 'google-cloud-resource-manager>=1.12.0,<2', 'google-cloud-dataflow-client>=0.13.0,<0.14.0', # GCP packages required by tests - 'google-cloud-bigquery>=2.0.0,<4', + 'google-cloud-bigquery>=3.0.0,<4', 'google-cloud-bigquery-storage>=2.6.3,<3', 'google-cloud-core>=2.0.0,<3', # 2.42.0 improves MutationsBatcher error handling: it surfaces @@ -600,7 +599,7 @@ def get_portability_package_data(): 'tf2onnx>=1.17.0,<1.18', ] + ml_base_core, 'p310_ml_test': [ - 'datatable', + 'datatable', ] + ml_base + qdrant_dependency, 'p312_ml_test': [ 'datatable', From 962918d05d1babacdcc1bdaed69f01c99d0bc7d8 Mon Sep 17 00:00:00 2001 From: Jack McCluskey <34928439+jrmccluskey@users.noreply.github.com> Date: Wed, 26 Aug 2026 09:10:13 -0400 Subject: [PATCH 02/12] Update bigquery_tools.py Fix type imports --- sdks/python/apache_beam/io/gcp/bigquery_tools.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index a2eb356c7ce8..f6da17d0d8b5 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -37,9 +37,10 @@ import sys import time import uuid +from collections.abc import Sequence from json.decoder import JSONDecodeError +from typing import Any from typing import Optional -from typing import Sequence from typing import TypeVar from typing import Union From 16c18d80d5264d2d58ff73330362cddb7e56cd8a Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 13:41:31 +0000 Subject: [PATCH 03/12] Fix typing.Any --- sdks/python/apache_beam/io/gcp/bigquery_tools.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index f6da17d0d8b5..2e8b816e49c7 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -36,10 +36,10 @@ import re import sys import time +import typing import uuid from collections.abc import Sequence from json.decoder import JSONDecodeError -from typing import Any from typing import Optional from typing import TypeVar from typing import Union @@ -2517,8 +2517,8 @@ def generate_bq_job_name(job_name, step_id, job_type, random=None): def check_schema_equal( - left: Union['TableSchema', 'TableFieldSchema', Any], - right: Union['TableSchema', 'TableFieldSchema', Any], + left: Union['TableSchema', 'TableFieldSchema', typing.Any], + right: Union['TableSchema', 'TableFieldSchema', typing.Any], *, ignore_descriptions: bool = False, ignore_field_order: bool = False) -> bool: From 04355c963e319bf0c99f7573085497e5bdbea290 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 13:57:45 +0000 Subject: [PATCH 04/12] Restore docstrings --- .../apache_beam/io/gcp/bigquery_tools.py | 399 ++++++++++++++++-- 1 file changed, 366 insertions(+), 33 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 2e8b816e49c7..6efedd0233cc 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -38,9 +38,9 @@ import time import typing import uuid -from collections.abc import Sequence from json.decoder import JSONDecodeError from typing import Optional +from typing import Sequence from typing import TypeVar from typing import Union @@ -118,6 +118,11 @@ # ----------------------------------------------------------------------------- # Compatibility Models for TableReference, DatasetReference, Schema, and Jobs. +# +# These classes and monkey patches bridge between legacy apitools structures +# and modern google.cloud.bigquery objects, providing camelCase attribute access +# (e.g. projectId, datasetId, tableId, tableReference) for backwards +# compatibility across pipelines, transforms, and test suites. if bigquery is not None and hasattr(bigquery, 'TableReference'): TableReference = bigquery.TableReference @@ -421,6 +426,7 @@ def totalBytesProcessed(self): def _to_gcp_table_ref(table_ref, default_project=None): + """Converts a TableReference or string into a google.cloud.bigquery.TableReference.""" if table_ref is None: return None if gcp_bigquery is not None and isinstance( @@ -447,6 +453,7 @@ def _to_gcp_table_ref(table_ref, default_project=None): def _to_gcp_dataset_ref(dataset_ref, project=None): + """Converts a DatasetReference or string into a google.cloud.bigquery.DatasetReference.""" if dataset_ref is None: return None if gcp_bigquery is not None and isinstance( @@ -471,6 +478,7 @@ def _to_gcp_dataset_ref(dataset_ref, project=None): def _to_gcp_schema(schema): + """Converts a TableSchema, list of fields, dict, or string into a list of google.cloud.bigquery.SchemaField.""" if schema is None: return None if isinstance(schema, (list, tuple)): @@ -634,6 +642,14 @@ def parse_table_schema_from_json(schema_string): 'Unable to parse JSON schema: %s - %r' % (schema_string, e)) def _parse_schema_field(field): + """Parse a single schema field from dictionary. + + Args: + field: Dictionary object containing serialized schema. + + Returns: + A TableFieldSchema for a single column in BigQuery. + """ name = field['name'] field_type = field.get('type') or field.get('type_') or 'STRING' mode = field.get('mode', 'NULLABLE') @@ -652,15 +668,31 @@ def _parse_schema_field(field): def parse_table_reference(table, dataset=None, project=None): - """Parses a table reference into a TableReference instance. + """Parses a table reference into a (project, dataset, table) tuple. Args: - table: The ID of the table, or full reference string / TableReference. - dataset: The ID of the dataset containing this table. - project: The ID of the project containing this table. + table: The ID of the table. The ID must contain only letters + (a-z, A-Z), numbers (0-9), connectors (-_). If dataset argument is None + then the table argument must contain the entire table reference: + 'DATASET.TABLE' or 'PROJECT:DATASET.TABLE'. This argument can be a + TableReference instance in which case dataset and project are + ignored and the reference is returned as a result. Additionally, for date + partitioned tables, appending '$YYYYmmdd' to the table name is supported, + e.g. 'DATASET.TABLE$YYYYmmdd'. + dataset: The ID of the dataset containing this table or null if the table + reference is specified entirely by the table argument. + project: The ID of the project containing this table or null if the table + reference is specified entirely by the table (and possibly dataset) + argument. Returns: - A TableReference object. + A TableReference object from the bigquery API. The object has the following + attributes: projectId, datasetId, and tableId. + If the input is a TableReference object, a new object will be returned. + + Raises: + ValueError: if the table reference as a string does not match the expected + format. """ if isinstance(table, TableReference): return TableReference( @@ -677,6 +709,9 @@ def parse_table_reference(table, dataset=None, project=None): elif isinstance(table, value_provider.ValueProvider): return table + # If dataset argument is not specified, the expectation is that the + # table argument will contain a full table reference instead of just a + # table name. if dataset is None: pattern = ( f'((?P{_PROJECT_PATTERN})[:\\.])?' @@ -755,6 +790,8 @@ class BigQueryWrapper(object): should not be used in production code. """ + # If updating following names, also update the corresponding pydocs in + # bigquery.py. TEMP_TABLE = 'beam_temp_table_' TEMP_DATASET = 'beam_temp_dataset_' @@ -825,7 +862,16 @@ def _is_modern_client(self): @property def unique_row_id(self): - """Returns a unique row ID (str) used to avoid multiple insertions.""" + """Returns a unique row ID (str) used to avoid multiple insertions. + + If the row ID is provided, BigQuery will make a best effort to not insert + the same row multiple times for fail and retry scenarios in which the insert + request may be issued several times. This comes into play for sinks executed + in a local runner. + + Returns: + a unique row ID string + """ self._unique_row_id += 1 return '%s_%d' % (self._row_id_prefix, self._unique_row_id) @@ -839,7 +885,11 @@ def _get_temp_table(self, project_id): project=project_id) def _get_temp_table_project(self, fallback_project_id): - """Returns the project ID for temporary table operations.""" + """Returns the project ID for temporary table operations. + + If temp_table_ref exists, returns its projectId. + Otherwise, returns the fallback_project_id. + """ if self.temp_table_ref: return self.temp_table_ref.projectId else: @@ -854,7 +904,13 @@ def _get_temp_dataset(self): num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def get_query_location(self, project_id, query, use_legacy_sql): - """Get the location of tables referenced in a query.""" + """Get the location of tables referenced in a query. + + This method returns the location of the first available referenced + table for user in the query and depends on the BigQuery service to + provide error handling for queries that reference tables in multiple + locations. + """ if self._is_modern_client: job_config = gcp_bigquery.QueryJobConfig( dry_run=True, @@ -864,6 +920,8 @@ def get_query_location(self, project_id, query, use_legacy_sql): response = self.client.query( query, job_config=job_config, project=project_id) except (Forbidden, HttpForbiddenError): + # Permission access for table (i.e. from authorized_view), + # try next one return None except Exception: raise @@ -874,11 +932,13 @@ def get_query_location(self, project_id, query, use_legacy_sql): referenced_tables = getattr( response.statistics.query, 'referencedTables', None) else: + # This behavior is only expected in tests _LOGGER.warning( - "Unable to get location, missing statistics. Query: %s", query) + "Unable to get location, missing response.statistics. Query: %s", + query) return None - if referenced_tables: + if referenced_tables: # Guards against both non-empty and non-None for table in referenced_tables: try: p = getattr(table, 'project', None) or getattr( @@ -889,6 +949,8 @@ def get_query_location(self, project_id, query, use_legacy_sql): table, 'tableId', None) location = self.get_table_location(p, d, t) except (Forbidden, HttpForbiddenError, ClientError): + # Permission access for table (i.e. from authorized_view), + # try next one continue if location: _LOGGER.info( @@ -918,19 +980,31 @@ def get_query_location(self, project_id, query, use_legacy_sql): jobReference=reference)) response = self.client.jobs.Insert(request) if response.statistics is None: + # This behavior is only expected in tests _LOGGER.warning( "Unable to get location, missing response.statistics. Query: %s", query) return None referenced_tables = response.statistics.query.referencedTables - if referenced_tables: + if referenced_tables: # Guards against both non-empty and non-None for table in referenced_tables: try: location = self.get_table_location( table.projectId, table.datasetId, table.tableId) except (HttpForbiddenError, Forbidden): + # Permission access for table (i.e. from authorized_view), + # try next one continue + _LOGGER.info( + "Using location %r from table %r referenced by query %s", + location, + table, + query) return location + _LOGGER.debug( + "Query %s does not reference any tables or " + "you don't have permission to inspect them.", + query) return None @retry.with_exponential_backoff( @@ -969,7 +1043,10 @@ def _insert_copy_job( except (Conflict, HttpError) as exn: if getattr(exn, 'code', None) == 409 or getattr( exn, 'status_code', None) == 409 or isinstance(exn, Conflict): - _LOGGER.info("BigQuery copy job %s already exists.", job_id) + _LOGGER.info( + "BigQuery copy job %s already exists, will not retry inserting it: %s", + job_id, + exn) return JobReference(job_id=job_id, project=project_id) raise @@ -1007,11 +1084,17 @@ def _insert_load_job( additional_load_parameters=None, source_format=None, job_labels=None): + if not source_uris and not source_stream: _LOGGER.warning( 'Both source URIs and source stream are not provided. BigQuery load ' 'job will not load any data.') + if source_uris and source_stream: + raise ValueError( + 'Only one of source_uris and source_stream may be specified. ' + 'Got both.') + if self._is_modern_client: dst_ref = _to_gcp_table_ref(table_reference, default_project=project_id) job_schema = None if schema == 'SCHEMA_AUTODETECT' else _to_gcp_schema( @@ -1048,7 +1131,10 @@ def _insert_load_job( except (Conflict, HttpError) as exn: if getattr(exn, 'code', None) == 409 or getattr( exn, 'status_code', None) == 409 or isinstance(exn, Conflict): - _LOGGER.info("BigQuery load job %s already exists.", job_id) + _LOGGER.info( + "BigQuery load job %s already exists, will not retry inserting it: %s", + job_id, + exn) return JobReference(job_id=job_id, project=project_id) raise @@ -1091,6 +1177,7 @@ def _parse_location_from_exc(content, job_id): """Parse job location from Exception content.""" if isinstance(content, bytes): content = content.decode('ascii', 'replace') + # search for "Already Exists: Job :." m = re.search(r"Already Exists: Job \S+\:(\S+)\." + job_id, content) if not m: _LOGGER.warning( @@ -1106,6 +1193,10 @@ def _start_job( """Inserts a BigQuery job for legacy apitools client. If the job exists already, it returns it. + + Args: + request (bigquery.BigqueryJobsInsertRequest): An insert job request. + stream (IO[bytes]): A bytes IO object open for reading. """ try: upload = None @@ -1129,6 +1220,7 @@ def _start_job( job_location = self._parse_location_from_exc(exn.content, jobId) response = request.job if not response.jobReference.location and job_location: + # Request not constructed with location response.jobReference.location = job_location return response else: @@ -1214,7 +1306,18 @@ def _start_query_job( 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.""" + """Poll job until it is DONE. + + Args: + job_reference: JobReference instance or job ID string. + 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. + + Raises: + `RuntimeError`: If the job is FAILED or the number of retries has been + reached. + """ retry = 0 project = getattr(job_reference, 'projectId', None) or getattr( job_reference, 'project', None) @@ -1286,12 +1389,20 @@ def _insert_all_rows( insert_ids, skip_invalid_rows=False, ignore_unknown_values=False): - """Calls the insertAll BigQuery API endpoint.""" + """Calls the insertAll BigQuery API endpoint. + + Docs for this BQ call: https://cloud.google.com/bigquery/docs/reference\ + /rest/v2/tabledata/insertAll.""" + # The rows argument is a list of plain Python dictionaries or rows. resource = resource_identifiers.BigQueryTable( project_id, dataset_id, table_id) labels = { + # TODO(ajamato): Add Ptransform label. monitoring_infos.SERVICE_LABEL: 'BigQuery', + # Refer to any method which writes elements to BigQuery in batches + # as "BigQueryBatchWrite". I.e. storage API's insertAll, or future + # APIs introduced. monitoring_infos.METHOD_LABEL: 'BigQueryBatchWrite', monitoring_infos.RESOURCE_LABEL: resource, monitoring_infos.BIGQUERY_PROJECT_ID_LABEL: project_id, @@ -1328,9 +1439,11 @@ def _insert_all_rows( err, 'reason', None) service_call_metric.call(reason or 'unknown') except (ClientError, GoogleAPICallError, HttpError) as e: + # e.code contains the numeric http status code. status_code = getattr(e, 'code', None) or getattr( e, 'status_code', None) or 500 service_call_metric.call(status_code) + # Package exception with required fields reason = None if hasattr(e, 'response') and getattr(e.response, 'reason', None): reason = e.response.reason @@ -1341,6 +1454,7 @@ def _insert_all_rows( reason = e.errors[0].get('reason') if not reason: reason = e.__class__.__name__ + # Add all rows to the errors list along with the error errors = [{ 'index': i, 'errors': [{ 'reason': reason, 'message': str(e) @@ -1358,6 +1472,9 @@ def _insert_all_rows( return not errors, errors # Legacy apitools path + # The rows argument is a list of + # bigquery.TableDataInsertAllRequest.RowsValueListEntry instances as + # required by the InsertAll() method. row_list = [] for row, insert_id in zip(rows, insert_ids): row_list.append( @@ -1389,6 +1506,7 @@ def _insert_all_rows( service_call_metric.call(error.reason) except HttpError as e: service_call_metric.call(e) + # Re-raise the exception so that we re-try appropriately. raise finally: self._latency_histogram_metric.update( @@ -1399,7 +1517,18 @@ 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.""" + """Lookup a table's metadata object. + + Args: + project_id: table lookup parameter + dataset_id: table lookup parameter + table_id: table lookup parameter + + Returns: + Table instance (bigquery.Table or google.cloud.bigquery.Table). + Raises: + NotFound or HttpError: if lookup failed. + """ if self._is_modern_client: table_ref = TableReference( projectId=project_id, datasetId=dataset_id, tableId=table_id) @@ -1408,7 +1537,8 @@ def get_table(self, project_id, dataset_id, table_id): # Fallback for legacy client request = apitools_bigquery.BigqueryTablesGetRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) - return self.client.tables.Get(request) + response = self.client.tables.Get(request) + return response def _create_table( self, @@ -1417,6 +1547,7 @@ def _create_table( table_id, schema, additional_parameters=None): + valid_tablename = regex.fullmatch(_TABLE_PATTERN, table_id, regex.ASCII) if not valid_tablename: raise ValueError( @@ -1434,6 +1565,7 @@ def _create_table( setattr(table, k, v) response = self.client.create_table(table) _LOGGER.debug("Created the table with id %s", table_id) + # The response is a Table instance. return response # Fallback for legacy client @@ -1447,6 +1579,7 @@ def _create_table( projectId=project_id, datasetId=dataset_id, table=table) response = self.client.tables.Insert(request) _LOGGER.debug("Created the table with id %s", table_id) + # The response is a bigquery.Table instance. return response @retry.with_exponential_backoff( @@ -1460,6 +1593,7 @@ def get_or_create_dataset( labels=None, kms_key=None, default_table_expiration_ms=None): + # Check if dataset already exists otherwise create it if self._is_modern_client: dataset_ref = DatasetReference(projectId=project_id, datasetId=dataset_id) try: @@ -1487,6 +1621,7 @@ def get_or_create_dataset( dataset.default_table_expiration_ms = default_table_expiration_ms response = self.client.create_dataset(dataset) self.created_temp_dataset = True + # The response is a Dataset instance. return response raise @@ -1521,6 +1656,7 @@ def get_or_create_dataset( projectId=project_id, dataset=dataset) response = self.client.datasets.Insert(request) self.created_temp_dataset = True + # The response is a bigquery.Dataset instance. return response else: raise @@ -1544,6 +1680,7 @@ def _is_table_empty(self, project_id, dataset_id, table_id): tableId=table_id, maxResults=1) response = self.client.tabledata.List(request) + # The response is a bigquery.TableDataList instance. return response.totalRows == 0 @retry.with_exponential_backoff( @@ -1616,8 +1753,8 @@ def get_table_location(self, project_id, dataset_id, table_id): table = self.get_table(project_id, dataset_id, table_id) return table.location + # Returns true if the temporary dataset was provided by the user. def is_user_configured_dataset(self): - """Returns true if the temporary dataset was provided by the user.""" return ( self.temp_dataset_id and not self.temp_dataset_id.startswith(self.TEMP_DATASET)) @@ -1637,6 +1774,8 @@ def create_temporary_dataset( if (project_id is not None and not self.is_user_configured_dataset() and not self.created_temp_dataset): + # Unittests don't pass projectIds so they can be run without error + # User configured datasets are allowed to pre-exist. raise RuntimeError( 'Dataset %s:%s already exists so cannot be used as temporary.' % (project_id, self.temp_dataset_id)) @@ -1659,6 +1798,8 @@ def clean_up_temporary_dataset(self, project_id): return raise try: + # We do not want to delete temporary datasets configured by the user hence + # we just delete the temporary table in that case. if not self.is_user_configured_dataset(): self._delete_dataset(temp_table.projectId, temp_table.datasetId, True) else: @@ -1689,6 +1830,8 @@ def clean_up_temporary_dataset(self, project_id): else: raise try: + # We do not want to delete temporary datasets configured by the user hence + # we just delete the temporary table in that case. if not self.is_user_configured_dataset(): self._delete_dataset(temp_table.projectId, temp_table.datasetId, True) else: @@ -1726,7 +1869,8 @@ def _clean_up_beam_labelled_temporary_datasets( if getattr(exn, 'code', None) == 403 or getattr( exn, 'status_code', None) == 403 or isinstance(exn, Forbidden): _LOGGER.warning( - 'Permission denied to delete temporary dataset for clean up.') + 'Permission denied to delete temporary dataset %s for clean up.', + project_id) return raise return @@ -1796,7 +1940,7 @@ def perform_load_job( """Starts a job to load data into BigQuery. Returns: - JobReference with the information about the job that was started. + JobReference or bigquery.JobReference with the information about the job that was started. """ if source_uris and source_stream: raise ValueError( @@ -1838,7 +1982,7 @@ def perform_extract_job( """Starts a job to export data from BigQuery. Returns: - JobReference with the information about the job that was started. + JobReference or bigquery.JobReference with the information about the job that was started. """ job_project = project or getattr(table_reference, 'projectId', None) or getattr( @@ -1903,7 +2047,28 @@ def get_or_create_table( create_disposition, write_disposition, additional_create_parameters=None): - """Gets or creates a table based on create and write dispositions.""" + """Gets or creates a table based on create and write dispositions. + + The function mimics the behavior of BigQuery import jobs when using the + same create and write dispositions. + + Args: + project_id: The project id owning the table. + dataset_id: The dataset id owning the table. + table_id: The table id. + schema: A TableSchema instance or None. + create_disposition: CREATE_NEVER or CREATE_IF_NEEDED. + write_disposition: WRITE_APPEND, WRITE_EMPTY or WRITE_TRUNCATE. + + Returns: + A Table instance if table was found or created. + + Raises: + `RuntimeError`: For various mismatches between the state of the table and + the create/write dispositions passed in. For example if the table is not + empty and WRITE_EMPTY was specified then an error will be raised since + the table was expected to be empty. + """ from apache_beam.io.gcp.bigquery import BigQueryDisposition found_table = None @@ -1919,8 +2084,12 @@ def get_or_create_table( else: raise + # If table exists already then handle the semantics for WRITE_EMPTY and + # WRITE_TRUNCATE write dispositions. if found_table and write_disposition in ( BigQueryDisposition.WRITE_EMPTY, BigQueryDisposition.WRITE_TRUNCATE): + # Delete the table and recreate it (later) if WRITE_TRUNCATE was + # specified. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: self._delete_table(project_id, dataset_id, table_id) elif (write_disposition == BigQueryDisposition.WRITE_EMPTY and @@ -1929,6 +2098,8 @@ def get_or_create_table( 'Table %s:%s.%s is not empty but write disposition is WRITE_EMPTY.' % (project_id, dataset_id, table_id)) + # Create a new table potentially reusing the schema from a previously + # found table in case the schema was not specified. if schema is None and found_table is None: raise RuntimeError( 'Table %s:%s.%s requires a schema. None can be inferred because the ' @@ -1961,11 +2132,16 @@ def get_or_create_table( table_id, schema or (found_table.schema if found_table else None), created_table) + # if write_disposition == BigQueryDisposition.WRITE_TRUNCATE we delete + # the table before this point. if write_disposition == BigQueryDisposition.WRITE_TRUNCATE: + # BigQuery can route data to the old table for 2 mins max so wait + # that much time before creating the table and writing it _LOGGER.warning( 'Sleeping for 150 seconds before the write as ' + 'BigQuery inserts can be routed to deleted table ' + 'for 2 mins after the delete and create.') + # TODO(BEAM-2673): Remove this sleep by migrating to load api time.sleep(150) return created_table else: @@ -1990,6 +2166,8 @@ def run_query( dry_run=dry_run, job_labels=job_labels) if dry_run: + # If this was a dry run then the fact that we get here means the + # query has no errors. The start_query_job would raise an error otherwise. return job_id = getattr(job, 'job_id', None) or getattr( @@ -2003,9 +2181,15 @@ def run_query( project_id, job_id, page_token, location=location) if hasattr(response, 'jobComplete'): if not response.jobComplete: + # The jobComplete field can be False if the query request times out + # (default is 10 seconds). Note that this is a timeout for the query + # request not for the actual execution of the query in the service. If + # the request times out we keep trying. This situation is quite possible + # if the query will return a large number of rows. _LOGGER.info('Waiting on response from query: %s ...', query) time.sleep(1.0) continue + # We got some results. The last page is signalled by a missing pageToken. yield response.rows, response.schema if not response.pageToken: break @@ -2024,7 +2208,31 @@ def insert_rows( insert_ids=None, skip_invalid_rows=False, ignore_unknown_values=False): - """Inserts rows into the specified table.""" + """Inserts rows into the specified table. + + Args: + project_id: The project id owning the table. + dataset_id: The dataset id owning the table. + table_id: The table id. + rows: A list of plain Python dictionaries. Each dictionary is a row and + each key in it is the name of a field. + insert_ids: Optional list of unique row IDs to avoid duplicate inserts. + skip_invalid_rows: If there are rows with insertion errors, whether they + should be skipped, and all others should be inserted successfully. + ignore_unknown_values: Set this option to true to ignore unknown column + names. If the input rows contain columns that are not + part of the existing table's schema, those columns are ignored, and + the rows are successfully inserted. + + Returns: + A tuple (bool, errors). If first element is False then the second element + will be a list containing specific errors. + """ + # Prepare rows for insertion. Of special note is the row ID that we add to + # each row in order to help BigQuery avoid inserting a row multiple times. + # BigQuery will do a best-effort if unique IDs are provided. This situation + # can happen during retries on failures. + # TODO(silviuc): Must add support to writing TableRow's instead of dicts. insert_ids = [ str(self.unique_row_id) if not insert_ids else insert_ids[i] for i, _ in enumerate(rows) @@ -2047,28 +2255,43 @@ def insert_rows( def _convert_cell_value_to_dict(self, value, field): f_type = field.type if hasattr(field, 'type') else field.field_type if f_type == 'STRING': + # Input: "XYZ" --> Output: "XYZ" return value elif f_type in ('BOOLEAN', 'BOOL'): + # Input: "true" --> Output: True return value == 'true' if isinstance(value, str) else bool(value) elif f_type in ('INTEGER', 'INT64'): + # Input: "123" --> Output: 123 return int(value) elif f_type in ('FLOAT', 'FLOAT64'): + # Input: "1.23" --> Output: 1.23 return float(value) elif f_type == 'TIMESTAMP': + # The UTC should come from the timezone library but this is a known + # issue in python 2.7 so we'll just hardcode it as we're reading using + # utcfromtimestamp. + # Input: 1478134176.985864 --> Output: "2016-11-03 00:49:36.985864 UTC" if isinstance(value, (int, float)): dt = datetime.datetime.fromtimestamp( float(value), tz=datetime.timezone.utc) return dt.strftime('%Y-%m-%d %H:%M:%S.%f UTC') return str(value) elif f_type == 'BYTES': + # Input: "YmJi" --> Output: "YmJi" return value elif f_type == 'DATE': + # Input: "2016-11-03" --> Output: "2016-11-03" return str(value) elif f_type == 'DATETIME': + # Input: "2016-11-03T00:49:36" --> Output: "2016-11-03T00:49:36" return str(value) elif f_type == 'TIME': + # Input: "00:49:36" --> Output: "00:49:36" return str(value) elif f_type in ('RECORD', 'STRUCT'): + # Note that a schema field object supports also a RECORD type. However + # when querying, the repeated and/or record fields are flattened + # unless we pass the flatten_results flag as False to the source return self.convert_row_to_dict(value, field) elif f_type in ('NUMERIC', 'BIGNUMERIC'): return decimal.Decimal(str(value)) @@ -2099,6 +2322,8 @@ def convert_row_to_dict(self, row, schema): mode = getattr(field, 'mode', 'NULLABLE') if mode == 'REPEATED': if value is None: + # Ideally this should never happen as repeated fields default to + # returning an empty list result[field.name] = [] else: result[field.name] = [ @@ -2355,7 +2580,18 @@ def process(self, element, *side_inputs): def beam_row_from_dict(row: dict, schema): - """Converts a dictionary row to a Beam Row.""" + """Converts a dictionary row to a Beam Row. + Nested records and lists are supported. + + Args: + row (dict): + The row to convert. + schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.bigquery_v2_messages.TableSchema, TableSchema): + The table schema. Will be used to help convert the row. + + Returns: + ~apache_beam.pvalue.Row: The converted row. + """ if not isinstance(schema, (TableSchema, TableFieldSchema)) and not hasattr(schema, 'fields'): @@ -2368,13 +2604,25 @@ def beam_row_from_dict(row: dict, schema): field_type = (field.type if hasattr(field, 'type') else field.field_type).upper() + # When writing with Storage Write API via xlang, we give the Beam Row + # PCollection a hint on the schema using `with_output_types`. + # This requires that each row has all the fields in the schema. + # However, it's possible that some nullable fields don't appear in the row. + # For this case, we create the field with a `None` value + # None is also set when a repeated field is missing as BigQuery + # converts Null Repeated fields to empty lists + if row and name not in row and mode != "REQUIRED": + row[name] = None + value = row.get(name) if row else None if field_type in ["RECORD", "STRUCT"] and value: + # if this is a list of records, we create a list of Beam Rows if mode == "REPEATED": list_of_beam_rows = [] for record in value: list_of_beam_rows.append(beam_row_from_dict(record, field)) beam_row[name] = list_of_beam_rows + # otherwise, create a Beam Row from this record else: beam_row[name] = beam_row_from_dict(value, field) else: @@ -2383,7 +2631,17 @@ def beam_row_from_dict(row: dict, schema): def get_table_schema_from_string(schema): - """Transform string table schema into TableSchema instance.""" + """Transform the string table schema into a TableSchema instance. + + Args: + schema (str): The string schema to be used if the BigQuery table to write + has to be created. + + Returns: + TableSchema: + The schema to be used if the BigQuery table to write has to be created + but in the TableSchema format. + """ table_schema = TableSchema() schema_list = [s.strip() for s in schema.split(',')] for field_and_type in schema_list: @@ -2395,7 +2653,14 @@ def get_table_schema_from_string(schema): def table_field_to_dict(field): - """Create a dictionary representation of a TableFieldSchema.""" + """Create a dictionary representation of a table field for serialization. + + Args: + field: A TableFieldSchema or SchemaField instance. + + Returns: + dict: A dictionary representation of the field. + """ if isinstance(field, dict): return field result = {} @@ -2411,7 +2676,14 @@ def table_field_to_dict(field): def table_schema_to_dict(table_schema): - """Create a dictionary representation of table schema for serialization.""" + """Create a dictionary representation of table schema for serialization. + + Args: + table_schema: A TableSchema or list of SchemaField instances. + + Returns: + dict: A dictionary representation of the schema with 'fields' list. + """ if not isinstance(table_schema, (TableSchema, list, tuple)) and not hasattr(table_schema, 'fields'): @@ -2422,7 +2694,17 @@ def table_schema_to_dict(table_schema): def get_dict_table_schema(schema): - """Transform table schema into dictionary representation.""" + """Transform the table schema into a dictionary instance. + + Args: + schema (str, dict, TableSchema, SchemaField list): + The schema to be used if the BigQuery table to write has to be created. + This can either be a dict or string or in the TableSchema format. + + Returns: + Dict[str, Any]: The schema to be used if the BigQuery table to write has + to be created but in the dictionary format. + """ if (isinstance(schema, (dict, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema @@ -2436,7 +2718,16 @@ def get_dict_table_schema(schema): def get_bq_tableschema(schema): - """Convert table schema to TableSchema object.""" + """Convert the table schema to a TableSchema object. + + Args: + schema (str, dict, TableSchema, SchemaField list): + The schema to be used if the BigQuery table to write has to be created. + This can either be a dict or string or in the TableSchema format. + + Returns: + TableSchema: The schema as a TableSchema object. + """ if (isinstance(schema, (TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema @@ -2454,14 +2745,36 @@ def get_bq_tableschema(schema): def get_avro_schema_from_table_schema(schema): - """Transform table schema into Avro schema.""" + """Transform the table schema into an Avro schema. + + Args: + schema (str, dict, TableSchema, SchemaField list): + The TableSchema to convert to Avro schema. This can either be a dict or + string or in the TableSchema format. + + Returns: + Dict[str, Any]: An Avro schema, which can be used by fastavro. + """ dict_table_schema = get_dict_table_schema(schema) return bigquery_avro_tools.get_record_schema_from_dict_table_schema( "root", dict_table_schema) def get_beam_typehints_from_tableschema(schema, type_overrides=None): - """Extracts Beam Python type hints from the schema.""" + """Extracts Beam Python type hints from the schema. + + Args: + schema (TableSchema, SchemaField list, dict, str): + The TableSchema to extract type hints from. + type_overrides (dict): Optional mapping of BigQuery type names (uppercase) + to Python types. These override the default mappings in + BIGQUERY_TYPE_TO_PYTHON_TYPE. For example: + ``{'DATE': datetime.date, 'JSON': dict}`` + + Returns: + List[Tuple[str, Any]]: A list of type hints that describe the input schema. + Nested and repeated fields are supported. + """ normalized_overrides = { k.upper(): v for k, v in (type_overrides or {}).items() @@ -2480,6 +2793,7 @@ def get_beam_typehints_from_tableschema(schema, type_overrides=None): mode = (getattr(field, 'mode', None) or 'NULLABLE').upper() if field_type in ["STRUCT", "RECORD"]: + # Structs can be represented as Beam Rows. typehint = RowTypeConstraint.from_fields( get_beam_typehints_from_tableschema(field, type_overrides)) elif field_type in effective_types: @@ -2522,7 +2836,25 @@ def check_schema_equal( *, ignore_descriptions: bool = False, ignore_field_order: bool = False) -> bool: - """Check whether schemas are equivalent.""" + """Check whether schemas are equivalent. + + This comparison function differs from using == to compare TableSchema + because it ignores categories, policy tags, descriptions (optionally), and + field ordering (optionally). + + Args: + left (TableSchema, TableFieldSchema, SchemaField, list, dict): + One schema to compare. + right (TableSchema, TableFieldSchema, SchemaField, list, dict): + The other schema to compare. + ignore_descriptions (bool): (optional) Whether or not to ignore field + descriptions when comparing. Defaults to False. + ignore_field_order (bool): (optional) Whether or not to ignore struct field + order when comparing. Defaults to False. + + Returns: + bool: True if the schemas are equivalent, False otherwise. + """ if left is None and right is None: return True if left is None or right is None: @@ -2548,6 +2880,7 @@ def check_schema_equal( r_type = (right.type if hasattr(right, 'type') else right.field_type).upper() if l_type != r_type: + # Check for type aliases if sorted((l_type, r_type)) not in (["BOOL", "BOOLEAN"], ["FLOAT", "FLOAT64"], ["INT64", "INTEGER"], ["RECORD", From d8e0442fc1c157925a588d31aecb16c067656b47 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 15:33:03 +0000 Subject: [PATCH 05/12] add extra compat layer, allow use of legacy client via experiment flag --- .../apache_beam/io/gcp/bigquery_tools.py | 258 ++++++++++++++++-- .../apache_beam/io/gcp/bigquery_tools_test.py | 22 ++ 2 files changed, 254 insertions(+), 26 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 6efedd0233cc..0fceae61648f 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -519,6 +519,180 @@ def _to_gcp_schema(schema): return schema +if gcp_bigquery: + + class _ClientTablesCompat: + def __init__(self, client): + self._client = client + + def Get(self, request): + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + tbl_id = getattr(request, 'tableId', None) + if ds_id and tbl_id: + table_ref = gcp_bigquery.TableReference( + gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', + ds_id), + tbl_id) + else: + t_ref = getattr(request, 'tableReference', None) or getattr( + request, 'tableId', None) or request + table_ref = _to_gcp_table_ref( + t_ref, + default_project=proj or getattr(self._client, 'project', None)) + return self._client.get_table(table_ref) + + def Insert(self, request): + table = getattr(request, 'table', None) + if table is not None: + t_ref = getattr(table, 'tableReference', None) + proj = getattr(t_ref, 'projectId', None) or getattr( + request, 'projectId', None) + ds_id = getattr(t_ref, 'datasetId', None) or getattr( + request, 'datasetId', None) + tbl_id = getattr(t_ref, 'tableId', None) + schema = getattr(table, 'schema', None) + else: + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + tbl_id = getattr(request, 'tableId', None) + schema = getattr(request, 'schema', None) + gcp_tbl_ref = gcp_bigquery.TableReference( + gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', + ds_id), + tbl_id) + gcp_table = gcp_bigquery.Table(gcp_tbl_ref, schema=_to_gcp_schema(schema)) + return self._client.create_table(gcp_table, exists_ok=True) + + def Delete(self, request): + t_ref = getattr(request, 'tableReference', None) + proj = getattr(t_ref, 'projectId', None) or getattr( + request, 'projectId', None) + ds_id = getattr(t_ref, 'datasetId', None) or getattr( + request, 'datasetId', None) + tbl_id = getattr(t_ref, 'tableId', None) or getattr( + request, 'tableId', None) + gcp_tbl_ref = gcp_bigquery.TableReference( + gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', + ds_id), + tbl_id) + return self._client.delete_table(gcp_tbl_ref, not_found_ok=True) + + def List(self, request): + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + ds_ref = gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', ds_id) + return self._client.list_tables(ds_ref) + + def Patch(self, request): + table = getattr(request, 'table', None) + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + tbl_id = getattr(request, 'tableId', None) + gcp_tbl_ref = gcp_bigquery.TableReference( + gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', + ds_id), + tbl_id) + gcp_table = gcp_bigquery.Table(gcp_tbl_ref) + if table and getattr(table, 'schema', None): + gcp_table.schema = _to_gcp_schema(table.schema) + return self._client.update_table(gcp_table, ['schema']) + + def Update(self, request): + return self.Patch(request) + + class _ClientDatasetsCompat: + def __init__(self, client): + self._client = client + + def Get(self, request): + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + ds_ref = gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', ds_id) + return self._client.get_dataset(ds_ref) + + def Insert(self, request): + dataset = getattr(request, 'dataset', None) + ds_ref_raw = getattr( + dataset, 'datasetReference', None) if dataset else None + proj = getattr(ds_ref_raw, 'projectId', None) or getattr( + request, 'projectId', None) + ds_id = getattr(ds_ref_raw, 'datasetId', None) or getattr( + request, 'datasetId', None) + ds_ref = gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', ds_id) + gcp_ds = gcp_bigquery.Dataset(ds_ref) + if dataset: + if getattr(dataset, 'location', None): + gcp_ds.location = dataset.location + if getattr(dataset, 'defaultTableExpirationMs', None): + gcp_ds.default_table_expiration_ms = dataset.defaultTableExpirationMs + return self._client.create_dataset(gcp_ds, exists_ok=True) + + def Delete(self, request): + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + delete_contents = getattr(request, 'deleteContents', True) + ds_ref = gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', ds_id) + return self._client.delete_dataset( + ds_ref, delete_contents=delete_contents, not_found_ok=True) + + def List(self, request): + proj = getattr(request, 'projectId', None) or getattr( + self._client, 'project', None) + return self._client.list_datasets(project=proj) + + def Patch(self, request): + dataset = getattr(request, 'dataset', None) + proj = getattr(request, 'projectId', None) + ds_id = getattr(request, 'datasetId', None) + ds_ref = gcp_bigquery.DatasetReference( + proj or getattr(self._client, 'project', None) or 'default', ds_id) + gcp_ds = gcp_bigquery.Dataset(ds_ref) + fields_to_update = [] + if dataset: + if getattr(dataset, 'defaultTableExpirationMs', None): + gcp_ds.default_table_expiration_ms = dataset.defaultTableExpirationMs + fields_to_update.append('default_table_expiration_ms') + return self._client.update_dataset(gcp_ds, fields_to_update) + + def Update(self, request): + return self.Patch(request) + + class _ClientJobsCompat: + def __init__(self, client): + self._client = client + + def Get(self, request): + proj = getattr(request, 'projectId', None) + job_id = getattr(request, 'jobId', None) + loc = getattr(request, 'location', None) + return self._client.get_job(job_id, project=proj, location=loc) + + def GetQueryResults(self, request): + proj = getattr(request, 'projectId', None) + job_id = getattr(request, 'jobId', None) + loc = getattr(request, 'location', None) + page_token = getattr(request, 'pageToken', None) + job = self._client.get_job(job_id, project=proj, location=loc) + return job.result(page_token=page_token) + + if not hasattr(gcp_bigquery.Client, 'tables'): + gcp_bigquery.Client.tables = property( + lambda self: _ClientTablesCompat(self)) + if not hasattr(gcp_bigquery.Client, 'datasets'): + gcp_bigquery.Client.datasets = property( + lambda self: _ClientDatasetsCompat(self)) + if not hasattr(gcp_bigquery.Client, 'jobs'): + gcp_bigquery.Client.jobs = property(lambda self: _ClientJobsCompat(self)) + _LOGGER = logging.getLogger(__name__) JSON_COMPLIANCE_ERROR = 'NAN, INF and -INF values are not JSON compliant.' @@ -797,12 +971,18 @@ class BigQueryWrapper(object): HISTOGRAM_METRIC_LOGGER = MetricLogger() - def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): + def __init__( + self, + client=None, + temp_dataset_id=None, + temp_table_ref=None, + use_legacy_client=False): if client is not None: self.client = client self.gcp_bq_client = client else: - self.client = BigQueryWrapper._bigquery_client(PipelineOptions()) + self.client = BigQueryWrapper._bigquery_client( + PipelineOptions(), use_legacy_client=use_legacy_client) self.gcp_bq_client = self.client self._unique_row_id = 0 @@ -1530,8 +1710,10 @@ def get_table(self, project_id, dataset_id, table_id): NotFound or HttpError: if lookup failed. """ if self._is_modern_client: - table_ref = TableReference( - projectId=project_id, datasetId=dataset_id, tableId=table_id) + table_ref = _to_gcp_table_ref( + TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id), + default_project=project_id) return self.client.get_table(table_ref) # Fallback for legacy client @@ -1556,8 +1738,10 @@ def _create_table( table_id) if self._is_modern_client: - table_ref = TableReference( - projectId=project_id, datasetId=dataset_id, tableId=table_id) + table_ref = _to_gcp_table_ref( + TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id), + default_project=project_id) gcp_schema = _to_gcp_schema(schema) table = gcp_bigquery.Table(table_ref, schema=gcp_schema) if additional_parameters: @@ -1595,7 +1779,9 @@ def get_or_create_dataset( default_table_expiration_ms=None): # Check if dataset already exists otherwise create it if self._is_modern_client: - dataset_ref = DatasetReference(projectId=project_id, datasetId=dataset_id) + dataset_ref = _to_gcp_dataset_ref( + DatasetReference(projectId=project_id, datasetId=dataset_id), + project=project_id) try: dataset = self.client.get_dataset(dataset_ref) self.created_temp_dataset = False @@ -1666,8 +1852,10 @@ def get_or_create_dataset( retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _is_table_empty(self, project_id, dataset_id, table_id): if self._is_modern_client: - table_ref = TableReference( - projectId=project_id, datasetId=dataset_id, tableId=table_id) + table_ref = _to_gcp_table_ref( + TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id), + default_project=project_id) rows = self.client.list_rows(table_ref, max_results=1) if hasattr(rows, 'total_rows') and rows.total_rows is not None: return rows.total_rows == 0 @@ -1688,8 +1876,10 @@ def _is_table_empty(self, project_id, dataset_id, table_id): retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_table(self, project_id, dataset_id, table_id): if self._is_modern_client: - table_ref = TableReference( - projectId=project_id, datasetId=dataset_id, tableId=table_id) + table_ref = _to_gcp_table_ref( + TableReference( + projectId=project_id, datasetId=dataset_id, tableId=table_id), + default_project=project_id) try: self.client.delete_table(table_ref, not_found_ok=True) except (NotFound, HttpError, ClientError) as exn: @@ -1719,7 +1909,9 @@ def _delete_table(self, project_id, dataset_id, table_id): retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_dataset(self, project_id, dataset_id, delete_contents=True): if self._is_modern_client: - dataset_ref = DatasetReference(projectId=project_id, datasetId=dataset_id) + dataset_ref = _to_gcp_dataset_ref( + DatasetReference(projectId=project_id, datasetId=dataset_id), + project=project_id) try: self.client.delete_dataset( dataset_ref, delete_contents=delete_contents, not_found_ok=True) @@ -1786,8 +1978,10 @@ def create_temporary_dataset( def clean_up_temporary_dataset(self, project_id): temp_table = self._get_temp_table(project_id) if self._is_modern_client: - dataset_ref = DatasetReference( - projectId=project_id, datasetId=temp_table.datasetId) + dataset_ref = _to_gcp_dataset_ref( + DatasetReference( + projectId=project_id, datasetId=temp_table.datasetId), + project=project_id) try: self.client.get_dataset(dataset_ref) except (NotFound, HttpError, ClientError) as exn: @@ -2347,13 +2541,15 @@ def from_pipeline_options(pipeline_options: PipelineOptions): client=BigQueryWrapper._bigquery_client(pipeline_options)) @staticmethod - def _bigquery_client(pipeline_options: PipelineOptions): + def _bigquery_client( + pipeline_options: PipelineOptions, use_legacy_client: bool = False): raw_credentials = auth.get_service_credentials(pipeline_options) google_credentials = ( raw_credentials.get_google_auth_credentials() if hasattr( raw_credentials, 'get_google_auth_credentials') else raw_credentials) project = None + experiments = [] if pipeline_options: try: from apache_beam.options.pipeline_options import GoogleCloudOptions @@ -2362,19 +2558,29 @@ def _bigquery_client(pipeline_options: PipelineOptions): project = project.get() except Exception: project = None - - client_cls = getattr(gcp_bigquery, 'Client', None) - if isinstance(client_cls, type): try: - client_info = ClientInfo( - user_agent="apache-beam-%s" % - apache_beam.__version__) if ClientInfo else None - return gcp_bigquery.Client( - project=project, - credentials=google_credentials, - client_info=client_info) + from apache_beam.options.pipeline_options import DebugOptions + experiments = pipeline_options.view_as(DebugOptions).experiments or [] except Exception: - pass + experiments = [] + + use_legacy = ( + use_legacy_client or 'use_legacy_bigquery_client' in experiments or + 'use_legacy_bq_client' in experiments) + + if not use_legacy: + client_cls = getattr(gcp_bigquery, 'Client', None) + if isinstance(client_cls, type): + try: + client_info = ClientInfo( + user_agent="apache-beam-%s" % + apache_beam.__version__) if ClientInfo else None + return gcp_bigquery.Client( + project=project, + credentials=google_credentials, + client_info=client_info) + except Exception: + pass return apitools_bigquery.BigqueryV2( http=get_new_http(), diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index f78d59047076..cf17e764c27c 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -39,6 +39,7 @@ from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import AvroRowWriter from apache_beam.io.gcp.bigquery_tools import BigQueryJobTypes +from apache_beam.io.gcp.bigquery_tools import BigQueryWrapper from apache_beam.io.gcp.bigquery_tools import JsonRowWriter from apache_beam.io.gcp.bigquery_tools import RowAsDictJsonCoder from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict @@ -50,6 +51,7 @@ from apache_beam.io.gcp.internal.clients import bigquery from apache_beam.metrics import monitoring_infos from apache_beam.metrics.execution import MetricsEnvironment +from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import StaticValueProvider from apache_beam.typehints.row_type import RowTypeConstraint from apache_beam.utils.timestamp import Timestamp @@ -1411,6 +1413,26 @@ def test_type_overrides_json_to_dict(self): self.assertEqual(typehints_dict, [("data", Optional[dict])]) +class TestBigQueryClientExperimentFallback(unittest.TestCase): + def test_default_client_is_modern(self): + wrapper = BigQueryWrapper.from_pipeline_options(PipelineOptions([])) + self.assertTrue(wrapper._is_modern_client) + + def test_experiment_flag_use_legacy_bigquery_client(self): + options = PipelineOptions(['--experiments=use_legacy_bigquery_client']) + wrapper = BigQueryWrapper.from_pipeline_options(options) + self.assertFalse(wrapper._is_modern_client) + + def test_experiment_flag_use_legacy_bq_client(self): + options = PipelineOptions(['--experiments=use_legacy_bq_client']) + wrapper = BigQueryWrapper.from_pipeline_options(options) + self.assertFalse(wrapper._is_modern_client) + + def test_kwarg_use_legacy_client(self): + wrapper = BigQueryWrapper(use_legacy_client=True) + self.assertFalse(wrapper._is_modern_client) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() From 7529003eec56a8adadf84741b8e81cffea2f4f10 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 15:44:53 +0000 Subject: [PATCH 06/12] revert setup.py formatting change --- sdks/python/setup.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 16ab4b63d780..fe23337abc4b 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -150,10 +150,10 @@ def cythonize(*args, **kwargs): pyarrow_dependency = [''] else: pyarrow_dependency = [ - # Generally try to cover versions released in the last two years. - # Update python/sdks/tox.ini to cover the same pyarrow versions - # when updating the bounds here. - 'pyarrow>=14.0.1,<26.0.0', + # Generally try to cover versions released in the last two years. + # Update python/sdks/tox.ini to cover the same pyarrow versions + # when updating the bounds here. + 'pyarrow>=14.0.1,<26.0.0', ] # Exclude pandas<=1.4.2 since it doesn't work with numpy 1.24.x. @@ -199,7 +199,8 @@ def cythonize(*args, **kwargs): # `cls_token=` or tokenizers reinstates `cls=` as a deprecation alias. 'tokenizers<0.23', # Avoid torch 2.12.0+ which fails to run unit tests with segfault. - 'torch<2.12.0', # Match tested transformers range. + 'torch<2.12.0', + # Match tested transformers range. 'transformers>=4.28.0,<4.56.0', # Keep tokenizers compatible with this transformers range. 'tokenizers>=0.13.3,<0.22.0', @@ -459,9 +460,9 @@ def get_portability_package_data(): # BEAM-8840: Do NOT use tests_require or setup_requires. extras_require={ 'dev': [ - 'pyrefly==1.1.1', - 'ruff==0.15.22', - 'yapf==0.43.0', + 'pyrefly==1.1.1', + 'ruff==0.15.22', + 'yapf==0.43.0', ], 'dill': [ # Dill doesn't have forwards-compatibility guarantees within minor @@ -599,7 +600,7 @@ def get_portability_package_data(): 'tf2onnx>=1.17.0,<1.18', ] + ml_base_core, 'p310_ml_test': [ - 'datatable', + 'datatable', ] + ml_base + qdrant_dependency, 'p312_ml_test': [ 'datatable', From bf871a132b6534fd912a45c8164bc3acf43a7713 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 17:10:10 +0000 Subject: [PATCH 07/12] skip fallback check if missing imports --- sdks/python/apache_beam/io/gcp/bigquery_tools_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index cf17e764c27c..07ecff58ce07 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -1413,6 +1413,7 @@ def test_type_overrides_json_to_dict(self): self.assertEqual(typehints_dict, [("data", Optional[dict])]) +@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQueryClientExperimentFallback(unittest.TestCase): def test_default_client_is_modern(self): wrapper = BigQueryWrapper.from_pipeline_options(PipelineOptions([])) From 3f9e010f26e3207298ad1e7f814efaffa626c31e Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 18:26:31 +0000 Subject: [PATCH 08/12] address integration test failures, expand testing --- sdks/python/apache_beam/io/gcp/bigquery.py | 25 ++- .../apache_beam/io/gcp/bigquery_tools.py | 165 ++++++++++++++---- .../apache_beam/io/gcp/bigquery_tools_test.py | 109 ++++++++++++ 3 files changed, 253 insertions(+), 46 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 09e957d25e31..0b9902e441f9 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -780,8 +780,10 @@ def _get_temp_dataset_id(self): elif hasattr(self.temp_dataset, 'dataset_id'): return self.temp_dataset.dataset_id elif isinstance(self.temp_dataset, str): - if ':' in self.temp_dataset or '.' in self.temp_dataset: - return bigquery_tools.parse_table_reference(self.temp_dataset).datasetId + if ':' in self.temp_dataset: + return self.temp_dataset.split(':', 1)[1] + elif '.' in self.temp_dataset: + return self.temp_dataset.split('.', 1)[1] return self.temp_dataset return None @@ -790,9 +792,11 @@ def _get_temp_dataset_project(self): return self.temp_dataset.projectId elif hasattr(self.temp_dataset, 'project') and self.temp_dataset.project: return self.temp_dataset.project - elif isinstance(self.temp_dataset, str) and (':' in self.temp_dataset or - '.' in self.temp_dataset): - return bigquery_tools.parse_table_reference(self.temp_dataset).projectId + elif isinstance(self.temp_dataset, str): + if ':' in self.temp_dataset: + return self.temp_dataset.split(':', 1)[0] + elif '.' in self.temp_dataset: + return self.temp_dataset.split('.', 1)[0] return None def _get_project(self): @@ -3119,9 +3123,14 @@ def _expand_direct_read(self, pcoll): temp_dataset = self._kwargs.get('temp_dataset') if temp_dataset is not None: if isinstance(temp_dataset, str): - ds_ref = bigquery_tools.parse_table_reference(temp_dataset) - project_id = ds_ref.projectId - dataset_id = ds_ref.datasetId + if ':' in temp_dataset: + project_id, dataset_id = temp_dataset.split(':', 1) + elif '.' in temp_dataset: + project_id, dataset_id = temp_dataset.split('.', 1) + else: + project_id = pcoll.pipeline.options.view_as( + GoogleCloudOptions).project + dataset_id = temp_dataset else: project_id = getattr(temp_dataset, 'projectId', None) or getattr( temp_dataset, 'project', None) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 0fceae61648f..9f7c02119a35 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -133,7 +133,6 @@ TableCell = getattr(bigquery, 'TableCell', None) Table = getattr(bigquery, 'Table', None) Dataset = getattr(bigquery, 'Dataset', None) - JobReference = getattr(bigquery, 'JobReference', None) Job = getattr(bigquery, 'Job', None) JobConfiguration = getattr(bigquery, 'JobConfiguration', None) JobConfigurationLoad = getattr(bigquery, 'JobConfigurationLoad', None) @@ -292,34 +291,6 @@ def fields(self, value): if value: self.extend(value) - class JobReference(object): - def __init__( - self, - jobId=None, - projectId=None, - location=None, - job_id=None, - project=None): - self.jobId = jobId if jobId is not None else job_id - self.projectId = projectId if projectId is not None else project - self.location = location - - @property - def job_id(self): - return self.jobId - - @job_id.setter - def job_id(self, val): - self.jobId = val - - @property - def project(self): - return self.projectId - - @project.setter - def project(self, val): - self.projectId = val - class TableCell(object): def __init__(self, v=None): self.v = v @@ -341,18 +312,120 @@ def __init__(self, f=None): JobStatistics4 = None ErrorProto = None + +class JobReference(object): + """Compatibility model for BigQuery JobReference. + + Supports both camelCase (jobId, projectId) and snake_case (job_id, project, project_id) + initialization and attribute access. + """ + def __init__( + self, + jobId=None, + projectId=None, + location=None, + job_id=None, + project=None, + project_id=None): + self.jobId = jobId if jobId is not None else job_id + self.projectId = ( + projectId if projectId is not None else + (project if project is not None else project_id)) + self.location = location + + @property + def job_id(self): + return self.jobId + + @job_id.setter + def job_id(self, val): + self.jobId = val + + @property + def project(self): + return self.projectId + + @project.setter + def project(self, val): + self.projectId = val + + @property + def project_id(self): + return self.projectId + + @project_id.setter + def project_id(self, val): + self.projectId = val + + def __eq__(self, other): + if other is None: + return False + if isinstance(other, JobReference): + return ( + self.jobId == other.jobId and self.projectId == other.projectId and + self.location == other.location) + if apitools_bigquery and hasattr(apitools_bigquery, + 'JobReference') and isinstance( + other, apitools_bigquery.JobReference): + return ( + self.jobId == getattr(other, 'jobId', None) and + self.projectId == getattr(other, 'projectId', None) and + self.location == getattr(other, 'location', None)) + return NotImplemented + + def __hash__(self): + return hash((self.jobId, self.projectId, self.location)) + + def __repr__(self): + return ( + f"JobReference(jobId={self.jobId!r}, " + f"projectId={self.projectId!r}, " + f"location={self.location!r})") + + +try: + from apitools.base.protorpclite import messages as _protorpclite_messages + if hasattr(_protorpclite_messages, 'Message'): + _orig_message_eq = _protorpclite_messages.Message.__eq__ + + def _message_compat_eq(self, other): + if isinstance(other, JobReference) and apitools_bigquery and hasattr( + apitools_bigquery, 'JobReference') and isinstance( + self, apitools_bigquery.JobReference): + return ( + getattr(self, 'jobId', None) == other.jobId and + getattr(self, 'projectId', None) == other.projectId and + getattr(self, 'location', None) == other.location) + return _orig_message_eq(self, other) + + _protorpclite_messages.Message.__eq__ = _message_compat_eq +except ImportError: + pass + + +def _set_table_ref_prop(ref, prop, val): + if hasattr(ref, '_properties') and isinstance(ref._properties, dict): + ref._properties[prop] = val + if prop == 'projectId': + setattr(ref, '_project', val) + elif prop == 'datasetId': + setattr(ref, '_dataset_id', val) + elif prop == 'tableId': + setattr(ref, '_table_id', val) + + # Monkey-patch gcp_bigquery classes to ensure full backward compatibility if gcp_bigquery: if not hasattr(gcp_bigquery.TableReference, 'projectId'): gcp_bigquery.TableReference.projectId = property( lambda self: self.project, - lambda self, val: setattr(self, '_project', val)) + lambda self, val: _set_table_ref_prop(self, 'projectId', val)) gcp_bigquery.TableReference.datasetId = property( lambda self: self.dataset_id, - lambda self, val: setattr(self, '_dataset_id', val)) + lambda self, val: _set_table_ref_prop(self, 'datasetId', val)) gcp_bigquery.TableReference.tableId = property( lambda self: self.table_id, - lambda self, val: setattr(self, '_table_id', val)) + lambda self, val: _set_table_ref_prop(self, 'tableId', val)) if not hasattr(gcp_bigquery.DatasetReference, 'projectId'): gcp_bigquery.DatasetReference.projectId = property( @@ -464,7 +537,7 @@ def _to_gcp_dataset_ref(dataset_ref, project=None): if gcp_bigquery is not None and hasattr(gcp_bigquery.DatasetReference, 'from_string'): return gcp_bigquery.DatasetReference.from_string( - dataset_ref, default_project=project) + dataset_ref.replace(':', '.'), default_project=project) proj = project or 'default' if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): return gcp_bigquery.DatasetReference(proj, dataset_ref) @@ -1147,7 +1220,11 @@ def get_query_location(self, project_id, query, use_legacy_sql): return None # Fallback if legacy client.jobs.Insert is mocked - reference = JobReference(jobId=uuid.uuid4().hex, projectId=project_id) + reference = ( + apitools_bigquery.JobReference( + jobId=uuid.uuid4().hex, projectId=project_id) + if apitools_bigquery and hasattr(apitools_bigquery, 'JobReference') else + JobReference(jobId=uuid.uuid4().hex, projectId=project_id)) request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=apitools_bigquery.Job( @@ -1231,7 +1308,10 @@ def _insert_copy_job( raise # Fallback if legacy client.jobs.Insert is mocked - reference = JobReference(jobId=job_id, projectId=project_id) + reference = ( + apitools_bigquery.JobReference(jobId=job_id, projectId=project_id) + if apitools_bigquery and hasattr(apitools_bigquery, 'JobReference') else + JobReference(jobId=job_id, projectId=project_id)) request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=apitools_bigquery.Job( @@ -1331,7 +1411,10 @@ def _insert_load_job( {'fields': [table_field_to_dict(f) for f in schema]}) else: job_schema = schema - reference = JobReference(jobId=job_id, projectId=project_id) + reference = ( + apitools_bigquery.JobReference(jobId=job_id, projectId=project_id) + if apitools_bigquery and hasattr(apitools_bigquery, 'JobReference') else + JobReference(jobId=job_id, projectId=project_id)) request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=apitools_bigquery.Job( @@ -1460,7 +1543,10 @@ def _start_query_job( raise # Fallback if legacy client.jobs.Insert is mocked - reference = JobReference(jobId=job_id, projectId=project_id) + reference = ( + apitools_bigquery.JobReference(jobId=job_id, projectId=project_id) + if apitools_bigquery and hasattr(apitools_bigquery, 'JobReference') else + JobReference(jobId=job_id, projectId=project_id)) request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=project_id, job=apitools_bigquery.Job( @@ -2209,7 +2295,10 @@ def perform_extract_job( raise # Fallback for legacy client - job_reference = JobReference(jobId=job_id, projectId=job_project) + job_reference = ( + apitools_bigquery.JobReference(jobId=job_id, projectId=job_project) + if apitools_bigquery and hasattr(apitools_bigquery, 'JobReference') else + JobReference(jobId=job_id, projectId=job_project)) request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=job_project, job=apitools_bigquery.Job( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index 07ecff58ce07..c8b9ce70666c 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -35,6 +35,7 @@ from parameterized import parameterized import apache_beam as beam +from apache_beam.io.gcp import bigquery_tools from apache_beam.io.gcp import resource_identifiers from apache_beam.io.gcp.bigquery_tools import JSON_COMPLIANCE_ERROR from apache_beam.io.gcp.bigquery_tools import AvroRowWriter @@ -1434,6 +1435,114 @@ def test_kwarg_use_legacy_client(self): self.assertFalse(wrapper._is_modern_client) +class TestJobReferenceCompatibility(unittest.TestCase): + def test_init_camel_case(self): + ref = bigquery_tools.JobReference( + jobId='test-job', projectId='test-proj', location='US') + self.assertEqual(ref.jobId, 'test-job') + self.assertEqual(ref.job_id, 'test-job') + self.assertEqual(ref.projectId, 'test-proj') + self.assertEqual(ref.project, 'test-proj') + self.assertEqual(ref.project_id, 'test-proj') + self.assertEqual(ref.location, 'US') + + def test_init_snake_case(self): + ref = bigquery_tools.JobReference( + job_id='test-job', project='test-proj', location='EU') + self.assertEqual(ref.jobId, 'test-job') + self.assertEqual(ref.job_id, 'test-job') + self.assertEqual(ref.projectId, 'test-proj') + self.assertEqual(ref.project, 'test-proj') + self.assertEqual(ref.project_id, 'test-proj') + self.assertEqual(ref.location, 'EU') + + def test_init_with_project_id(self): + ref = bigquery_tools.JobReference( + job_id='test-job', project_id='test-proj', location='EU') + self.assertEqual(ref.jobId, 'test-job') + self.assertEqual(ref.job_id, 'test-job') + self.assertEqual(ref.projectId, 'test-proj') + self.assertEqual(ref.project, 'test-proj') + self.assertEqual(ref.project_id, 'test-proj') + + def test_attribute_setters(self): + ref = bigquery_tools.JobReference() + ref.job_id = 'j1' + ref.project = 'p1' + self.assertEqual(ref.jobId, 'j1') + self.assertEqual(ref.job_id, 'j1') + self.assertEqual(ref.projectId, 'p1') + self.assertEqual(ref.project, 'p1') + self.assertEqual(ref.project_id, 'p1') + ref.project_id = 'p2' + self.assertEqual(ref.projectId, 'p2') + self.assertEqual(ref.project, 'p2') + self.assertEqual(ref.project_id, 'p2') + + def test_equality_with_custom_and_apitools(self): + ref1 = bigquery_tools.JobReference(job_id='j1', project='p1', location='US') + ref2 = bigquery_tools.JobReference( + jobId='j1', projectId='p1', location='US') + ref3 = bigquery_tools.JobReference( + jobId='j2', projectId='p1', location='US') + self.assertEqual(ref1, ref2) + self.assertNotEqual(ref1, ref3) + self.assertEqual(hash(ref1), hash(ref2)) + + if bigquery is not None and hasattr(bigquery, 'JobReference'): + ap_ref = bigquery.JobReference(jobId='j1', projectId='p1', location='US') + self.assertEqual(ref1, ap_ref) + self.assertEqual(ap_ref, ref1) + + def test_equality_type_safety(self): + empty_ref = bigquery_tools.JobReference() + self.assertNotEqual(empty_ref, None) + self.assertNotEqual(empty_ref, 123) + self.assertNotEqual(empty_ref, "") + self.assertNotEqual(empty_ref, {}) + + ref = bigquery_tools.JobReference(project='p1') + table_ref = bigquery_tools.TableReference( + projectId='p1', datasetId='d1', tableId='t1') + self.assertNotEqual(ref, table_ref) + self.assertNotEqual(empty_ref, table_ref) + + def test_pickle_and_coder_roundtrip(self): + import pickle + ref = bigquery_tools.JobReference(job_id='j1', project='p1', location='US') + pickled = pickle.dumps(ref) + unpickled = pickle.loads(pickled) + self.assertEqual(ref, unpickled) + self.assertEqual(unpickled.jobId, 'j1') + self.assertEqual(unpickled.projectId, 'p1') + self.assertEqual(unpickled.location, 'US') + + coder = beam.coders.FastPrimitivesCoder() + encoded = coder.encode(ref) + decoded = coder.decode(encoded) + self.assertEqual(ref, decoded) + + def test_table_reference_property_mutability(self): + from google.cloud import bigquery as gcp_bigquery + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'TableReference'): + ds = gcp_bigquery.DatasetReference('p1', 'd1') + table = gcp_bigquery.TableReference(ds, 't1') + table.tableId = 't2' + self.assertEqual(table.tableId, 't2') + self.assertEqual(table.table_id, 't2') + table.datasetId = 'd2' + self.assertEqual(table.datasetId, 'd2') + self.assertEqual(table.dataset_id, 'd2') + table.projectId = 'p2' + self.assertEqual(table.projectId, 'p2') + self.assertEqual(table.project, 'p2') + + def test_to_gcp_dataset_ref_colon_format(self): + ds_ref = bigquery_tools._to_gcp_dataset_ref('my-project:my_dataset') + self.assertEqual(ds_ref.project, 'my-project') + self.assertEqual(ds_ref.dataset_id, 'my_dataset') + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() From 795877c8ce6fd30cd241ee82b3c3da9ec2b7ad35 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 19:02:17 +0000 Subject: [PATCH 09/12] refactor compat classes, fix fallback --- .../apache_beam/io/gcp/bigquery_tools.py | 477 ++++++++++++------ .../apache_beam/io/gcp/bigquery_tools_test.py | 65 ++- 2 files changed, 366 insertions(+), 176 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 9f7c02119a35..194fe56fdd52 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -124,13 +124,289 @@ # (e.g. projectId, datasetId, tableId, tableReference) for backwards # compatibility across pipelines, transforms, and test suites. + +class _DatasetReferenceCompat(object): + """Compatibility model for BigQuery DatasetReference when google-cloud-bigquery is unavailable. + + Supports both camelCase (projectId, datasetId) and snake_case (project, dataset_id, project_id). + """ + def __init__( + self, + project=None, + dataset_id=None, + projectId=None, + datasetId=None, + project_id=None): + p = ( + projectId if projectId is not None else + (project_id if project_id is not None else project)) + d = datasetId if datasetId is not None else dataset_id + self._project = p or '' + self._dataset_id = d or '' + + @classmethod + def from_string(cls, dataset_ref, default_project=None): + if ':' in dataset_ref: + p, d = dataset_ref.split(':', 1) + elif '.' in dataset_ref: + parts = dataset_ref.split('.', 1) + p, d = parts[0], parts[1] + else: + p, d = default_project or 'default', dataset_ref + return cls(project=p, dataset_id=d) + + @property + def projectId(self): + return self._project + + @projectId.setter + def projectId(self, val): + self._project = val + + @property + def project(self): + return self._project + + @project.setter + def project(self, val): + self._project = val + + @property + def project_id(self): + return self._project + + @project_id.setter + def project_id(self, val): + self._project = val + + @property + def datasetId(self): + return self._dataset_id + + @datasetId.setter + def datasetId(self, val): + self._dataset_id = val + + @property + def dataset_id(self): + return self._dataset_id + + @dataset_id.setter + def dataset_id(self, val): + self._dataset_id = val + + def __repr__(self): + return f"DatasetReference('{self.project}', '{self.dataset_id}')" + + def __eq__(self, other): + if other is None: + return False + if not hasattr(other, 'project') and not hasattr(other, 'projectId'): + return NotImplemented + other_p = getattr(other, 'projectId', None) or getattr( + other, 'project', None) + other_d = getattr(other, 'datasetId', None) or getattr( + other, 'dataset_id', None) + return (self.projectId, self.datasetId) == (other_p, other_d) + + def __hash__(self): + return hash((self.projectId, self.datasetId)) + + +class _TableReferenceCompat(object): + """Compatibility model for BigQuery TableReference when google-cloud-bigquery is unavailable. + + Supports both camelCase (projectId, datasetId, tableId) and snake_case + (project, dataset_id, table_id, project_id). + """ + def __init__( + self, + dataset_ref=None, + table_id=None, + projectId=None, + datasetId=None, + tableId=None, + project=None, + dataset_id=None, + project_id=None): + p = ( + projectId if projectId is not None else + (project_id if project_id is not None else project)) + d = datasetId if datasetId is not None else dataset_id + t = tableId if tableId is not None else table_id + if p is not None or d is not None or t is not None: + self._project = p + self._dataset_id = d + self._table_id = t + elif dataset_ref is not None: + self._project = getattr(dataset_ref, 'projectId', None) or getattr( + dataset_ref, 'project', None) + self._dataset_id = getattr(dataset_ref, 'datasetId', None) or getattr( + dataset_ref, 'dataset_id', None) + self._table_id = table_id or '' + else: + self._project = None + self._dataset_id = None + self._table_id = None + + @classmethod + def from_string(cls, table_ref, default_project=None): + parsed = parse_table_reference(table_ref, project=default_project) + return cls( + projectId=parsed.projectId, + datasetId=parsed.datasetId, + tableId=parsed.tableId) + + @property + def projectId(self): + return self._project + + @projectId.setter + def projectId(self, val): + self._project = val + + @property + def project(self): + return self._project + + @project.setter + def project(self, val): + self._project = val + + @property + def project_id(self): + return self._project + + @project_id.setter + def project_id(self, val): + self._project = val + + @property + def datasetId(self): + return self._dataset_id + + @datasetId.setter + def datasetId(self, val): + self._dataset_id = val + + @property + def dataset_id(self): + return self._dataset_id + + @dataset_id.setter + def dataset_id(self, val): + self._dataset_id = val + + @property + def tableId(self): + return self._table_id + + @tableId.setter + def tableId(self, val): + self._table_id = val + + @property + def table_id(self): + return self._table_id + + @table_id.setter + def table_id(self, val): + self._table_id = val + + @property + def dataset_reference(self): + return _DatasetReferenceCompat( + projectId=self.projectId, datasetId=self.datasetId) + + @property + def datasetReference(self): + return self.dataset_reference + + def __repr__(self): + return ( + f"TableReference(projectId='{self.projectId}', " + f"datasetId='{self.datasetId}', tableId='{self.tableId}')") + + def __eq__(self, other): + if other is None: + return False + if not hasattr(other, 'tableId') and not hasattr(other, 'table_id'): + return NotImplemented + other_p = getattr(other, 'projectId', None) or getattr( + other, 'project', None) + other_d = getattr(other, 'datasetId', None) or getattr( + other, 'dataset_id', None) + other_t = getattr(other, 'tableId', None) or getattr( + other, 'table_id', None) + return (self.projectId, self.datasetId, + self.tableId) == (other_p, other_d, other_t) + + def __hash__(self): + return hash((self.projectId, self.datasetId, self.tableId)) + + +class _TableFieldSchemaCompat(object): + def __init__( + self, + name='', + type='STRING', + mode='NULLABLE', + description=None, + fields=(), + field_type=None, + **kwargs): + ft = type or field_type or 'STRING' + self.name = name + self.field_type = ft + self.mode = mode or 'NULLABLE' + self.description = description + self.fields = list(fields) if fields else [] + + @property + def type(self): + return self.field_type + + @type.setter + def type(self, val): + self.field_type = val + + +class _TableSchemaCompat(list): + def __init__(self, fields=None): + if fields: + super().__init__(fields) + else: + super().__init__() + + @property + def fields(self): + return self + + @fields.setter + def fields(self, value): + self.clear() + if value: + self.extend(value) + + +class _TableCellCompat(object): + def __init__(self, v=None): + self.v = v + + +class _TableRowCompat(object): + def __init__(self, f=None): + self.f = f or [] + + if bigquery is not None and hasattr(bigquery, 'TableReference'): TableReference = bigquery.TableReference - DatasetReference = getattr(bigquery, 'DatasetReference', None) + DatasetReference = getattr( + bigquery, 'DatasetReference', None) or _DatasetReferenceCompat TableFieldSchema = bigquery.TableFieldSchema TableSchema = bigquery.TableSchema - TableRow = getattr(bigquery, 'TableRow', None) - TableCell = getattr(bigquery, 'TableCell', None) + TableRow = getattr(bigquery, 'TableRow', None) or _TableRowCompat + TableCell = getattr(bigquery, 'TableCell', None) or _TableCellCompat Table = getattr(bigquery, 'Table', None) Dataset = getattr(bigquery, 'Dataset', None) Job = getattr(bigquery, 'Job', None) @@ -145,160 +421,12 @@ JobStatistics4 = getattr(bigquery, 'JobStatistics4', None) ErrorProto = getattr(bigquery, 'ErrorProto', None) else: - - class DatasetReference( - gcp_bigquery.DatasetReference if gcp_bigquery else object): - def __init__( - self, project=None, dataset_id=None, projectId=None, datasetId=None): - p = projectId if projectId is not None else project - d = datasetId if datasetId is not None else dataset_id - if gcp_bigquery: - super().__init__(p or '', d or '') - else: - self._project = p or '' - self._dataset_id = d or '' - - @property - def projectId(self): - return self._project - - @projectId.setter - def projectId(self, val): - self._project = val - - @property - def datasetId(self): - return self._dataset_id - - @datasetId.setter - def datasetId(self, val): - self._dataset_id = val - - class TableReference(gcp_bigquery.TableReference if gcp_bigquery else object): - def __init__( - self, - dataset_ref=None, - table_id=None, - projectId=None, - datasetId=None, - tableId=None, - project=None, - dataset_id=None): - p = projectId if projectId is not None else project - d = datasetId if datasetId is not None else dataset_id - t = tableId if tableId is not None else table_id - if p is not None or d is not None or t is not None: - ds_ref = DatasetReference(p, d) if (p or d) else DatasetReference( - '', '') - if gcp_bigquery: - super().__init__(ds_ref, t or '') - else: - self._project = p - self._dataset_id = d - self._table_id = t - elif dataset_ref is not None: - if gcp_bigquery: - super().__init__(dataset_ref, table_id or '') - else: - self._project = getattr(dataset_ref, 'projectId', None) or getattr( - dataset_ref, 'project', None) - self._dataset_id = getattr(dataset_ref, 'datasetId', None) or getattr( - dataset_ref, 'dataset_id', None) - self._table_id = table_id or '' - else: - if gcp_bigquery: - super().__init__(DatasetReference('', ''), '') - else: - self._project = None - self._dataset_id = None - self._table_id = None - - @property - def projectId(self): - return self.project if hasattr(self, 'project') else getattr( - self, '_project', None) - - @projectId.setter - def projectId(self, val): - self._project = val - - @property - def datasetId(self): - return self.dataset_id if hasattr(self, 'dataset_id') else getattr( - self, '_dataset_id', None) - - @datasetId.setter - def datasetId(self, val): - self._dataset_id = val - - @property - def tableId(self): - return self.table_id if hasattr(self, 'table_id') else getattr( - self, '_table_id', None) - - @tableId.setter - def tableId(self, val): - self._table_id = val - - class TableFieldSchema(gcp_bigquery.SchemaField if gcp_bigquery else object): - def __init__( - self, - name='', - type='STRING', - mode='NULLABLE', - description=None, - fields=(), - field_type=None, - **kwargs): - ft = type or field_type or 'STRING' - if gcp_bigquery: - super().__init__( - name=name, - field_type=ft, - mode=mode or 'NULLABLE', - description=description, - fields=fields or (), - **kwargs) - else: - self.name = name - self.field_type = ft - self.mode = mode or 'NULLABLE' - self.description = description - self.fields = list(fields) if fields else [] - - @property - def type(self): - return self.field_type - - @type.setter - def type(self, val): - self._field_type = val - - class TableSchema(list): - def __init__(self, fields=None): - if fields: - super().__init__(fields) - else: - super().__init__() - - @property - def fields(self): - return self - - @fields.setter - def fields(self, value): - self.clear() - if value: - self.extend(value) - - class TableCell(object): - def __init__(self, v=None): - self.v = v - - class TableRow(object): - def __init__(self, f=None): - self.f = f or [] - + TableReference = _TableReferenceCompat + DatasetReference = _DatasetReferenceCompat + TableFieldSchema = _TableFieldSchemaCompat + TableSchema = _TableSchemaCompat + TableRow = _TableRowCompat + TableCell = _TableCellCompat Table = None Dataset = None Job = None @@ -396,6 +524,19 @@ def _message_compat_eq(self, other): getattr(self, 'jobId', None) == other.jobId and getattr(self, 'projectId', None) == other.projectId and getattr(self, 'location', None) == other.location) + if isinstance(other, TableReference) and apitools_bigquery and hasattr( + apitools_bigquery, 'TableReference') and isinstance( + self, apitools_bigquery.TableReference): + return ( + getattr(self, 'projectId', None) == other.projectId and + getattr(self, 'datasetId', None) == other.datasetId and + getattr(self, 'tableId', None) == other.tableId) + if isinstance(other, DatasetReference) and apitools_bigquery and hasattr( + apitools_bigquery, 'DatasetReference') and isinstance( + self, apitools_bigquery.DatasetReference): + return ( + getattr(self, 'projectId', None) == other.projectId and + getattr(self, 'datasetId', None) == other.datasetId) return _orig_message_eq(self, other) _protorpclite_messages.Message.__eq__ = _message_compat_eq @@ -520,7 +661,7 @@ def _to_gcp_table_ref(table_ref, default_project=None): 'DatasetReference'): return gcp_bigquery.TableReference( gcp_bigquery.DatasetReference(proj, dataset_id), table_id) - return TableReference( + return _TableReferenceCompat( projectId=proj, datasetId=dataset_id, tableId=table_id) return table_ref @@ -538,15 +679,25 @@ def _to_gcp_dataset_ref(dataset_ref, project=None): 'from_string'): return gcp_bigquery.DatasetReference.from_string( dataset_ref.replace(':', '.'), default_project=project) + if ':' in dataset_ref: + proj, ds_id = dataset_ref.split(':', 1) + else: + parts = dataset_ref.split('.', 1) + proj, ds_id = parts[0], parts[1] + return _DatasetReferenceCompat(projectId=proj, datasetId=ds_id) proj = project or 'default' if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): return gcp_bigquery.DatasetReference(proj, dataset_ref) - return DatasetReference(projectId=proj, datasetId=dataset_ref) - if hasattr(dataset_ref, 'projectId') and hasattr(dataset_ref, 'datasetId'): - proj = getattr(dataset_ref, 'projectId', None) or project or 'default' + return _DatasetReferenceCompat(projectId=proj, datasetId=dataset_ref) + if hasattr(dataset_ref, 'projectId') or hasattr(dataset_ref, 'project'): + proj = getattr(dataset_ref, 'projectId', None) or getattr( + dataset_ref, 'project', None) or getattr( + dataset_ref, 'project_id', None) or project or 'default' + ds_id = getattr(dataset_ref, 'datasetId', None) or getattr( + dataset_ref, 'dataset_id', None) if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): - return gcp_bigquery.DatasetReference(proj, dataset_ref.datasetId) - return DatasetReference(projectId=proj, datasetId=dataset_ref.datasetId) + return gcp_bigquery.DatasetReference(proj, ds_id) + return _DatasetReferenceCompat(projectId=proj, datasetId=ds_id) return dataset_ref diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index c8b9ce70666c..6961332308bf 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -1522,25 +1522,64 @@ def test_pickle_and_coder_roundtrip(self): decoded = coder.decode(encoded) self.assertEqual(ref, decoded) + @unittest.skipIf( + bigquery_tools.gcp_bigquery is None, + 'google-cloud-bigquery not installed') def test_table_reference_property_mutability(self): - from google.cloud import bigquery as gcp_bigquery - if gcp_bigquery is not None and hasattr(gcp_bigquery, 'TableReference'): - ds = gcp_bigquery.DatasetReference('p1', 'd1') - table = gcp_bigquery.TableReference(ds, 't1') - table.tableId = 't2' - self.assertEqual(table.tableId, 't2') - self.assertEqual(table.table_id, 't2') - table.datasetId = 'd2' - self.assertEqual(table.datasetId, 'd2') - self.assertEqual(table.dataset_id, 'd2') - table.projectId = 'p2' - self.assertEqual(table.projectId, 'p2') - self.assertEqual(table.project, 'p2') + gcp_bigquery = bigquery_tools.gcp_bigquery + ds = gcp_bigquery.DatasetReference('p1', 'd1') + table = gcp_bigquery.TableReference(ds, 't1') + table.tableId = 't2' + self.assertEqual(table.tableId, 't2') + self.assertEqual(table.table_id, 't2') + table.datasetId = 'd2' + self.assertEqual(table.datasetId, 'd2') + self.assertEqual(table.dataset_id, 'd2') + table.projectId = 'p2' + self.assertEqual(table.projectId, 'p2') + self.assertEqual(table.project, 'p2') + + def test_table_reference_compat_model(self): + table = bigquery_tools._TableReferenceCompat( + projectId='p1', datasetId='d1', tableId='t1') + self.assertEqual(table.projectId, 'p1') + self.assertEqual(table.project, 'p1') + self.assertEqual(table.project_id, 'p1') + self.assertEqual(table.datasetId, 'd1') + self.assertEqual(table.dataset_id, 'd1') + self.assertEqual(table.tableId, 't1') + self.assertEqual(table.table_id, 't1') + table.tableId = 't2' + self.assertEqual(table.tableId, 't2') + self.assertEqual(table.table_id, 't2') + table.datasetId = 'd2' + self.assertEqual(table.datasetId, 'd2') + self.assertEqual(table.dataset_id, 'd2') + table.projectId = 'p2' + self.assertEqual(table.projectId, 'p2') + self.assertEqual(table.project, 'p2') + self.assertEqual(table.project_id, 'p2') + + def test_dataset_reference_compat_model(self): + ds = bigquery_tools._DatasetReferenceCompat(projectId='p1', datasetId='d1') + self.assertEqual(ds.projectId, 'p1') + self.assertEqual(ds.project, 'p1') + self.assertEqual(ds.project_id, 'p1') + self.assertEqual(ds.datasetId, 'd1') + self.assertEqual(ds.dataset_id, 'd1') + ds.projectId = 'p2' + self.assertEqual(ds.projectId, 'p2') + self.assertEqual(ds.project, 'p2') + ds.datasetId = 'd2' + self.assertEqual(ds.datasetId, 'd2') + self.assertEqual(ds.dataset_id, 'd2') def test_to_gcp_dataset_ref_colon_format(self): ds_ref = bigquery_tools._to_gcp_dataset_ref('my-project:my_dataset') self.assertEqual(ds_ref.project, 'my-project') self.assertEqual(ds_ref.dataset_id, 'my_dataset') + self.assertEqual(ds_ref.projectId, 'my-project') + self.assertEqual(ds_ref.datasetId, 'my_dataset') if __name__ == '__main__': From 97eb9b1a27513ae78183db87e9e2c2bb2d930736 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Wed, 26 Aug 2026 19:54:51 +0000 Subject: [PATCH 10/12] fix fields list --- sdks/python/apache_beam/io/gcp/bigquery.py | 8 +++- .../io/gcp/bigquery_read_internal.py | 14 ++++++- .../apache_beam/io/gcp/bigquery_test.py | 38 +++++++++++++++++++ .../apache_beam/io/gcp/bigquery_tools.py | 8 +++- 4 files changed, 63 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 0b9902e441f9..04eafa2cf53a 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -529,8 +529,12 @@ def __init__(self, table_schema=None): self.table_schema = table_schema # Precompute field names since we need them for row encoding. if self.table_schema: - self.field_names = tuple(fs.name for fs in self.table_schema.fields) - self.field_types = tuple(fs.type for fs in self.table_schema.fields) + fields = ( + self.table_schema.fields + if hasattr(self.table_schema, 'fields') else self.table_schema) + self.field_names = tuple(fs.name for fs in fields) + self.field_types = tuple( + getattr(fs, 'type', getattr(fs, 'field_type', None)) for fs in fields) def encode(self, table_row): if self.table_schema is None: diff --git a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py index e06df43e8d98..15449e6da780 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal.py @@ -423,7 +423,13 @@ def _get_project(self): class _JsonToDictCoder(coders.Coder): """A coder for a JSON string to a Python dict.""" def __init__(self, table_schema): - self.fields = self._convert_to_tuple(table_schema.fields) + if table_schema is not None: + fields = ( + table_schema.fields + if hasattr(table_schema, 'fields') else table_schema) + self.fields = self._convert_to_tuple(fields) + else: + self.fields = [] self._converters = { 'INTEGER': int, 'INT64': int, @@ -452,7 +458,11 @@ def _convert_to_tuple(cls, table_field_schemas): return [] return [ - FieldSchema(cls._convert_to_tuple(x.fields), x.mode, x.name, x.type) + FieldSchema( + cls._convert_to_tuple(getattr(x, 'fields', ())), + getattr(x, 'mode', 'NULLABLE'), + getattr(x, 'name', ''), + getattr(x, 'field_type', getattr(x, 'type', 'STRING'))) for x in table_field_schemas ] diff --git a/sdks/python/apache_beam/io/gcp/bigquery_test.py b/sdks/python/apache_beam/io/gcp/bigquery_test.py index bd32f4657589..a53c10d18a77 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_test.py @@ -230,6 +230,26 @@ def test_invalid_json_inf(self): def test_invalid_json_neg_inf(self): self.json_compliance_exception(float('-inf')) + def test_schema_as_list(self): + if gcp_bigquery is not None: + schema_fields = [ + gcp_bigquery.SchemaField('s', 'STRING'), + gcp_bigquery.SchemaField('i', 'INTEGER'), + ] + else: + schema_fields = [ + bigquery_tools._TableFieldSchemaCompat('s', 'STRING'), + bigquery_tools._TableFieldSchemaCompat('i', 'INTEGER'), + ] + coder = TableRowJsonCoder(table_schema=schema_fields) + test_row = bigquery.TableRow( + f=[ + bigquery.TableCell(v=to_json_value('abc')), + bigquery.TableCell(v=to_json_value(123)) + ]) + encoded = coder.encode(test_row) + self.assertEqual('{"s": "abc", "i": 123}', encoded) + @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestJsonToDictCoder(unittest.TestCase): @@ -331,6 +351,24 @@ def test_repeatable_field_is_properly_converted(self): actual = coder.decode(input_row) self.assertEqual(expected_row, actual) + def test_schema_as_list_of_schema_fields(self): + if gcp_bigquery is not None: + schema_fields = [ + gcp_bigquery.SchemaField('float', 'FLOAT', mode='NULLABLE'), + gcp_bigquery.SchemaField('string', 'STRING', mode='NULLABLE'), + ] + else: + schema_fields = [ + bigquery_tools._TableFieldSchemaCompat('float', 'FLOAT', 'NULLABLE'), + bigquery_tools._TableFieldSchemaCompat( + 'string', 'STRING', 'NULLABLE'), + ] + coder = _JsonToDictCoder(schema_fields) + input_row = b'{"float": "10.5", "string": "abc"}' + expected_row = {'float': 10.5, 'string': 'abc'} + actual = coder.decode(input_row) + self.assertEqual(expected_row, actual) + @unittest.skipIf( HttpError is None or HttpForbiddenError is None, diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 194fe56fdd52..bcafa58259ca 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -3177,7 +3177,10 @@ def get_bq_tableschema(schema): if (isinstance(schema, (TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema - elif isinstance(schema, list): + elif isinstance(schema, (list, tuple)): + dict_schema = get_dict_table_schema(schema) + if isinstance(dict_schema, dict): + return parse_table_schema_from_json(json.dumps(dict_schema)) return TableSchema(fields=schema) elif isinstance(schema, str): return get_table_schema_from_string(schema) @@ -3185,6 +3188,9 @@ def get_bq_tableschema(schema): schema_string = json.dumps(schema) return parse_table_schema_from_json(schema_string) elif hasattr(schema, 'fields'): + dict_schema = get_dict_table_schema(schema) + if isinstance(dict_schema, dict): + return parse_table_schema_from_json(json.dumps(dict_schema)) return TableSchema(fields=list(schema.fields)) else: raise TypeError('Unexpected schema argument: %s.' % schema) From 668753c1828842c9f85db965b60144d6865e82e1 Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 27 Aug 2026 13:50:45 +0000 Subject: [PATCH 11/12] fix nested job configs --- .../apache_beam/io/gcp/bigquery_tools.py | 137 ++++++++++++++++-- .../apache_beam/io/gcp/bigquery_tools_test.py | 89 ++++++++++++ 2 files changed, 215 insertions(+), 11 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index bcafa58259ca..09d41e0901fa 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -541,7 +541,7 @@ def _message_compat_eq(self, other): _protorpclite_messages.Message.__eq__ = _message_compat_eq except ImportError: - pass + _protorpclite_messages = None def _set_table_ref_prop(ref, prop, val): @@ -597,6 +597,70 @@ def _set_table_ref_prop(ref, prop, val): lambda self: self.default_table_expiration_ms, lambda self, val: setattr(self, 'default_table_expiration_ms', val)) + if hasattr(gcp_bigquery, 'LoadJobConfig'): + if not hasattr(gcp_bigquery.LoadJobConfig, 'schemaUpdateOptions'): + gcp_bigquery.LoadJobConfig.schemaUpdateOptions = property( + lambda self: self.schema_update_options, + lambda self, val: setattr(self, 'schema_update_options', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'ignoreUnknownValues'): + gcp_bigquery.LoadJobConfig.ignoreUnknownValues = property( + lambda self: self.ignore_unknown_values, + lambda self, val: setattr(self, 'ignore_unknown_values', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'maxBadRecords'): + gcp_bigquery.LoadJobConfig.maxBadRecords = property( + lambda self: self.max_bad_records, + lambda self, val: setattr(self, 'max_bad_records', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'nullMarker'): + gcp_bigquery.LoadJobConfig.nullMarker = property( + lambda self: self.null_marker, + lambda self, val: setattr(self, 'null_marker', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'fieldDelimiter'): + gcp_bigquery.LoadJobConfig.fieldDelimiter = property( + lambda self: self.field_delimiter, + lambda self, val: setattr(self, 'field_delimiter', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'skipLeadingRows'): + gcp_bigquery.LoadJobConfig.skipLeadingRows = property( + lambda self: self.skip_leading_rows, + lambda self, val: setattr(self, 'skip_leading_rows', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'allowJaggedRows'): + gcp_bigquery.LoadJobConfig.allowJaggedRows = property( + lambda self: self.allow_jagged_rows, + lambda self, val: setattr(self, 'allow_jagged_rows', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'allowQuotedNewlines'): + gcp_bigquery.LoadJobConfig.allowQuotedNewlines = property( + lambda self: self.allow_quoted_newlines, + lambda self, val: setattr(self, 'allow_quoted_newlines', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'decimalTargetTypes'): + gcp_bigquery.LoadJobConfig.decimalTargetTypes = property( + lambda self: self.decimal_target_types, + lambda self, val: setattr(self, 'decimal_target_types', val)) + if not hasattr(gcp_bigquery.LoadJobConfig, 'useAvroLogicalTypes'): + gcp_bigquery.LoadJobConfig.useAvroLogicalTypes = property( + lambda self: self.use_avro_logical_types, + lambda self, val: setattr(self, 'use_avro_logical_types', val)) + + if hasattr(gcp_bigquery, 'QueryJobConfig'): + if not hasattr(gcp_bigquery.QueryJobConfig, 'schemaUpdateOptions'): + gcp_bigquery.QueryJobConfig.schemaUpdateOptions = property( + lambda self: self.schema_update_options, + lambda self, val: setattr(self, 'schema_update_options', val)) + if not hasattr(gcp_bigquery.QueryJobConfig, 'useLegacySql'): + gcp_bigquery.QueryJobConfig.useLegacySql = property( + lambda self: self.use_legacy_sql, + lambda self, val: setattr(self, 'use_legacy_sql', val)) + if not hasattr(gcp_bigquery.QueryJobConfig, 'flattenResults'): + gcp_bigquery.QueryJobConfig.flattenResults = property( + lambda self: self.flatten_results, + lambda self, val: setattr(self, 'flatten_results', val)) + if not hasattr(gcp_bigquery.QueryJobConfig, 'allowLargeResults'): + gcp_bigquery.QueryJobConfig.allowLargeResults = property( + lambda self: self.allow_large_results, + lambda self, val: setattr(self, 'allow_large_results', val)) + if not hasattr(gcp_bigquery.QueryJobConfig, 'maximumBytesBilled'): + gcp_bigquery.QueryJobConfig.maximumBytesBilled = property( + lambda self: self.maximum_bytes_billed, + lambda self, val: setattr(self, 'maximum_bytes_billed', val)) + if hasattr(gcp_job, '_AsyncJob') and not hasattr(gcp_job._AsyncJob, 'jobReference'): @@ -639,6 +703,29 @@ def totalBytesProcessed(self): gcp_job._AsyncJob.statistics = property(lambda self: _JobStatsCompat(self)) +def _to_json_compatible(obj): + """Converts an object or nested structure to JSON/API-compatible dicts/types.""" + if obj is None: + return None + if isinstance(obj, (str, int, float, bool)): + return obj + if isinstance(obj, (list, tuple, set)): + return [_to_json_compatible(item) for item in obj] + if isinstance(obj, dict): + return {k: _to_json_compatible(v) for k, v in obj.items()} + if hasattr(obj, 'to_api_repr') and callable(obj.to_api_repr): + return obj.to_api_repr() + if _protorpclite_messages is not None and hasattr( + _protorpclite_messages, 'Message') and isinstance( + obj, _protorpclite_messages.Message): + try: + from apitools.base.py import encoding + return encoding.MessageToDict(obj) + except Exception: + pass + return obj + + def _to_gcp_table_ref(table_ref, default_project=None): """Converts a TableReference or string into a google.cloud.bigquery.TableReference.""" if table_ref is None: @@ -1511,16 +1598,44 @@ def _insert_load_job( job_schema = None if schema == 'SCHEMA_AUTODETECT' else _to_gcp_schema( schema) autodetect = schema == 'SCHEMA_AUTODETECT' - additional_load_parameters = additional_load_parameters or {} - job_config = gcp_bigquery.LoadJobConfig( - schema=job_schema, - autodetect=autodetect, - create_disposition=create_disposition, - write_disposition=write_disposition, - source_format=source_format, - use_avro_logical_types=True, - labels=job_labels, - **additional_load_parameters) + clean_params = _to_json_compatible(additional_load_parameters or {}) + api_repr_dict = {} + direct_attrs = {} + if isinstance(clean_params, dict): + for k, v in clean_params.items(): + if '_' in k and hasattr(gcp_bigquery.LoadJobConfig, k): + direct_attrs[k] = v + else: + api_repr_dict[k] = v + + if api_repr_dict: + job_config = gcp_bigquery.LoadJobConfig.from_api_repr( + {'load': api_repr_dict}) + else: + job_config = gcp_bigquery.LoadJobConfig() + + for k, v in direct_attrs.items(): + try: + setattr(job_config, k, v) + except Exception: + try: + job_config._set_sub_prop(k, v) + except Exception: + pass + + if job_schema is not None: + job_config.schema = job_schema + if autodetect: + job_config.autodetect = True + if create_disposition is not None: + job_config.create_disposition = create_disposition + if write_disposition is not None: + job_config.write_disposition = write_disposition + if source_format is not None: + job_config.source_format = source_format + job_config.use_avro_logical_types = True + if job_labels is not None: + job_config.labels = job_labels try: if source_stream: job = self.client.load_table_from_file( diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index 6961332308bf..017184775b95 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -1582,6 +1582,95 @@ def test_to_gcp_dataset_ref_colon_format(self): self.assertEqual(ds_ref.datasetId, 'my_dataset') +class TestJobConfigCompatibility(unittest.TestCase): + def test_load_job_config_camel_case_properties(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + config = gcp_bigquery.LoadJobConfig( + schemaUpdateOptions=['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION'], + ignoreUnknownValues=True, + maxBadRecords=10, + nullMarker='\\N', + fieldDelimiter='\t', + skipLeadingRows=1, + allowJaggedRows=True, + allowQuotedNewlines=True, + decimalTargetTypes=['NUMERIC'], + useAvroLogicalTypes=True, + ) + self.assertEqual( + config.schema_update_options, + ['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION']) + self.assertEqual( + config.schemaUpdateOptions, + ['ALLOW_FIELD_ADDITION', 'ALLOW_FIELD_RELAXATION']) + self.assertTrue(config.ignore_unknown_values) + self.assertTrue(config.ignoreUnknownValues) + self.assertEqual(config.max_bad_records, 10) + self.assertEqual(config.maxBadRecords, 10) + + def test_query_job_config_camel_case_properties(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + config = gcp_bigquery.QueryJobConfig( + schemaUpdateOptions=['ALLOW_FIELD_ADDITION'], + useLegacySql=False, + flattenResults=True, + allowLargeResults=True, + maximumBytesBilled=1000000, + ) + self.assertEqual(config.schema_update_options, ['ALLOW_FIELD_ADDITION']) + self.assertEqual(config.schemaUpdateOptions, ['ALLOW_FIELD_ADDITION']) + self.assertFalse(config.use_legacy_sql) + self.assertFalse(config.useLegacySql) + self.assertTrue(config.flatten_results) + self.assertTrue(config.flattenResults) + self.assertEqual(config.maximum_bytes_billed, 1000000) + self.assertEqual(config.maximumBytesBilled, 1000000) + + def test_perform_load_job_with_modern_client_additional_parameters(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=gcp_bigquery.Client) + mock_job = mock.MagicMock() + mock_job.job_id = 'test_job_id' + mock_job.project = 'test-project' + mock_job.location = 'US' + client.load_table_from_uri.return_value = mock_job + + wrapper = bigquery_tools.BigQueryWrapper(client) + job_ref = wrapper.perform_load_job( + destination='test-project:test_dataset.test_table', + source_uris=['gs://test-bucket/data.csv'], + job_id='test_job_id', + additional_load_parameters={ + 'schemaUpdateOptions': ['ALLOW_FIELD_ADDITION'], + 'timePartitioning': { + 'type': 'DAY', 'field': 'date' + }, + 'ignoreUnknownValues': True, + }) + + self.assertEqual(job_ref.jobId, 'test_job_id') + self.assertEqual(job_ref.projectId, 'test-project') + client.load_table_from_uri.assert_called_once() + called_config = client.load_table_from_uri.call_args.kwargs['job_config'] + self.assertEqual( + called_config.schema_update_options, ['ALLOW_FIELD_ADDITION']) + self.assertEqual(called_config.time_partitioning.type_, 'DAY') + self.assertEqual(called_config.time_partitioning.field, 'date') + self.assertTrue(called_config.ignore_unknown_values) + + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) unittest.main() From ca73d78280cdbd60de41f67602339cedaebcae7c Mon Sep 17 00:00:00 2001 From: jrmccluskey Date: Thu, 27 Aug 2026 17:46:20 +0000 Subject: [PATCH 12/12] IT test fixes --- .../io/gcp/bigquery_change_history.py | 22 +- .../io/gcp/bigquery_change_history_it_test.py | 3 +- .../io/gcp/bigquery_change_history_test.py | 20 ++ .../apache_beam/io/gcp/bigquery_tools.py | 336 +++++++++++++++++- .../apache_beam/io/gcp/bigquery_tools_test.py | 228 ++++++++++++ .../io/gcp/tests/bigquery_matcher.py | 6 +- 6 files changed, 596 insertions(+), 19 deletions(-) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py index d6583c18de26..80489daa80a0 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history.py @@ -335,18 +335,26 @@ def build_changes_query( Returns: SQL string. """ - # Normalize 'project:dataset.table' to 'project.dataset.table' - table = table.replace(':', '.') + # Normalize 'project:dataset.table' or 'project.dataset.table' + table_ref = bigquery_tools.parse_table_reference(table) + if table_ref.projectId: + table = f"{table_ref.projectId}.{table_ref.datasetId}.{table_ref.tableId}" + else: + table = f"{table_ref.datasetId}.{table_ref.tableId}" start_iso = start.to_rfc3339() end_iso = end.to_rfc3339() - # Pseudo-columns (_CHANGE_TYPE, _CHANGE_TIMESTAMP) can't be written to - # destination tables with their original names. Rename them so they can + # Pseudo-columns (_CHANGE_TYPE, _CHANGE_TIMESTAMP, and for CHANGES() _CHANGE_IS_FOR_UPDATE) + # can't be written to destination tables with their original names. Rename them so they can # be persisted to the temp table for Storage Read API reading. pseudo = ( f"_CHANGE_TYPE AS {change_type_column}, " f"_CHANGE_TIMESTAMP AS {change_timestamp_column}") + except_cols = ( + "_CHANGE_TYPE, _CHANGE_TIMESTAMP, _CHANGE_IS_FOR_UPDATE" + if change_function.upper() == 'CHANGES' else + "_CHANGE_TYPE, _CHANGE_TIMESTAMP") if columns is None: - select = f"SELECT * EXCEPT(_CHANGE_TYPE, _CHANGE_TIMESTAMP), {pseudo}" + select = f"SELECT * EXCEPT({except_cols}), {pseudo}" else: select = f"SELECT {', '.join(columns)}, {pseudo}" from_clause = ( @@ -527,7 +535,7 @@ def _emit_query_ranges( start_ts: Timestamp, end_ts: Timestamp, watermark_estimator: _PollWatermarkEstimator - ) -> Iterable[TimestampedValue[_QueryRange]]: + ) -> Iterable[Any]: """Compute and yield _QueryRange elements, advancing estimator state.""" ranges = compute_ranges(start_ts, end_ts, self._change_function) _LOGGER.info( @@ -555,7 +563,7 @@ def process( restriction_tracker=beam.DoFn.RestrictionParam(), watermark_estimator=beam.DoFn.WatermarkEstimatorParam( _PollWatermarkEstimatorProvider()) - ) -> Iterable[TimestampedValue[_QueryRange]]: + ) -> Iterable[Any]: now = time.time() start_ts = watermark_estimator.poll_cursor() diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py index ef41fc393af7..cc9a8abb87f0 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_it_test.py @@ -307,7 +307,8 @@ def test_cleanup_signal_emitted(self): # Verify cleanup signal cleanup_table_keys = ( outputs['cleanup'] - | 'ExtractKey' >> beam.Map(lambda x: x[0])) + | 'ExtractKey' >> beam.Map(lambda x: x[0]) + | 'Distinct' >> beam.Distinct()) assert_that( cleanup_table_keys, equal_to([bigquery_tools.get_hashable_destination(table_ref)]), diff --git a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py index 11b467f26d49..f68e63b66632 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_change_history_test.py @@ -89,6 +89,26 @@ def test_row_filter(self): 'proj.ds.tbl', start, end, 'APPENDS', row_filter='status = "active"') self.assertIn('WHERE status = "active"', sql) + def test_unqualified_table_does_not_emit_none_prefix(self): + start = _ts(2025, 1, 1) + end = _ts(2025, 1, 2) + sql = build_changes_query('ds.tbl', start, end, 'APPENDS') + self.assertIn('TABLE `ds.tbl`', sql) + self.assertNotIn('None.', sql) + + def test_domain_scoped_project_table(self): + start = _ts(2025, 1, 1) + end = _ts(2025, 1, 2) + sql = build_changes_query( + 'google.com:clouddfe:ds.tbl', start, end, 'APPENDS') + self.assertIn('TABLE `google.com:clouddfe.ds.tbl`', sql) + + def test_changes_query_excepts_is_for_update(self): + start = _ts(2025, 1, 1) + end = _ts(2025, 1, 2) + sql = build_changes_query('proj.ds.tbl', start, end, 'CHANGES') + self.assertIn('_CHANGE_IS_FOR_UPDATE', sql) + def test_no_row_filter(self): start = _ts(2025, 1, 1) end = _ts(2025, 1, 2) diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools.py b/sdks/python/apache_beam/io/gcp/bigquery_tools.py index 09d41e0901fa..9a8c96cf1312 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -590,6 +590,26 @@ def _set_table_ref_prop(ref, prop, val): gcp_bigquery.Table.rangePartitioning = property( lambda self: self.range_partitioning) + if hasattr(gcp_bigquery, 'TimePartitioning'): + if not hasattr(gcp_bigquery.TimePartitioning, 'type'): + gcp_bigquery.TimePartitioning.type = property( + lambda self: self.type_, + lambda self, val: setattr(self, 'type_', val)) + if not hasattr(gcp_bigquery.TimePartitioning, 'expirationMs'): + gcp_bigquery.TimePartitioning.expirationMs = property( + lambda self: self.expiration_ms, + lambda self, val: setattr(self, 'expiration_ms', val)) + if not hasattr(gcp_bigquery.TimePartitioning, 'requirePartitionFilter'): + gcp_bigquery.TimePartitioning.requirePartitionFilter = property( + lambda self: self.require_partition_filter, + lambda self, val: setattr(self, 'require_partition_filter', val)) + + if hasattr(gcp_bigquery, 'RangePartitioning'): + if not hasattr(gcp_bigquery.RangePartitioning, 'range'): + gcp_bigquery.RangePartitioning.range = property( + lambda self: self.range_, + lambda self, val: setattr(self, 'range_', val)) + if not hasattr(gcp_bigquery.Dataset, 'datasetReference'): gcp_bigquery.Dataset.datasetReference = property( lambda self: self.reference) @@ -661,6 +681,60 @@ def _set_table_ref_prop(ref, prop, val): lambda self: self.maximum_bytes_billed, lambda self, val: setattr(self, 'maximum_bytes_billed', val)) + if hasattr(gcp_bigquery, 'Table') and hasattr(gcp_bigquery.Table, 'labels'): + _orig_tbl_labels_setter = gcp_bigquery.Table.labels.fset + if _orig_tbl_labels_setter: + + def _safe_tbl_labels_setter(self, value): + if value is None: + value = {} + elif not isinstance(value, dict) and hasattr(value, + 'additionalProperties'): + from apitools.base.py import encoding + value = encoding.MessageToDict(value) + _orig_tbl_labels_setter(self, value) + + gcp_bigquery.Table.labels = gcp_bigquery.Table.labels.setter( + _safe_tbl_labels_setter) + + if hasattr(gcp_bigquery, 'Dataset') and hasattr(gcp_bigquery.Dataset, + 'labels'): + _orig_ds_labels_setter = gcp_bigquery.Dataset.labels.fset + if _orig_ds_labels_setter: + + def _safe_ds_labels_setter(self, value): + if value is None: + value = {} + elif not isinstance(value, dict) and hasattr(value, + 'additionalProperties'): + from apitools.base.py import encoding + value = encoding.MessageToDict(value) + _orig_ds_labels_setter(self, value) + + gcp_bigquery.Dataset.labels = gcp_bigquery.Dataset.labels.setter( + _safe_ds_labels_setter) + + try: + from google.cloud.bigquery.job.base import _JobConfig as _GcpJobConfig + if hasattr(_GcpJobConfig, 'labels') and hasattr(_GcpJobConfig.labels, + 'fset'): + _orig_job_labels_setter = _GcpJobConfig.labels.fset + if _orig_job_labels_setter: + + def _safe_job_labels_setter(self, value): + if value is None: + value = {} + elif not isinstance(value, dict) and hasattr(value, + 'additionalProperties'): + from apitools.base.py import encoding + value = encoding.MessageToDict(value) + _orig_job_labels_setter(self, value) + + _GcpJobConfig.labels = _GcpJobConfig.labels.setter( + _safe_job_labels_setter) + except ImportError: + pass + if hasattr(gcp_job, '_AsyncJob') and not hasattr(gcp_job._AsyncJob, 'jobReference'): @@ -696,6 +770,18 @@ def totalBytesBilled(self): def totalBytesProcessed(self): return getattr(self._job, 'total_bytes_processed', None) + @property + def referencedTables(self): + tables = getattr(self._job, 'referenced_tables', None) + if tables is not None: + return [ + TableReference( + projectId=t.project, + datasetId=t.dataset_id, + tableId=t.table_id) for t in tables + ] + return None + gcp_job._AsyncJob.jobReference = property( lambda self: JobReference( job_id=self.job_id, project=self.project, location=self.location)) @@ -726,6 +812,16 @@ def _to_json_compatible(obj): return obj +def _extract_dict_labels(labels): + """Converts labels to a non-empty dictionary or returns None.""" + if not labels: + return None + labels = _to_json_compatible(labels) + if isinstance(labels, dict) and labels: + return labels + return None + + def _to_gcp_table_ref(table_ref, default_project=None): """Converts a TableReference or string into a google.cloud.bigquery.TableReference.""" if table_ref is None: @@ -830,6 +926,59 @@ def _to_gcp_schema(schema): return schema +def _to_table_schema(schema): + """Converts a list of google.cloud.bigquery.SchemaField, dict, or TableSchema into a TableSchema.""" + if schema is None: + return TableSchema() + if isinstance(schema, TableSchema): + return schema + if isinstance(schema, dict): + return _to_table_schema(schema.get('fields', [])) + if hasattr(schema, 'fields') and not isinstance(schema, (list, tuple)): + return _to_table_schema(schema.fields) + + def _to_field_schema(f): + if isinstance(f, TableFieldSchema): + return f + if isinstance(f, dict): + f_dict = f + elif hasattr(f, 'to_api_repr'): + f_dict = f.to_api_repr() + else: + f_dict = None + + if f_dict is not None: + name = f_dict.get('name', '') + field_type = f_dict.get('type') or f_dict.get('type_') or 'STRING' + mode = f_dict.get('mode', 'NULLABLE') + description = f_dict.get('description', None) + sub_fields = [_to_field_schema(sf) for sf in f_dict.get('fields', [])] + return TableFieldSchema( + name=name, + type=field_type, + mode=mode, + description=description, + fields=sub_fields) + + name = getattr(f, 'name', '') + field_type = getattr(f, 'field_type', None) or getattr(f, 'type', + None) or 'STRING' + mode = getattr(f, 'mode', 'NULLABLE') + description = getattr(f, 'description', None) + sub = getattr(f, 'fields', ()) + sub_fields = [_to_field_schema(sf) for sf in sub] if sub else () + return TableFieldSchema( + name=name, + type=field_type, + mode=mode, + description=description, + fields=sub_fields) + + if isinstance(schema, (list, tuple)): + return TableSchema(fields=[_to_field_schema(f) for f in schema]) + return TableSchema() + + if gcp_bigquery: class _ClientTablesCompat: @@ -875,6 +1024,65 @@ def Insert(self, request): ds_id), tbl_id) gcp_table = gcp_bigquery.Table(gcp_tbl_ref, schema=_to_gcp_schema(schema)) + if table is not None: + tp = getattr(table, 'timePartitioning', None) or getattr( + table, 'time_partitioning', None) + if tp is not None: + if isinstance(tp, gcp_bigquery.TimePartitioning): + gcp_table.time_partitioning = tp + else: + tp_field = getattr(tp, 'field', None) + tp_type = getattr(tp, 'type', None) or getattr(tp, 'type_', None) + tp_exp = getattr(tp, 'expirationMs', None) or getattr( + tp, 'expiration_ms', None) + tp_req = getattr(tp, 'requirePartitionFilter', None) or getattr( + tp, 'require_partition_filter', None) + gcp_table.time_partitioning = gcp_bigquery.TimePartitioning( + type_=tp_type, + field=tp_field, + expiration_ms=tp_exp, + require_partition_filter=tp_req) + rp = getattr(table, 'rangePartitioning', None) or getattr( + table, 'range_partitioning', None) + if rp is not None: + if isinstance(rp, gcp_bigquery.RangePartitioning): + gcp_table.range_partitioning = rp + else: + rp_field = getattr(rp, 'field', None) + rp_range = getattr(rp, 'range', None) or getattr(rp, 'range_', None) + if rp_range is not None and hasattr(gcp_bigquery, 'PartitionRange'): + start = getattr(rp_range, 'start', None) + end = getattr(rp_range, 'end', None) + interval = getattr(rp_range, 'interval', None) + rp_range = gcp_bigquery.PartitionRange( + start=start, end=end, interval=interval) + gcp_table.range_partitioning = gcp_bigquery.RangePartitioning( + field=rp_field, range_=rp_range) + clustering = getattr(table, 'clustering', None) + if clustering is not None: + fields = getattr(clustering, 'fields', clustering) + if isinstance(fields, (list, tuple)): + gcp_table.clustering_fields = list(fields) + if getattr(table, 'description', None): + gcp_table.description = table.description + if getattr(table, 'friendlyName', None) or getattr( + table, 'friendly_name', None): + gcp_table.friendly_name = getattr( + table, 'friendlyName', None) or getattr( + table, 'friendly_name', None) + dict_labels = _extract_dict_labels(getattr(table, 'labels', None)) + if dict_labels: + gcp_table.labels = dict_labels + kms = getattr( + getattr(table, 'encryptionConfiguration', None), + 'kmsKeyName', + None) or getattr( + getattr(table, 'encryption_configuration', None), + 'kms_key_name', + None) + if kms: + gcp_table.encryption_configuration = ( + gcp_bigquery.EncryptionConfiguration(kms_key_name=kms)) return self._client.create_table(gcp_table, exists_ok=True) def Delete(self, request): @@ -992,8 +1200,82 @@ def GetQueryResults(self, request): job_id = getattr(request, 'jobId', None) loc = getattr(request, 'location', None) page_token = getattr(request, 'pageToken', None) + max_results = getattr(request, 'maxResults', None) job = self._client.get_job(job_id, project=proj, location=loc) - return job.result(page_token=page_token) + if page_token is not None: + return self._client.list_rows( + job, page_token=page_token, max_results=max_results) + return job.result(max_results=max_results) + + def Insert(self, request, upload=None): + job_obj = getattr(request, 'job', None) + job_ref = ( + getattr(job_obj, 'jobReference', None) + if job_obj else getattr(request, 'jobReference', None)) + job_id = getattr(job_ref, 'jobId', None) or getattr( + job_ref, 'job_id', None) + proj = ( + getattr(request, 'projectId', None) or + getattr(job_ref, 'projectId', None) or + getattr(job_ref, 'project', None)) + config = getattr(job_obj, 'configuration', None) if job_obj else None + if config and getattr(config, 'query', None): + q = config.query + dest = None + if getattr(q, 'destinationTable', None): + dest = _to_gcp_table_ref(q.destinationTable, default_project=proj) + dict_labels = _extract_dict_labels(getattr(config, 'labels', None)) + job_config = gcp_bigquery.QueryJobConfig( + dry_run=getattr(q, 'dryRun', False), + use_legacy_sql=getattr(q, 'useLegacySql', False) + if getattr(q, 'useLegacySql', None) is not None else False, + flatten_results=getattr(q, 'flattenResults', None), + priority=getattr(q, 'priority', 'INTERACTIVE'), + destination=dest, + ) + if dict_labels: + job_config.labels = dict_labels + kms = getattr( + getattr(q, 'destinationEncryptionConfiguration', None), + 'kmsKeyName', + None) + if kms: + job_config.destination_encryption_configuration = ( + gcp_bigquery.EncryptionConfiguration(kms_key_name=kms)) + return self._client.query( + q.query, + job_config=job_config, + job_id=job_id, + project=proj, + job_retry=None, + ) + elif config and getattr(config, 'load', None): + ld = config.load + dest = _to_gcp_table_ref( + getattr(ld, 'destinationTable', None), default_project=proj) + uris = list(getattr(ld, 'sourceUris', [])) + if uris: + return self._client.load_table_from_uri( + uris, dest, job_id=job_id, project=proj) + elif config and getattr(config, 'copy', None): + cp = config.copy + sources = [ + _to_gcp_table_ref(s, default_project=proj) + for s in getattr(cp, 'sourceTables', []) + ] + dest = _to_gcp_table_ref( + getattr(cp, 'destinationTable', None), default_project=proj) + return self._client.copy_table( + sources, dest, job_id=job_id, project=proj) + elif config and getattr(config, 'extract', None): + ex = config.extract + src = _to_gcp_table_ref( + getattr(ex, 'sourceTable', None), default_project=proj) + uris = list(getattr(ex, 'destinationUris', [])) + return self._client.extract_table( + src, uris, job_id=job_id, project=proj) + + return self._client.get_job(job_id, project=proj) if not hasattr(gcp_bigquery.Client, 'tables'): gcp_bigquery.Client.tables = property( @@ -1515,11 +1797,13 @@ def _insert_copy_job( write_disposition=None, job_labels=None): if self._is_modern_client: + dict_labels = _extract_dict_labels(job_labels) job_config = gcp_bigquery.CopyJobConfig( create_disposition=create_disposition, write_disposition=write_disposition, - labels=job_labels, ) + if dict_labels: + job_config.labels = dict_labels src = from_table_reference if isinstance( from_table_reference, list) else [from_table_reference] src_refs = [_to_gcp_table_ref(t, default_project=project_id) for t in src] @@ -1634,8 +1918,9 @@ def _insert_load_job( if source_format is not None: job_config.source_format = source_format job_config.use_avro_logical_types = True - if job_labels is not None: - job_config.labels = job_labels + dict_labels = _extract_dict_labels(job_labels) + if dict_labels: + job_config.labels = dict_labels try: if source_stream: job = self.client.load_table_from_file( @@ -1780,6 +2065,7 @@ def _start_query_job( self._get_temp_table(self._get_temp_table_project(project_id)), default_project=project_id or getattr(self.client, 'project', None)) + dict_labels = _extract_dict_labels(job_labels) job_config = gcp_bigquery.QueryJobConfig( dry_run=dry_run, use_legacy_sql=use_legacy_sql, @@ -1787,8 +2073,9 @@ def _start_query_job( destination=dest_table, flatten_results=flatten_results, priority=priority, - labels=job_labels, ) + if dict_labels: + job_config.labels = dict_labels if kms_key: job_config.destination_encryption_configuration = ( gcp_bigquery.EncryptionConfiguration(kms_key_name=kms_key)) @@ -1907,7 +2194,10 @@ def _get_query_results( return self.client.jobs.GetQueryResults(request) job = self.client.get_job(job_id, project=project_id, location=location) - return job.result(max_results=max_results, page_token=page_token) + if page_token is not None: + return self.client.list_rows( + job, page_token=page_token, max_results=max_results) + return job.result(max_results=max_results) @retry.with_exponential_backoff( num_retries=MAX_RETRIES, @@ -2150,8 +2440,9 @@ def get_or_create_dataset( dataset = gcp_bigquery.Dataset(dataset_ref) if location is not None: dataset.location = location - if labels is not None: - dataset.labels = labels + dict_labels = _extract_dict_labels(labels) + if dict_labels: + dataset.labels = dict_labels if kms_key is not None: dataset.default_encryption_configuration = ( gcp_bigquery.EncryptionConfiguration(kms_key_name=kms_key)) @@ -2537,13 +2828,15 @@ def perform_extract_job( src_ref = _to_gcp_table_ref(table_reference, default_project=job_project) dest_uris = destination if isinstance(destination, list) else [destination] + dict_labels = _extract_dict_labels(job_labels) job_config = gcp_bigquery.ExtractJobConfig( destination_format=destination_format, print_header=include_header, compression=compression, use_avro_logical_types=use_avro_logical_types, - labels=job_labels, ) + if dict_labels: + job_config.labels = dict_labels try: job = self.client.extract_table( src_ref, @@ -2705,6 +2998,28 @@ def run_query( priority, dry_run=False, job_labels=None): + if self._is_modern_client: + job_config = gcp_bigquery.QueryJobConfig( + dry_run=dry_run, + use_legacy_sql=use_legacy_sql, + flatten_results=flatten_results, + priority=priority, + ) + if job_labels: + dict_labels = _extract_dict_labels(job_labels) + if dict_labels: + job_config.labels = dict_labels + job = self.client.query( + query, + job_config=job_config, + project=project_id, + ) + if dry_run: + return + rows = job.result() + yield list(rows), _to_table_schema(rows.schema) + return + job = self._start_query_job( project_id, query, @@ -2745,9 +3060,10 @@ def run_query( page_token = response.pageToken else: # Modern RowIterator - yield list(response), TableSchema(fields=response.schema) + yield list(response), _to_table_schema(response.schema) break + def insert_rows( self, project_id, diff --git a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py index 017184775b95..058ff18dc1ee 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -1670,6 +1670,234 @@ def test_perform_load_job_with_modern_client_additional_parameters(self): self.assertEqual(called_config.time_partitioning.field, 'date') self.assertTrue(called_config.ignore_unknown_values) + def test_job_configs_with_none_labels(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + # Instantiating configs with labels=None or empty dict must not fail + q_cfg = gcp_bigquery.QueryJobConfig(labels=None) + self.assertEqual(q_cfg.labels, {}) + l_cfg = gcp_bigquery.LoadJobConfig(labels=None) + self.assertEqual(l_cfg.labels, {}) + c_cfg = gcp_bigquery.CopyJobConfig(labels=None) + self.assertEqual(c_cfg.labels, {}) + e_cfg = gcp_bigquery.ExtractJobConfig(labels=None) + self.assertEqual(e_cfg.labels, {}) + + # Setting labels = None on Dataset and Table must not fail + ds = gcp_bigquery.Dataset('test-project.test_dataset') + ds.labels = None + self.assertEqual(ds.labels, {}) + tbl = gcp_bigquery.Table('test-project.test_dataset.test_table') + tbl.labels = None + self.assertEqual(tbl.labels, {}) + + def test_start_query_job_with_none_labels(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=gcp_bigquery.Client) + mock_job = mock.MagicMock() + mock_job.job_id = 'query_job_id' + client.query.return_value = mock_job + + wrapper = bigquery_tools.BigQueryWrapper(client) + job = wrapper._start_query_job( + project_id='test-project', + query='SELECT 1', + use_legacy_sql=False, + flatten_results=False, + job_id='query_job_id', + priority='BATCH', + dry_run=False, + job_labels=None, + ) + self.assertEqual(job, mock_job) + client.query.assert_called_once() + called_config = client.query.call_args.kwargs['job_config'] + self.assertEqual(called_config.labels, {}) + + def test_insert_copy_job_with_none_labels(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=gcp_bigquery.Client) + mock_job = mock.MagicMock() + mock_job.job_id = 'copy_job_id' + mock_job.project = 'test-project' + mock_job.location = 'US' + client.copy_table.return_value = mock_job + + wrapper = bigquery_tools.BigQueryWrapper(client) + job_ref = wrapper._insert_copy_job( + project_id='test-project', + job_id='copy_job_id', + from_table_reference='test-project:test_dataset.src_table', + to_table_reference='test-project:test_dataset.dst_table', + job_labels=None, + ) + self.assertEqual(job_ref.jobId, 'copy_job_id') + client.copy_table.assert_called_once() + called_config = client.copy_table.call_args.kwargs['job_config'] + self.assertEqual(called_config.labels, {}) + + def test_perform_extract_job_with_none_labels(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=gcp_bigquery.Client) + mock_job = mock.MagicMock() + mock_job.job_id = 'extract_job_id' + mock_job.project = 'test-project' + mock_job.location = 'US' + client.extract_table.return_value = mock_job + + wrapper = bigquery_tools.BigQueryWrapper(client) + job_ref = wrapper.perform_extract_job( + destination='gs://test-bucket/output.csv', + job_id='extract_job_id', + table_reference='test-project:test_dataset.src_table', + destination_format='CSV', + job_labels=None, + ) + self.assertEqual(job_ref.jobId, 'extract_job_id') + client.extract_table.assert_called_once() + called_config = client.extract_table.call_args.kwargs['job_config'] + self.assertEqual(called_config.labels, {}) + + def test_to_table_schema_nested_records(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + schema = [ + gcp_bigquery.SchemaField( + 'person', + 'RECORD', + mode='NULLABLE', + fields=[ + gcp_bigquery.SchemaField('name', 'STRING', mode='REQUIRED'), + gcp_bigquery.SchemaField('age', 'INTEGER', mode='NULLABLE'), + gcp_bigquery.SchemaField( + 'address', + 'RECORD', + fields=[ + gcp_bigquery.SchemaField( + 'city', 'STRING', mode='NULLABLE'), + ]), + ]), + ] + table_schema = bigquery_tools._to_table_schema(schema) + self.assertEqual(len(table_schema.fields), 1) + person = table_schema.fields[0] + self.assertEqual(person.name, 'person') + self.assertEqual(person.type, 'RECORD') + self.assertEqual(len(person.fields), 3) + self.assertEqual(person.fields[0].name, 'name') + self.assertEqual(person.fields[0].type, 'STRING') + self.assertEqual(person.fields[1].name, 'age') + self.assertEqual(person.fields[1].type, 'INTEGER') + self.assertEqual(person.fields[2].name, 'address') + self.assertEqual(person.fields[2].type, 'RECORD') + self.assertEqual(len(person.fields[2].fields), 1) + self.assertEqual(person.fields[2].fields[0].name, 'city') + self.assertEqual(person.fields[2].fields[0].type, 'STRING') + + def test_to_table_schema_dict(self): + dict_schema = { + 'fields': [ + {'name': 'id', 'type': 'INTEGER', 'mode': 'REQUIRED'}, + {'name': 'val', 'type': 'STRING', 'mode': 'NULLABLE'}, + ] + } + table_schema = bigquery_tools._to_table_schema(dict_schema) + self.assertEqual(len(table_schema.fields), 2) + self.assertEqual(table_schema.fields[0].name, 'id') + self.assertEqual(table_schema.fields[0].type, 'INTEGER') + self.assertEqual(table_schema.fields[1].name, 'val') + self.assertEqual(table_schema.fields[1].type, 'STRING') + + def test_job_stats_referenced_tables(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + mock_job = mock.MagicMock(spec=gcp_bigquery.QueryJob) + mock_job.job_id = 'query_job_id' + mock_job.project = 'test-project' + mock_job.location = 'US' + t1 = gcp_bigquery.TableReference.from_string('test-project.dataset.table1') + t2 = gcp_bigquery.TableReference.from_string('test-project.dataset.table2') + mock_job.referenced_tables = [t1, t2] + + stats = bigquery_tools._JobStatsCompat(mock_job) + ref_tables = stats.referencedTables + self.assertEqual(len(ref_tables), 2) + self.assertEqual(ref_tables[0].projectId, 'test-project') + self.assertEqual(ref_tables[0].datasetId, 'dataset') + self.assertEqual(ref_tables[0].tableId, 'table1') + self.assertEqual(ref_tables[1].projectId, 'test-project') + self.assertEqual(ref_tables[1].datasetId, 'dataset') + self.assertEqual(ref_tables[1].tableId, 'table2') + + def test_client_tables_compat_insert_labels_and_metadata(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=gcp_bigquery.Client) + client.project = 'test-project' + created_table_mock = mock.MagicMock(spec=gcp_bigquery.Table) + client.create_table.return_value = created_table_mock + + tables_compat = bigquery_tools._ClientTablesCompat(client) + request = mock.MagicMock() + request.projectId = 'test-project' + request.datasetId = 'test_dataset' + request.table = mock.MagicMock() + request.table.tableReference = mock.MagicMock() + request.table.tableReference.projectId = 'test-project' + request.table.tableReference.datasetId = 'test_dataset' + request.table.tableReference.tableId = 'test_table' + request.table.schema = None + request.table.labels = {'env': 'test', 'tier': 'frontend'} + request.table.friendlyName = 'My Test Table' + request.table.description = 'A test table description' + request.table.timePartitioning = None + request.table.rangePartitioning = None + request.table.clustering = None + request.table.encryptionConfiguration = None + + tables_compat.Insert(request) + client.create_table.assert_called_once() + passed_table = client.create_table.call_args.args[0] + self.assertEqual(passed_table.labels, {'env': 'test', 'tier': 'frontend'}) + self.assertEqual(passed_table.friendly_name, 'My Test Table') + self.assertEqual(passed_table.description, 'A test table description') + + def test_labels_setter_clears_on_none(self): + try: + from google.cloud import bigquery as gcp_bigquery + except ImportError: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + table = gcp_bigquery.Table('test-project.dataset.table') + table.labels = {'initial': 'label'} + self.assertEqual(table.labels, {'initial': 'label'}) + table.labels = None + self.assertEqual(table.labels, {}) + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) diff --git a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py index ff48a5644916..7d02768e16b3 100644 --- a/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py +++ b/sdks/python/apache_beam/io/gcp/tests/bigquery_matcher.py @@ -239,9 +239,13 @@ def _get_or_none(obj, attr): try: return obj.__getattribute__(attr) except AttributeError: + if attr == 'type' and hasattr(obj, 'type_'): + return obj.type_ + if attr == 'range' and hasattr(obj, 'range_'): + return obj.range_ try: return obj.get(attr, None) - except TypeError: + except (TypeError, AttributeError): return None @staticmethod