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..04eafa2cf53a 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__) @@ -524,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: @@ -767,14 +776,42 @@ 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: + 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 + + 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): + 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): """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 +831,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 +3124,25 @@ 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): + 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) + 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..80489daa80a0 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 @@ -326,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 = ( @@ -448,14 +465,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) @@ -497,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( @@ -525,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() @@ -627,27 +665,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_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_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..15449e6da780 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() @@ -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_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..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, @@ -448,8 +486,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 +556,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 +582,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 +678,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 +702,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 +743,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 +758,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 +772,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 +796,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 +2260,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..9a8c96cf1312 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools.py @@ -36,6 +36,7 @@ import re import sys import time +import typing import uuid from json.decoder import JSONDecodeError from typing import Optional @@ -49,13 +50,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 +77,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 +116,1175 @@ 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. +# +# 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. + + +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) or _DatasetReferenceCompat + TableFieldSchema = bigquery.TableFieldSchema + TableSchema = bigquery.TableSchema + 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) + 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: + TableReference = _TableReferenceCompat + DatasetReference = _DatasetReferenceCompat + TableFieldSchema = _TableFieldSchemaCompat + TableSchema = _TableSchemaCompat + TableRow = _TableRowCompat + TableCell = _TableCellCompat + Table = None + Dataset = None + Job = None + JobConfiguration = None + JobConfigurationLoad = None + JobConfigurationQuery = None + JobConfigurationExtract = None + JobConfigurationTableCopy = None + JobStatistics = None + JobStatistics2 = 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})") + -# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports try: - from apache_beam.io.gcp.internal.clients.bigquery import TableReference + 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) + 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 except ImportError: - TableReference = None -# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports + _protorpclite_messages = None + + +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: _set_table_ref_prop(self, 'projectId', val)) + gcp_bigquery.TableReference.datasetId = property( + lambda self: self.dataset_id, + lambda self, val: _set_table_ref_prop(self, 'datasetId', val)) + gcp_bigquery.TableReference.tableId = property( + lambda self: self.table_id, + lambda self, val: _set_table_ref_prop(self, 'tableId', 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 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) + 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_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_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'): + + 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) + + @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)) + gcp_job._AsyncJob.status = property(lambda self: _JobStatusCompat(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 _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: + 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 _TableReferenceCompat( + projectId=proj, datasetId=dataset_id, tableId=table_id) + return table_ref + + +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( + 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.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 _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, ds_id) + return _DatasetReferenceCompat(projectId=proj, datasetId=ds_id) + return dataset_ref + + +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)): + 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 + + +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: + 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)) + 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): + 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) + max_results = getattr(request, 'maxResults', None) + job = self._client.get_job(job_id, project=proj, location=loc) + 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( + 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__) @@ -153,6 +1343,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): @@ -225,21 +1417,21 @@ def _parse_schema_field(field): 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): @@ -269,18 +1461,21 @@ def parse_table_reference(table, dataset=None, project=None): ValueError: if the table reference as a string does not match the expected format. """ - 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. @@ -293,14 +1488,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 +1501,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 +1539,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): @@ -365,11 +1564,19 @@ class BigQueryWrapper(object): HISTOGRAM_METRIC_LOGGER = MetricLogger() - def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): - self.client = client or BigQueryWrapper._bigquery_client(PipelineOptions()) - self.gcp_bq_client = client or gcp_bigquery.Client( - client_info=ClientInfo( - user_agent="apache-beam-%s" % apache_beam.__version__)) + def __init__( + self, + client=None, + 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(), use_legacy_client=use_legacy_client) + 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 @@ -401,6 +1608,31 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): self.created_temp_dataset = False + @property + 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) + @property def unique_row_id(self): """Returns a unique row ID (str) used to avoid multiple insertions. @@ -445,43 +1677,98 @@ 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. """ - reference = bigquery.JobReference( - jobId=uuid.uuid4().hex, projectId=project_id) - request = bigquery.BigqueryJobsInsertRequest( + 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): + # Permission access for table (i.e. from authorized_view), + # try next one + return None + except Exception: + raise + + 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: + # This behavior is only expected in tests + _LOGGER.warning( + "Unable to get location, missing response.statistics. Query: %s", + query) + return None + + if referenced_tables: # Guards against both non-empty and non-None + 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): + # Permission access for table (i.e. from authorized_view), + # try next one + 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 = ( + 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=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 for table in referenced_tables: try: location = self.get_table_location( table.projectId, table.datasetId, table.tableId) - except HttpForbiddenError: + except (HttpForbiddenError, Forbidden): # Permission access for table (i.e. from authorized_view), # try next one continue @@ -491,7 +1778,6 @@ def get_query_location(self, project_id, query, use_legacy_sql): table, query) return location - _LOGGER.debug( "Query %s does not reference any tables or " "you don't have permission to inspect them.", @@ -510,14 +1796,49 @@ 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: + dict_labels = _extract_dict_labels(job_labels) + job_config = gcp_bigquery.CopyJobConfig( + create_disposition=create_disposition, + write_disposition=write_disposition, + ) + 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] + 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, will not retry inserting it: %s", + job_id, + exn) + return JobReference(job_id=job_id, project=project_id) + raise + + # Fallback if legacy client.jobs.Insert is mocked + 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=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 +1848,6 @@ def _insert_copy_job( ), jobReference=reference, )) - return self._start_job(request).jobReference @retry.with_exponential_backoff( @@ -557,17 +1877,100 @@ def _insert_load_job( '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' + 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 + 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( + 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, will not retry inserting it: %s", + job_id, + exn) + 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 = ( + 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=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, @@ -598,10 +2001,10 @@ 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. @@ -611,7 +2014,7 @@ def _start_job( """ 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( @@ -652,35 +2055,80 @@ 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)) + + dict_labels = _extract_dict_labels(job_labels) + 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, + ) + if dict_labels: + job_config.labels = dict_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 = ( + 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=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. + 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. @@ -690,16 +2138,35 @@ def wait_for_bq_job(self, job_reference, sleep_duration_sec=5, max_retries=0): reached. """ 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 +2183,21 @@ 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) + 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, @@ -741,9 +2215,7 @@ def _insert_all_rows( 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. + # The rows argument is a list of plain Python dictionaries or rows. resource = resource_identifiers.BigQueryTable( project_id, dataset_id, table_id) @@ -763,28 +2235,97 @@ 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: + # 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 + 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__ + # Add all rows to the errors list along with the error + 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 + # 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( + 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. @@ -801,17 +2342,24 @@ 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 + Table instance (bigquery.Table or google.cloud.bigquery.Table). Raises: - HttpError: if lookup failed. + NotFound or HttpError: if lookup failed. """ - request = bigquery.BigqueryTablesGetRequest( + if self._is_modern_client: + 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 + request = apitools_bigquery.BigqueryTablesGetRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id) response = self.client.tables.Get(request) return response @@ -831,13 +2379,29 @@ def _create_table( 'See https://cloud.google.com/bigquery/docs/tables#table_naming' % table_id) + if self._is_modern_client: + 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: + 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) + # The response is a Table instance. + 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) @@ -856,9 +2420,44 @@ def get_or_create_dataset( kms_key=None, default_table_expiration_ms=None): # Check if dataset already exists otherwise create it + if self._is_modern_client: + 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 + 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 + 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)) + 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 + # The response is a Dataset instance. + 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 +2469,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,7 +2481,7 @@ 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 @@ -895,7 +2494,18 @@ 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 = _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 + return len(list(rows)) == 0 + + # Fallback for legacy client + request = apitools_bigquery.BigqueryTabledataListRequest( projectId=project_id, datasetId=dataset_id, tableId=table_id, @@ -908,7 +2518,24 @@ def _is_table_empty(self, project_id, dataset_id, table_id): 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 = _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: + 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 +2551,24 @@ 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 = _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) + 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) @@ -976,9 +2620,44 @@ 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 = _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: + 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: + # 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: + 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: @@ -1013,31 +2692,51 @@ 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 %s for clean up.', + project_id) + 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 +2751,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 +2777,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 or bigquery.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 +2819,50 @@ 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 or bigquery.JobReference with the information about the job that was started. """ - job_project = project or table_reference.projectId - job_reference = bigquery.JobReference(jobId=job_id, projectId=job_project) - request = bigquery.BigqueryJobsInsertRequest( + 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] + 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, + ) + if dict_labels: + job_config.labels = dict_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 = ( + 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=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, @@ -1155,12 +2898,12 @@ def get_or_create_table( 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. + 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 bigquery.Table instance if table was found or created. + A Table instance if table was found or created. Raises: `RuntimeError`: For various mismatches between the state of the table and @@ -1173,8 +2916,9 @@ def get_or_create_table( 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.' @@ -1213,8 +2957,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,7 +2972,7 @@ 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. @@ -1253,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, @@ -1262,31 +3029,40 @@ 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: + # 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 + page_token = response.pageToken + else: + # Modern RowIterator + yield list(response), _to_table_schema(response.schema) break - page_token = response.pageToken + def insert_rows( self, @@ -1305,6 +3081,7 @@ def insert_rows( 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 @@ -1314,10 +3091,8 @@ def insert_rows( Returns: A tuple (bool, errors). If first element is False then the second element - will be a bigquery.InsertErrorsValueListEntry instance containing - specific errors. + 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 @@ -1333,77 +3108,96 @@ 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': + f_type = field.type if hasattr(field, 'type') else field.field_type + if f_type == 'STRING': # Input: "XYZ" --> Output: "XYZ" return value - elif field.type == 'BOOLEAN': + elif f_type in ('BOOLEAN', 'BOOL'): # Input: "true" --> Output: True - return value == 'true' - elif field.type == 'INTEGER': + return value == 'true' if isinstance(value, str) else bool(value) + elif f_type in ('INTEGER', 'INT64'): # Input: "123" --> Output: 123 return int(value) - elif field.type == 'FLOAT': + elif f_type in ('FLOAT', 'FLOAT64'): # Input: "1.23" --> Output: 1.23 return float(value) - elif field.type == 'TIMESTAMP': + 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" - dt = datetime.datetime.utcfromtimestamp(float(value)) - return dt.strftime('%Y-%m-%d %H:%M:%S.%f UTC') - elif field.type == 'BYTES': + 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 field.type == 'DATE': + elif f_type == 'DATE': # Input: "2016-11-03" --> Output: "2016-11-03" - return value - elif field.type == 'DATETIME': + return str(value) + elif f_type == 'DATETIME': # Input: "2016-11-03T00:49:36" --> Output: "2016-11-03T00:49:36" - return value - elif field.type == 'TIME': + return str(value) + elif f_type == 'TIME': # Input: "00:49:36" --> Output: "00:49:36" - return value - elif field.type == 'RECORD': + 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 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) @@ -1418,10 +3212,50 @@ def from_pipeline_options(pipeline_options: PipelineOptions): client=BigQueryWrapper._bigquery_client(pipeline_options)) @staticmethod - def _bigquery_client(pipeline_options: PipelineOptions): - return bigquery.BigqueryV2( + 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 + project = pipeline_options.view_as(GoogleCloudOptions).project + if hasattr(project, 'get') and callable(project.get): + project = project.get() + except Exception: + project = None + try: + from apache_beam.options.pipeline_options import DebugOptions + experiments = pipeline_options.view_as(DebugOptions).experiments or [] + except Exception: + 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(), - credentials=auth.get_service_credentials(pipeline_options), + credentials=raw_credentials, response_encoding='utf8', additional_http_headers={ "user-agent": "apache-beam-%s" % apache_beam.__version__ @@ -1629,20 +3463,24 @@ def beam_row_from_dict(row: dict, schema): Args: row (dict): The row to convert. - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): + 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, (bigquery.TableSchema, bigquery.TableFieldSchema)): + 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() + mode = (getattr(field, 'mode', None) or 'NULLABLE').upper() + 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. @@ -1650,11 +3488,11 @@ def beam_row_from_dict(row: dict, schema): # 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": + if row and name not in row and mode != "REQUIRED": row[name] = None - value = row[name] - if type in ["RECORD", "STRUCT"] and value: + 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 = [] @@ -1670,63 +3508,73 @@ 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. + """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: - ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema: + 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. + but in the TableSchema format. """ - table_schema = bigquery.TableSchema() + 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 table_field_to_dict(field): + """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. """ - 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 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 - 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_schema_to_dict(table_schema): + """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'): + 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]} def get_dict_table_schema(schema): """Transform the table schema into a dictionary instance. Args: - schema (str, dict, ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema): + 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. @@ -1740,7 +3588,7 @@ def get_dict_table_schema(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) @@ -1750,24 +3598,31 @@ 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): + 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: - ~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema: The schema as a TableSchema object. + TableSchema: The schema as a TableSchema object. """ - if (isinstance(schema, - (bigquery.TableSchema, value_provider.ValueProvider)) or + if (isinstance(schema, (TableSchema, value_provider.ValueProvider)) or callable(schema) or schema is None): return schema + 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) elif isinstance(schema, dict): 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) @@ -1776,8 +3631,7 @@ 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): + 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. @@ -1793,8 +3647,7 @@ 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): + 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 @@ -1805,12 +3658,22 @@ def get_beam_typehints_from_tableschema(schema, type_overrides=None): 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)): + 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. @@ -1851,8 +3714,8 @@ 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', typing.Any], + right: Union['TableSchema', 'TableFieldSchema', typing.Any], *, ignore_descriptions: bool = False, ignore_field_order: bool = False) -> bool: @@ -1863,13 +3726,9 @@ def check_schema_equal( field ordering (optionally). Args: - left (~apache_beam.io.gcp.internal.clients.bigquery.\ -bigquery_v2_messages.TableSchema, ~apache_beam.io.gcp.internal.clients.\ -bigquery.bigquery_v2_messages.TableFieldSchema): + left (TableSchema, TableFieldSchema, SchemaField, list, dict): 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): + 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. @@ -1879,44 +3738,68 @@ def check_schema_equal( Returns: bool: True if the schemas are equivalent, False otherwise. """ - if type(left) != type(right) or not isinstance( - left, (bigquery.TableSchema, bigquery.TableFieldSchema)): + if left is None and right is None: + return True + if left is None or right is None: return False - if isinstance(left, bigquery.TableFieldSchema): + 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'))) + + if is_field_left != is_field_right: + return False + + if is_field_left: if left.name != right.name: return False - if left.type != right.type: + 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: # Check for type aliases - if sorted( - (left.type, right.type)) not in (["BOOL", "BOOLEAN"], ["FLOAT", - "FLOAT64"], - ["INT64", "INTEGER"], ["RECORD", - "STRUCT"]): + 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..058ff18dc1ee 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -35,10 +35,12 @@ 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 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 +52,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 @@ -566,6 +569,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( @@ -1409,6 +1414,491 @@ 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([])) + 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) + + +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) + + @unittest.skipIf( + bigquery_tools.gcp_bigquery is None, + 'google-cloud-bigquery not installed') + def test_table_reference_property_mutability(self): + 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') + + +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) + + 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) unittest.main() 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/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 diff --git a/sdks/python/setup.py b/sdks/python/setup.py index 5b02dc9ab0ef..fe23337abc4b 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -528,7 +528,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