diff --git a/sdks/python/apache_beam/io/gcp/BIGQUERY_COMPAT.md b/sdks/python/apache_beam/io/gcp/BIGQUERY_COMPAT.md new file mode 100644 index 000000000000..e31024dbbdd9 --- /dev/null +++ b/sdks/python/apache_beam/io/gcp/BIGQUERY_COMPAT.md @@ -0,0 +1,236 @@ + + +# BigQuery Compatibility Shims Usage & Migration Guide + +## Overview + +Apache Beam's Python SDK has migrated its BigQuery I/O connector from the legacy, apitools-generated client (`apache_beam.io.gcp.internal.clients.bigquery`) to Google's official [`google-cloud-bigquery`](https://cloud.google.com/python/docs/reference/bigquery/latest) client library (version `3.0.0` or higher). + +To facilitate a smooth transition for pipelines that interact with internal or legacy BigQuery data structures, Beam provides temporary compatibility shims in `apache_beam.io.gcp.bigquery_compat`. + +--- + +## Enabling the Compatibility Shims (Explicit Opt-In) + +The compatibility shims are **disabled by default** to encourage modern client adoption. If your pipeline or test code directly imports or depends on legacy `bigquery_v2_client` or `bigquery_client_messages` interfaces, you must explicitly opt in. + +### Setting the Environment Variable + +Set `BEAM_USE_BIGQUERY_COMPAT_SHIMS=1` in your environment before importing Beam BigQuery modules: + +#### Shell / Terminal: +```bash +export BEAM_USE_BIGQUERY_COMPAT_SHIMS=1 +python my_pipeline.py +``` + +#### In Python Code (before importing Beam BigQuery modules): +```python +import os +os.environ["BEAM_USE_BIGQUERY_COMPAT_SHIMS"] = "1" + +from apache_beam.io.gcp import bigquery_compat +``` + +### What Happens Without Opt-In? + +Attempting to import `apache_beam.io.gcp.bigquery_compat` without this environment variable will raise an `ImportError`: + +```text +ImportError: BigQuery compatibility shims in apache_beam.io.gcp.bigquery_compat are disabled by default and require explicit opt-in. Set the environment variable BEAM_USE_BIGQUERY_COMPAT_SHIMS=1 to enable them. +``` + +When enabled, a `BeamDeprecationWarning` is emitted to remind users that these shims will be removed in a future release. + +--- + +## What the Compatibility Layer Provides + +When opted in, `bigquery_compat` provides: + +### 1. Dynamic Attribute Bridging (Monkey-Patching) +Augments modern `google.cloud.bigquery` classes with legacy camelCase property accessors: + +| Modern Class | Modern Attribute (snake_case) | Compatibility Accessor (camelCase) | +| :--- | :--- | :--- | +| `TableReference` | `.project` | `.projectId` | +| `TableReference` | `.dataset_id` | `.datasetId` | +| `TableReference` | `.table_id` | `.tableId` | +| `DatasetReference` | `.project` | `.projectId` | +| `DatasetReference` | `.dataset_id` | `.datasetId` | +| `SchemaField` | `.field_type` | `.type` | +| `Table` | `.reference` | `.tableReference` | +| `Table` | `.num_rows` | `.numRows` | +| `Table` | `.num_bytes` | `.numBytes` | +| `Table` | `.time_partitioning` | `.timePartitioning` (also accepts raw dicts) | +| `Table` | `.range_partitioning` | `.rangePartitioning` (also accepts raw dicts) | +| `Table` | `.clustering_fields` | `.clustering` (also accepts raw dicts/lists) | +| `_AsyncJob` | `.job_id`, `.project`, `.location` | `.jobReference` | +| `_AsyncJob` | `.state`, `.error_result` | `.status.state`, `.status.errorResult` | +| `_AsyncJob` | `.total_bytes_billed` | `.statistics.query.totalBytesBilled` | + +### 2. Client Emulation +Wraps modern `google.cloud.bigquery.Client` instances to support legacy apitools method patterns: +* `client.tables.Get(request)` $\rightarrow$ delegates to `client.get_table(...)` +* `client.tables.Insert(request)` $\rightarrow$ delegates to `client.create_table(...)` +* `client.datasets.Get(request)` $\rightarrow$ delegates to `client.get_dataset(...)` +* `client.jobs.GetQueryResults(request)` $\rightarrow$ delegates to `client.list_rows(...)` or `job.result(...)` +* `client.jobs.Insert(request)` $\rightarrow$ delegates to `client.query`, `client.load_table_from_uri`, `client.copy_table`, or `client.extract_table` + +### 3. Standalone Compatibility Models +When `google-cloud-bigquery` is not installed or when running in lightweight environments: +* `JobReference(jobId=..., projectId=..., location=...)` +* `TableReference(projectId=..., datasetId=..., tableId=...)` +* `TableSchema(fields=[...])`, `TableFieldSchema(name=..., type=...)` + +--- + +## Migration Guide: Modernizing Your Code + +To future-proof your pipeline and prepare for the removal of the compatibility shims, update your code according to the patterns below. + +### 1. Table References + +#### Legacy (Deprecated): +```python +from apache_beam.io.gcp.internal.clients.bigquery import TableReference + +table_ref = TableReference(projectId='my-project', datasetId='my_dataset', tableId='my_table') +print(table_ref.projectId) +``` + +#### Modern (Recommended): +```python +from google.cloud.bigquery import DatasetReference +from google.cloud.bigquery import TableReference + +# Using from_string (easiest): +table_ref = TableReference.from_string("my-project.my_dataset.my_table") + +# Or constructor: +table_ref = TableReference(DatasetReference("my-project", "my_dataset"), "my_table") + +# Modern snake_case attributes: +print(table_ref.project) +print(table_ref.dataset_id) +print(table_ref.table_id) +``` + +--- + +### 2. Table Schemas + +#### Legacy (Deprecated): +```python +from apache_beam.io.gcp.internal.clients.bigquery import TableFieldSchema +from apache_beam.io.gcp.internal.clients.bigquery import TableSchema + +schema = TableSchema(fields=[ + TableFieldSchema(name='id', type='INTEGER', mode='REQUIRED'), + TableFieldSchema(name='name', type='STRING', mode='NULLABLE'), +]) +``` + +#### Modern (Recommended): +```python +from google.cloud.bigquery import SchemaField + +# Modern SchemaField list: +schema = [ + SchemaField("id", "INTEGER", mode="REQUIRED"), + SchemaField("name", "STRING", mode="NULLABLE"), +] + +# Or pass a simple string schema to Beam transforms: +# schema = "id:INTEGER, name:STRING" +``` + +--- + +### 3. Additional Table Parameters (Partitioning & Clustering) + +In `WriteToBigQuery(additional_bq_parameters=...)`, you can continue passing dictionaries or use modern classes: + +```python +from google.cloud.bigquery import PartitionRange +from google.cloud.bigquery import RangePartitioning +from google.cloud.bigquery import TimePartitioning + +# Native Time Partitioning: +time_partitioning = TimePartitioning(type_="DAY", field="timestamp") + +# Native Range Partitioning: +range_partitioning = RangePartitioning( + field="customer_id", + range_=PartitionRange(start=0, end=100000, interval=100) +) +``` + +--- + +### 4. Exception Handling + +#### Legacy (Deprecated): +```python +from apitools.base.py.exceptions import HttpError +from apitools.base.py.exceptions import HttpForbiddenError + +try: + ... +except HttpError as e: + if e.status_code == 404: + ... +``` + +#### Modern (Recommended): +```python +from google.api_core import exceptions + +try: + ... +except exceptions.NotFound: + # Resource not found (404) + ... +except exceptions.Forbidden: + # Permission denied (403) + ... +except exceptions.GoogleAPICallError as e: + # Base class for all Google Cloud API call errors + ... +``` + +--- + +## Deprecation & Removal Roadmap + +```mermaid +flowchart LR + A["Phase 1: Opt-In Shims (Current)
BEAM_USE_BIGQUERY_COMPAT_SHIMS=1"] --> B["Phase 2: Deprecation Grace Period
(2-3 Minor Releases)"] + B --> C["Phase 3: Complete Removal
Apitools & bigquery_compat.py Deleted"] +``` + +1. **Current Release**: + * Modern `google-cloud-bigquery` is the default execution engine. + * `bigquery_compat.py` provides opt-in shims via `BEAM_USE_BIGQUERY_COMPAT_SHIMS=1`. + * `apache_beam.io.gcp.internal.clients.bigquery` is deprecated. +2. **Upcoming Release (Target: Beam 2.80.0)**: + * `apache_beam.io.gcp.internal.clients.bigquery` will be deleted. + * `apache_beam.io.gcp.bigquery_compat` and its test suite will be deleted. + * The `BEAM_USE_BIGQUERY_COMPAT_SHIMS` environment variable will be decommissioned. diff --git a/sdks/python/apache_beam/io/gcp/bigquery.py b/sdks/python/apache_beam/io/gcp/bigquery.py index 38acd29da7d9..cd083d26edb0 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery.py +++ b/sdks/python/apache_beam/io/gcp/bigquery.py @@ -390,8 +390,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 @@ -422,13 +428,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__) @@ -555,8 +560,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: @@ -824,14 +833,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 @@ -851,8 +888,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: @@ -2335,8 +2371,8 @@ def expand(self, pcoll): p = pcoll.pipeline if (isinstance(self.table_reference, TableReference) and - self.table_reference.projectId is None): - self.table_reference.projectId = pcoll.pipeline.options.view_as( + self.table_reference.project is None): + self.table_reference.project = pcoll.pipeline.options.view_as( GoogleCloudOptions).project # TODO(pabloem): Use a different method to determine if streaming or batch. @@ -3175,11 +3211,26 @@ 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, - tableId='beam_temp_table_' + uuid.uuid4().hex) + 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, 'project', None) or getattr( + temp_dataset, 'projectId', None) + dataset_id = getattr(temp_dataset, 'dataset_id', None) or getattr( + temp_dataset, 'datasetId', None) + temp_table_ref = TableReference( + project=project_id, + dataset_id=dataset_id, + table_id='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..128aa2e80972 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) @@ -496,8 +534,7 @@ def _emit_query_ranges( self, start_ts: Timestamp, end_ts: Timestamp, - watermark_estimator: _PollWatermarkEstimator - ) -> Iterable[TimestampedValue[_QueryRange]]: + watermark_estimator: _PollWatermarkEstimator) -> Iterable[Any]: """Compute and yield _QueryRange elements, advancing estimator state.""" ranges = compute_ranges(start_ts, end_ts, self._change_function) _LOGGER.info( @@ -525,7 +562,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() @@ -622,32 +659,45 @@ def process(self, qr: _QueryRange) -> Iterable[_QueryResult]: _utc(qr.chunk_start), _utc(qr.chunk_end)) - temp_table_ref = bigquery.TableReference( - projectId=self._project, - 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) + temp_table_ref = TableReference( + project=self._project, + dataset_id=self._temp_dataset, + table_id=temp_table_id) + + 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, 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..e3c189cf611d 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) @@ -185,7 +205,12 @@ def test_exact_day_boundary(self): self.assertEqual(len(ranges), 2) -@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') +from apache_beam.io.gcp import bigquery_change_history + + +@unittest.skipIf( + bigquery_change_history.bq_storage is None, + 'google-cloud-bigquery-storage is not installed') class ValidationTest(unittest.TestCase): """Tests for ReadBigQueryChangeHistory validation.""" def test_invalid_change_function(self): diff --git a/sdks/python/apache_beam/io/gcp/bigquery_compat.py b/sdks/python/apache_beam/io/gcp/bigquery_compat.py index 3ad76bbcc1d4..b928a24ca3d5 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_compat.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_compat.py @@ -25,6 +25,7 @@ Importing this module is disabled by default and requires explicit opt-in. To enable these compatibility shims, set the environment variable ``BEAM_USE_BIGQUERY_COMPAT_SHIMS=1``. + For detailed usage and migration guidance, see ``BIGQUERY_COMPAT.md``. .. note:: This module is intended to be removed in a future Beam release once the @@ -68,12 +69,13 @@ def _check_compat_opt_in(): raise ImportError( f"BigQuery compatibility shims in {__name__} are disabled by default " f"and require explicit opt-in. Set the environment variable " - f"{BIGQUERY_COMPAT_ENV_VAR}=1 to enable them.") + f"{BIGQUERY_COMPAT_ENV_VAR}=1 to enable them. For migration guidance, " + "see apache_beam/io/gcp/BIGQUERY_COMPAT.md.") warnings.warn( f"BigQuery compatibility shims in {__name__} are deprecated and will be " "removed in a future release of Apache Beam. Please migrate to using " - "google-cloud-bigquery directly.", + "google-cloud-bigquery directly (see apache_beam/io/gcp/BIGQUERY_COMPAT.md).", BeamDeprecationWarning, stacklevel=2) @@ -1422,6 +1424,8 @@ def Insert(self, request, upload=None): _to_gcp_table_ref(s, default_project=proj) for s in getattr(cp, "sourceTables", []) ] + if not sources and getattr(cp, "sourceTable", None): + sources = [_to_gcp_table_ref(cp.sourceTable, default_project=proj)] dest = _to_gcp_table_ref( getattr(cp, "destinationTable", None), default_project=proj) job_config = gcp_bigquery.CopyJobConfig() 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..ec09d243cea9 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 @@ -423,18 +449,20 @@ def process(self, element, schema_mod_job_name_prefix): return table_reference = bigquery_tools.parse_table_reference(destination) - if table_reference.projectId is None: - table_reference.projectId = vp.RuntimeValueProvider.get_value( + if table_reference.project is None: + table_reference.project = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project try: # Check if destination table exists destination_table = self.bq_wrapper.get_table( - project_id=table_reference.projectId, - dataset_id=table_reference.datasetId, - table_id=table_reference.tableId) - except HttpError as exn: - if exn.status_code == 404: + project_id=table_reference.project, + dataset_id=table_reference.dataset_id, + table_id=table_reference.table_id) + 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 @@ -442,10 +470,21 @@ def process(self, element, schema_mod_job_name_prefix): raise temp_table_load_job = self.bq_wrapper.get_job( - 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 + project=( + getattr(temp_table_load_job_reference, 'project', None) or + getattr(temp_table_load_job_reference, 'projectId', None)), + job_id=( + getattr(temp_table_load_job_reference, 'job_id', None) or + getattr(temp_table_load_job_reference, 'jobId', None)), + location=getattr(temp_table_load_job_reference, 'location', None)) + 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, @@ -457,9 +496,9 @@ def process(self, element, schema_mod_job_name_prefix): destination_hash = _bq_uuid( '%s:%s.%s' % ( - table_reference.projectId, - table_reference.datasetId, - table_reference.tableId)) + table_reference.project, + table_reference.dataset_id, + table_reference.table_id)) uid = _bq_uuid() job_name = '%s_%s_%s' % (schema_mod_job_name_prefix, destination_hash, uid) @@ -563,14 +602,16 @@ def process_one(self, element, job_name_prefix): destination, job_reference = element copy_to_reference = bigquery_tools.parse_table_reference(destination) - if copy_to_reference.projectId is None: - copy_to_reference.projectId = vp.RuntimeValueProvider.get_value( + if copy_to_reference.project is None: + copy_to_reference.project = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project copy_from_reference = bigquery_tools.parse_table_reference(destination) - copy_from_reference.tableId = job_reference.jobId - if copy_from_reference.projectId is None: - copy_from_reference.projectId = vp.RuntimeValueProvider.get_value( + copy_from_reference.table_id = ( + getattr(job_reference, 'job_id', None) or + getattr(job_reference, 'jobId', None)) + if copy_from_reference.project is None: + copy_from_reference.project = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project _LOGGER.info( @@ -585,15 +626,15 @@ def process_one(self, element, job_name_prefix): self.bq_io_metadata = create_bigquery_io_metadata(self._step_name) project_id = ( - copy_to_reference.projectId + copy_to_reference.project if self.load_job_project_id is None else self.load_job_project_id) copy_job_name = '%s_%s' % ( job_name_prefix, _bq_uuid( '%s:%s.%s' % ( - copy_from_reference.projectId, - copy_from_reference.datasetId, - copy_from_reference.tableId))) + copy_from_reference.project, + copy_from_reference.dataset_id, + copy_from_reference.table_id))) job_reference = self.bq_wrapper._insert_copy_job( project_id, copy_job_name, @@ -628,18 +669,18 @@ def _determine_write_disposition(self, copy_to_reference) -> tuple[bool, str]: complete and the write disposition to use for the job. """ full_table_ref = '%s:%s.%s' % ( - copy_to_reference.projectId, - copy_to_reference.datasetId, - copy_to_reference.tableId) + copy_to_reference.project, + copy_to_reference.dataset_id, + copy_to_reference.table_id) if full_table_ref not in self._observed_tables: write_disposition = self.write_disposition wait_for_job = write_disposition != 'WRITE_APPEND' self._observed_tables.add(full_table_ref) Lineage.sinks().add( 'bigquery', - copy_to_reference.projectId, - copy_to_reference.datasetId, - copy_to_reference.tableId) + copy_to_reference.project, + copy_to_reference.dataset_id, + copy_to_reference.table_id) else: wait_for_job = False write_disposition = 'WRITE_APPEND' @@ -746,17 +787,17 @@ def process( additional_parameters = dict(additional_parameters or {}) table_reference = bigquery_tools.parse_table_reference(destination) - if table_reference.projectId is None: - table_reference.projectId = vp.RuntimeValueProvider.get_value( + if table_reference.project is None: + table_reference.project = vp.RuntimeValueProvider.get_value( 'project', str, '') or self.project # Load jobs for a single destination are always triggered from the same # worker. This means that we can generate a deterministic numbered job id, # and not need to worry. destination_hash = _bq_uuid( '%s:%s.%s' % ( - table_reference.projectId, - table_reference.datasetId, - table_reference.tableId)) + table_reference.project, + table_reference.dataset_id, + table_reference.table_id)) job_name = '%s_%s_pane%s_partition%s' % ( load_job_name_prefix, destination_hash, pane_info.index, partition_key) _LOGGER.info('Load job has %s files. Job name is %s.', len(files), job_name) @@ -774,9 +815,9 @@ def process( destination_table = self.destination_table_cache[hashed_dest] else: destination_table = self.bq_wrapper.get_table( - project_id=table_reference.projectId, - dataset_id=table_reference.datasetId, - table_id=table_reference.tableId) + project_id=table_reference.project, + dataset_id=table_reference.dataset_id, + table_id=table_reference.table_id) self.destination_table_cache[hashed_dest] = destination_table except Exception as e: if need_schema: @@ -797,8 +838,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( @@ -817,16 +863,16 @@ def process( # temporary tables, so we replace the create_disposition. create_disposition = 'CREATE_IF_NEEDED' # For temporary tables, we create a new table with the name with JobId. - table_reference.tableId = job_name + table_reference.table_id = job_name yield pvalue.TaggedOutput( TriggerLoadJobs.TEMP_TABLES, bigquery_tools.get_hashable_destination(table_reference)) else: Lineage.sinks().add( 'bigquery', - table_reference.projectId, - table_reference.datasetId, - table_reference.tableId) + table_reference.project, + table_reference.dataset_id, + table_reference.table_id) _LOGGER.info( 'Triggering job %s to load data to BigQuery table %s.' @@ -938,9 +984,9 @@ def process(self, table_reference): _LOGGER.info("Deleting table %s", table_reference) table_reference = bigquery_tools.parse_table_reference(table_reference) self.bq_wrapper._delete_table( - table_reference.projectId, - table_reference.datasetId, - table_reference.tableId) + table_reference.project, + table_reference.dataset_id, + table_reference.table_id) class BigQueryBatchFileLoads(beam.PTransform): @@ -1238,7 +1284,7 @@ def _load_data( finished_temp_tables_load_job_ids_list_pc = ( finished_temp_tables_load_job_ids_pc | beam.MapTuple( lambda destination, job_reference: ( - bigquery_tools.parse_table_reference(destination).tableId, + bigquery_tools.parse_table_reference(destination).table_id, (destination, job_reference))) | beam.GroupByKey() | beam.MapTuple(lambda tableId, batch: list(batch))) 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..acfc32e1591e 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,23 +226,36 @@ 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, 'dataset_id'): + return self.temp_dataset.dataset_id + elif hasattr(self.temp_dataset, 'datasetId'): return self.temp_dataset.datasetId 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 else: raise ValueError("temp_dataset has to be either str or DatasetReference") def _get_temp_dataset_project(self): """Returns the project ID for temporary dataset operations. - - If temp_dataset is a DatasetReference, returns its projectId. + + If temp_dataset is a DatasetReference, returns its project. Otherwise, returns the pipeline project for billing. """ - if isinstance(self.temp_dataset, DatasetReference): + if hasattr(self.temp_dataset, 'project') and self.temp_dataset.project: + return self.temp_dataset.project + elif hasattr(self.temp_dataset, + 'projectId') and self.temp_dataset.projectId: return self.temp_dataset.projectId - else: - return self._get_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 self._get_project() def start_bundle(self): self.bq = bigquery_tools.BigQueryWrapper( @@ -423,7 +432,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 +467,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_read_internal_test.py b/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py index 46673b4ec2d2..eebe42d1afba 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_read_internal_test.py @@ -21,15 +21,11 @@ from unittest import mock from apache_beam.io.gcp import bigquery_read_internal +from apache_beam.io.gcp.bigquery_tools import DatasetReference from apache_beam.options.pipeline_options import GoogleCloudOptions from apache_beam.options.pipeline_options import PipelineOptions from apache_beam.options.value_provider import StaticValueProvider -try: - from apache_beam.io.gcp.internal.clients.bigquery import DatasetReference -except ImportError: - DatasetReference = None - class BigQueryReadSplitTest(unittest.TestCase): """Tests for _BigQueryReadSplit DoFn.""" 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 dcadee7f6a1a..c67abf82b0a2 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.TableFieldSchema(name='s', type='STRING'), + bigquery.TableFieldSchema(name='i', type='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,25 @@ 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.TableFieldSchema( + name='float', type='FLOAT', mode='NULLABLE'), + bigquery.TableFieldSchema( + name='string', type='STRING', mode='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 +487,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 +557,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 +583,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 +679,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 +703,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 +744,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 +759,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 +773,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 +797,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( @@ -2349,7 +2403,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..ddbcd6b0af7f 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,52 @@ # Protect against environments where bigquery library is not available. try: import regex +except ImportError: + regex = None + +try: + from apitools.base.protorpclite import messages as _protorpclite_messages +except ImportError: + _protorpclite_messages = None + +try: + 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 +except ImportError: + ClientInfo = None + +try: 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 +except ImportError: + GoogleAPICallError = type('GoogleAPICallError', (Exception, ), {}) + ClientError = type('ClientError', (GoogleAPICallError, ), {}) + Conflict = type('Conflict', (ClientError, ), {'code': 409}) + Forbidden = type('Forbidden', (ClientError, ), {'code': 403}) + NotFound = type('NotFound', (ClientError, ), {'code': 404}) + ServerError = type('ServerError', (GoogleAPICallError, ), {}) + +try: from google.cloud import bigquery as gcp_bigquery -except Exception: + from google.cloud.bigquery import job as gcp_job +except ImportError: gcp_bigquery = None - pass + gcp_job = None try: from orjson import dumps as fast_json_dumps @@ -87,14 +131,700 @@ fast_json_dumps = json.dumps fast_json_loads = json.loads -# pylint: enable=wrong-import-order, wrong-import-position +# ----------------------------------------------------------------------------- +# BigQuery Compatibility & Conversion Layer +# +# Helpers, models, and monkey-patches for interacting with google-cloud-bigquery +# and supporting legacy compatibility. +# ----------------------------------------------------------------------------- + + +class DatasetReference(object): + """Container for BigQuery dataset identification (dataset_id, project). + + Supports both modern snake_case (project, dataset_id) and legacy camelCase + (projectId, datasetId) initialization and attribute access. + """ + 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 + self._dataset_id = d + + @classmethod + def from_string(cls, dataset_ref, default_project=None): + last_sep = max(dataset_ref.rfind('.'), dataset_ref.rfind(':')) + if last_sep != -1: + p = dataset_ref[:last_sep] + d = dataset_ref[last_sep + 1:] + else: + p = default_project + d = dataset_ref + return cls(project=p, dataset_id=d) + + @property + def project(self): + return self._project + + @project.setter + def project(self, val): + self._project = val + + @property + def projectId(self): + return self._project + + @projectId.setter + def projectId(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 dataset_id(self): + return self._dataset_id + + @dataset_id.setter + def dataset_id(self, val): + self._dataset_id = val + + @property + def datasetId(self): + return self._dataset_id + + @datasetId.setter + def datasetId(self, val): + self._dataset_id = val + + def __repr__(self): + return ( + f"DatasetReference(project='{self._project}', " + f"dataset_id='{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, 'project', None) or getattr( + other, 'projectId', None) + other_d = getattr(other, 'dataset_id', None) or getattr( + other, 'datasetId', None) + return (self._project, self._dataset_id) == (other_p, other_d) + + def __hash__(self): + return hash((self._project, self._dataset_id)) + + +class TableReference(object): + """Container for BigQuery table identification (table_id, dataset_id, project). + + Supports both modern snake_case (project, dataset_id, table_id) and legacy + camelCase (projectId, datasetId, tableId) initialization and attribute access. + """ + 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 dataset_ref is not None: + self._project = getattr(dataset_ref, 'project', None) or getattr( + dataset_ref, 'projectId', None) or p + self._dataset_id = getattr(dataset_ref, 'dataset_id', None) or getattr( + dataset_ref, 'datasetId', None) or d + self._table_id = t + else: + self._project = p + self._dataset_id = d + self._table_id = t + + @classmethod + def from_string(cls, table_ref, default_project=None): + return parse_table_reference(table_ref, project=default_project) + + @property + def dataset_ref(self): + if self._dataset_id is not None or self._project is not None: + return DatasetReference( + project=self._project, dataset_id=self._dataset_id) + return None + + @property + def project(self): + return self._project + + @project.setter + def project(self, val): + self._project = val + + @property + def projectId(self): + return self._project + + @projectId.setter + def projectId(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 dataset_id(self): + return self._dataset_id + + @dataset_id.setter + def dataset_id(self, val): + self._dataset_id = val + + @property + def datasetId(self): + return self._dataset_id + + @datasetId.setter + def datasetId(self, val): + self._dataset_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 tableId(self): + return self._table_id + + @tableId.setter + def tableId(self, val): + self._table_id = val + + def __repr__(self): + return ( + f"TableReference(project='{self._project}', " + f"dataset_id='{self._dataset_id}', table_id='{self._table_id}')") + + def __eq__(self, other): + if other is None: + return False + if not hasattr(other, 'table_id') and not hasattr(other, 'tableId'): + return NotImplemented + other_p = getattr(other, 'project', None) or getattr( + other, 'projectId', None) + other_d = getattr(other, 'dataset_id', None) or getattr( + other, 'datasetId', None) + other_t = getattr(other, 'table_id', None) or getattr( + other, 'tableId', None) + return (self._project, self._dataset_id, + self._table_id) == (other_p, other_d, other_t) + + def __hash__(self): + return hash((self._project, self._dataset_id, self._table_id)) + + +TableCell = getattr(apitools_bigquery, 'TableCell', None) +TableFieldSchema = getattr(apitools_bigquery, 'TableFieldSchema', None) +TableRow = getattr(apitools_bigquery, 'TableRow', None) +TableSchema = getattr(apitools_bigquery, 'TableSchema', None) + + +class JobReference(object): + """Container for BigQuery job identification (jobId, projectId, location).""" + 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) + other_j = getattr(other, 'jobId', None) or getattr(other, 'job_id', None) + other_p = getattr(other, 'projectId', None) or getattr( + other, 'project', None) + other_l = getattr(other, 'location', None) + return (self.jobId, self.projectId, + self.location) == (other_j, other_p, other_l) + + 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})") + + +if _protorpclite_messages is not None and 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 + + +def _extract_field(obj, *field_names): + """Extracts the first matching non-None field from a dict or object.""" + if obj is None: + return None + for name in field_names: + if isinstance(obj, dict): + if name in obj and obj[name] is not None: + return obj[name] + elif hasattr(obj, name): + val = getattr(obj, name) + if val is not None and not hasattr(val, '_mock_methods'): + return val + return None + + +def _to_gcp_time_partitioning(tp): + """Converts a dict, apitools TimePartitioning, string, or google.cloud.bigquery.TimePartitioning.""" + if tp is None: + return None + if hasattr(tp, '_mock_methods') or hasattr(tp, '_mock_children'): + return tp + if gcp_bigquery is not None and isinstance( + tp, getattr(gcp_bigquery, 'TimePartitioning', ())): + return tp + if isinstance(tp, str): + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'TimePartitioning'): + return gcp_bigquery.TimePartitioning(type_=tp) + return tp + + tp_field = _extract_field(tp, 'field') + tp_type = _extract_field(tp, 'type', 'type_') + tp_exp = _extract_field(tp, 'expirationMs', 'expiration_ms') + tp_req = _extract_field( + tp, 'requirePartitionFilter', 'require_partition_filter') + + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'TimePartitioning'): + kwargs = {} + if tp_field is not None: + kwargs['field'] = tp_field + if tp_type is not None: + kwargs['type_'] = tp_type + if tp_exp is not None: + kwargs['expiration_ms'] = tp_exp + if tp_req is not None: + kwargs['require_partition_filter'] = tp_req + return gcp_bigquery.TimePartitioning(**kwargs) + return tp + + +def _to_gcp_range_partitioning(rp): + """Converts a dict, apitools RangePartitioning, or google.cloud.bigquery.RangePartitioning.""" + if rp is None: + return None + if hasattr(rp, '_mock_methods') or hasattr(rp, '_mock_children'): + return rp + if gcp_bigquery is not None and isinstance( + rp, getattr(gcp_bigquery, 'RangePartitioning', ())): + return rp + + rp_field = _extract_field(rp, 'field') + rp_range = _extract_field(rp, 'range') + + start = _extract_field(rp_range, 'start') if rp_range else None + end = _extract_field(rp_range, 'end') if rp_range else None + interval = _extract_field(rp_range, 'interval') if rp_range else None + + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'RangePartitioning'): + pr = None + if (start is not None or end is not None or + interval is not None) and hasattr(gcp_bigquery, 'PartitionRange'): + pr = gcp_bigquery.PartitionRange(start=start, end=end, interval=interval) + return gcp_bigquery.RangePartitioning(range_=pr, field=rp_field) + return rp + + +def _to_gcp_clustering_fields(clustering): + """Extracts clustering fields list from list, dict, or Clustering model.""" + if clustering is None: + return None + if hasattr(clustering, '_mock_methods') or hasattr(clustering, + '_mock_children'): + return clustering + if isinstance(clustering, (list, tuple)): + return list(clustering) + fields = _extract_field(clustering, 'fields') + if fields is not None: + return list(fields) + return None + + +def _to_json_compatible(obj): + """Converts an object or nested structure to JSON/API-compatible dicts/types. + + .. note:: + This is a temporary helper used to extract dictionary representations + from legacy apitools ``protorpclite.Message`` objects via + ``encoding.MessageToDict``. It will be removed once the legacy apitools + client code is excised. + """ + 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() + try: + from apitools.base.protorpclite import messages as _protorpclite_messages + if isinstance(obj, _protorpclite_messages.Message): + 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. + + .. note:: + This is a temporary helper to translate legacy apitools ``LabelsValue`` + protobuf messages into standard Python dictionaries. It will be removed + or simplified to a direct dictionary check once the legacy client code + is excised. + """ + 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 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): + last_sep = max(dataset_ref.rfind('.'), dataset_ref.rfind(':')) + if last_sep != -1: + proj = dataset_ref[:last_sep] + ds_id = dataset_ref[last_sep + 1:] + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): + return gcp_bigquery.DatasetReference(proj, ds_id) + return dataset_ref + proj = project or 'default' + if gcp_bigquery is not None and hasattr(gcp_bigquery, 'DatasetReference'): + return gcp_bigquery.DatasetReference(proj, dataset_ref) + return 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 dataset_ref + + +def _to_apitools_table_ref(table_ref, default_project=None): + """Converts a TableReference, dict, or string into an apitools TableReference. + + .. note:: + This helper supports backward compatibility with the legacy BigQuery + client and will be removed once the apitools client code is excised. + """ + if table_ref is None: + return None + if (apitools_bigquery is not None and + hasattr(apitools_bigquery, 'TableReference') and + isinstance(table_ref, apitools_bigquery.TableReference)): + return table_ref + if isinstance(table_ref, str): + table_ref = parse_table_reference(table_ref, project=default_project) + proj = getattr(table_ref, 'project', None) or getattr( + table_ref, 'projectId', None) or getattr( + table_ref, 'project_id', None) or default_project + dataset_id = getattr(table_ref, 'dataset_id', None) or getattr( + table_ref, 'datasetId', None) or getattr(table_ref, 'dataset', None) + table_id = getattr(table_ref, 'table_id', None) or getattr( + table_ref, 'tableId', None) or getattr(table_ref, 'table', None) + if apitools_bigquery and hasattr(apitools_bigquery, 'TableReference'): + return apitools_bigquery.TableReference( + projectId=proj, datasetId=dataset_id, tableId=table_id) + return table_ref + + +def _to_apitools_dataset_ref(dataset_ref, default_project=None): + """Converts a DatasetReference, dict, or string into an apitools DatasetReference. + + .. note:: + This helper supports backward compatibility with the legacy BigQuery + client and will be removed once the apitools client code is excised. + """ + if dataset_ref is None: + return None + if (apitools_bigquery is not None and + hasattr(apitools_bigquery, 'DatasetReference') and + isinstance(dataset_ref, apitools_bigquery.DatasetReference)): + return dataset_ref + if isinstance(dataset_ref, str): + last_sep = max(dataset_ref.rfind('.'), dataset_ref.rfind(':')) + if last_sep != -1: + proj = dataset_ref[:last_sep] + ds_id = dataset_ref[last_sep + 1:] + else: + proj = default_project + ds_id = dataset_ref + else: + proj = getattr(dataset_ref, 'project', None) or getattr( + dataset_ref, 'projectId', None) or getattr( + dataset_ref, 'project_id', None) or default_project + ds_id = getattr(dataset_ref, 'dataset_id', None) or getattr( + dataset_ref, 'datasetId', None) + if apitools_bigquery and hasattr(apitools_bigquery, 'DatasetReference'): + return apitools_bigquery.DatasetReference(projectId=proj, datasetId=ds_id) + return dataset_ref + + +def _table_field_to_dict(field): + """Create a dictionary representation of a table field for serialization.""" + if isinstance(field, dict): + return field + result = {} + result['name'] = getattr(field, 'name', '') + result['type'] = getattr(field, 'type', None) or getattr( + field, 'field_type', None) or 'STRING' + result['mode'] = getattr(field, 'mode', 'NULLABLE') or 'NULLABLE' + if hasattr(field, 'description') and field.description is not None: + result['description'] = field.description + if hasattr(field, 'fields') and field.fields: + result['fields'] = [_table_field_to_dict(f) for f in field.fields] + return result + + +def _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. + + .. note:: + This is a temporary backwards-compatibility helper that translates modern + ``google.cloud.bigquery.SchemaField`` instances back into the legacy + ``TableSchema`` model for operators like ``TableRowJsonCoder``. It will be + removed in a future release once Beam coders and schema utilities natively + consume modern schema types. + """ + if TableSchema is None: + return schema + 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 TableFieldSchema is None: + return 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() -# pylint: disable=wrong-import-order, wrong-import-position, ungrouped-imports -try: - from apache_beam.io.gcp.internal.clients.bigquery import TableReference -except ImportError: - TableReference = None -# pylint: enable=wrong-import-order, wrong-import-position, ungrouped-imports _LOGGER = logging.getLogger(__name__) @@ -107,7 +837,10 @@ _PROJECT_PATTERN = r'([a-z0-9.-]+:)?[a-z][a-z0-9-]*[a-z0-9]' _DATASET_PATTERN = r'\w{1,1024}' -_TABLE_PATTERN = r'[\p{L}\p{M}\p{N}\p{Pc}\p{Pd}\p{Zs}$]{1,1024}' +_TABLE_PATTERN_REGEX = r'[\p{L}\p{M}\p{N}\p{Pc}\p{Pd}\p{Zs}$]{1,1024}' +_TABLE_PATTERN_RE = r'[\w\s\-$]{1,1024}' +_TABLE_PATTERN = ( + _TABLE_PATTERN_REGEX if regex is not None else _TABLE_PATTERN_RE) # TODO(https://github.com/apache/beam/issues/25946): Add support for # more Beam portable schema types as Python types @@ -153,6 +886,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): @@ -167,17 +902,29 @@ def get_hashable_destination(destination): """Parses a table reference into a (project, dataset, table) tuple. Args: - destination: Either a TableReference object from the bigquery API. - The object has the following attributes: projectId, datasetId, and - tableId. Or a string representing the destination containing - 'PROJECT:DATASET.TABLE'. + destination: Either a TableReference object or a string representing the + destination containing 'PROJECT:DATASET.TABLE'. Returns: - A string representing the destination containing - 'PROJECT:DATASET.TABLE'. + A string representing the destination containing 'PROJECT:DATASET.TABLE'. """ - if isinstance(destination, TableReference): - return '%s:%s.%s' % ( - destination.projectId, destination.datasetId, destination.tableId) + if isinstance(destination, + TableReference) or (gcp_bigquery is not None and isinstance( + destination, getattr(gcp_bigquery, 'TableReference', ()))): + proj = getattr(destination, 'project', None) or getattr( + destination, 'projectId', None) + ds = getattr(destination, 'dataset_id', None) or getattr( + destination, 'datasetId', None) + tbl = getattr(destination, 'table_id', None) or getattr( + destination, 'tableId', None) + return '%s:%s.%s' % (proj, ds, tbl) + elif hasattr(destination, 'tableId') or hasattr(destination, 'table_id'): + proj = getattr(destination, 'project', None) or getattr( + destination, 'projectId', None) + ds = getattr(destination, 'dataset_id', None) or getattr( + destination, 'datasetId', None) + tbl = getattr(destination, 'table_id', None) or getattr( + destination, 'tableId', None) + return '%s:%s.%s' % (proj, ds, tbl) else: return destination @@ -225,21 +972,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): @@ -261,46 +1008,58 @@ def parse_table_reference(table, dataset=None, project=None): argument. Returns: - A TableReference object from the bigquery API. The object has the following - attributes: projectId, datasetId, and tableId. + A TableReference object. The object supports both modern snake_case + attributes (project, dataset_id, table_id) and legacy camelCase attributes + (projectId, datasetId, tableId). If the input is a TableReference object, a new object will be returned. Raises: ValueError: if the table reference as a string does not match the expected format. """ - if isinstance(table, TableReference): return TableReference( - projectId=table.projectId, - datasetId=table.datasetId, - tableId=table.tableId) + project=project or table.project, + dataset_id=table.dataset_id, + table_id=table.table_id) + elif isinstance(table, getattr(gcp_bigquery, 'TableReference', ())): + return TableReference( + project=project or getattr(table, 'project', None), + dataset_id=getattr(table, 'dataset_id', None), + table_id=getattr(table, 'table_id', None)) + elif (apitools_bigquery and hasattr(apitools_bigquery, 'TableReference') and + isinstance(table, apitools_bigquery.TableReference)): + return TableReference( + project=project or getattr(table, 'projectId', None), + dataset_id=getattr(table, 'datasetId', None), + table_id=getattr(table, 'tableId', None)) elif callable(table): return table elif isinstance(table, value_provider.ValueProvider): return table - table_reference = TableReference() # If dataset argument is not specified, the expectation is that the # table argument will contain a full table reference instead of just a # table name. if dataset is None: + table_pattern = ( + _TABLE_PATTERN_REGEX if regex is not None else _TABLE_PATTERN_RE) pattern = ( f'((?P{_PROJECT_PATTERN})[:\\.])?' - f'(?P{_DATASET_PATTERN})\\.(?P{_TABLE_PATTERN})') - match = regex.fullmatch(pattern, table) + f'(?P{_DATASET_PATTERN})\\.(?P
{table_pattern})') + match = ( + regex.fullmatch(pattern, table) if regex is not None else re.fullmatch( + pattern, table)) if not match: 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( + project=match.group('project') or project, + dataset_id=match.group('dataset'), + table_id=match.group('table')) else: - table_reference.projectId = project - table_reference.datasetId = dataset - table_reference.tableId = table - return table_reference + return TableReference(project=project, dataset_id=dataset, table_id=table) # ----------------------------------------------------------------------------- @@ -308,31 +1067,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 +1105,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 +1130,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 @@ -393,7 +1166,9 @@ def __init__(self, client=None, temp_dataset_id=None, temp_table_ref=None): if temp_table_ref is not None: self.temp_table_ref = temp_table_ref - self.temp_dataset_id = temp_table_ref.datasetId + self.temp_dataset_id = ( + getattr(temp_table_ref, 'dataset_id', None) or + getattr(temp_table_ref, 'datasetId', None)) else: self.temp_table_ref = None self._temporary_table_suffix = uuid.uuid4().hex @@ -401,6 +1176,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. @@ -432,56 +1232,115 @@ def _get_temp_table_project(self, fallback_project_id): Otherwise, returns the fallback_project_id. """ if self.temp_table_ref: - return self.temp_table_ref.projectId + return ( + getattr(self.temp_table_ref, 'project', None) or + getattr(self.temp_table_ref, 'projectId', None)) else: return fallback_project_id def _get_temp_dataset(self): if self.temp_table_ref: - return self.temp_table_ref.datasetId + return ( + getattr(self.temp_table_ref, 'dataset_id', None) or + getattr(self.temp_table_ref, 'datasetId', None)) return BigQueryWrapper.TEMP_DATASET + self._temporary_table_suffix @retry.with_exponential_backoff( 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 +1350,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,16 +1368,53 @@ 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( - destinationTable=to_table_reference, - sourceTable=from_table_reference, + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( + copy=apitools_bigquery.JobConfigurationTableCopy( + destinationTable=_to_apitools_table_ref( + to_table_reference, default_project=project_id), + sourceTable=_to_apitools_table_ref( + from_table_reference, default_project=project_id), createDisposition=create_disposition, writeDisposition=write_disposition, ), @@ -527,7 +1422,6 @@ def _insert_copy_job( ), jobReference=reference, )) - return self._start_job(request).jobReference @retry.with_exponential_backoff( @@ -557,26 +1451,125 @@ 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( + legacy_load_params = {} + if apitools_bigquery and hasattr(apitools_bigquery, 'JobConfigurationLoad'): + for k, v in additional_load_parameters.items(): + if '_' in k: + parts = k.split('_') + camel_k = parts[0] + ''.join(x.title() for x in parts[1:]) + else: + camel_k = k + if hasattr(apitools_bigquery.JobConfigurationLoad, camel_k): + legacy_load_params[camel_k] = v + elif hasattr(apitools_bigquery.JobConfigurationLoad, k): + legacy_load_params[k] = v + else: + legacy_load_params = additional_load_parameters + + 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, + destinationTable=_to_apitools_table_ref( + table_reference, default_project=project_id), schema=job_schema, writeDisposition=write_disposition, createDisposition=create_disposition, sourceFormat=source_format, useAvroLogicalTypes=True, autodetect=schema == 'SCHEMA_AUTODETECT', - **additional_load_parameters), + **legacy_load_params), labels=_build_job_labels(job_labels), ), jobReference=reference, @@ -598,10 +1591,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 +1604,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 +1645,81 @@ 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=_to_apitools_table_ref( + 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), + default_project=project_id), 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 +1729,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 +1774,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 +1806,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 +1826,95 @@ def _insert_all_rows( request_count_urn=monitoring_infos.API_REQUEST_COUNT_URN, base_labels=labels) + if self._is_modern_client: + started_millis = int(time.time() * 1000) + try: + table_ref_str = '%s.%s.%s' % (project_id, dataset_id, table_id) + row_ids_arg = insert_ids + if row_ids_arg is not None and all(x is None for x in row_ids_arg): + auto_uuid = getattr( + getattr(gcp_bigquery, 'AutoRowIDs', None), 'GENERATE_UUID', None) + if auto_uuid is not None: + row_ids_arg = auto_uuid + errors = self.gcp_bq_client.insert_rows_json( + table_ref_str, + json_rows=rows, + row_ids=row_ids_arg, + skip_invalid_rows=skip_invalid_rows, + ignore_unknown_values=ignore_unknown_values, + timeout=BQ_STREAMING_INSERT_TIMEOUT_SEC) + if not errors: + service_call_metric.call('ok') + else: + for insert_error in errors: + for err in insert_error.get('errors', []): + reason = err.get('reason') if isinstance(err, dict) else getattr( + err, 'reason', None) + service_call_metric.call(reason or 'unknown') + except (ClientError, GoogleAPICallError, HttpError) as e: + status_code = getattr(e, 'code', None) or getattr( + e, 'status_code', None) or 500 + service_call_metric.call(status_code) + # Package exception with required fields for BigQueryWriteFn + reason = None + if hasattr(e, 'response') and getattr(e.response, 'reason', None): + reason = e.response.reason + elif hasattr(e, 'reason') and e.reason: + reason = e.reason + elif hasattr(e, 'errors') and e.errors and isinstance( + e.errors, (list, tuple)) and isinstance(e.errors[0], dict): + reason = e.errors[0].get('reason') + if not reason: + reason = e.__class__.__name__ + errors = [{ + 'index': i, 'errors': [{ + 'reason': reason, 'message': str(e) + }] + } for i in range(len(rows))] + if not errors: + errors = [{ + 'index': 0, 'errors': [{ + 'reason': reason, 'message': str(e) + }] + }] + finally: + self._latency_histogram_metric.update( + int(time.time() * 1000) - started_millis) + return not errors, errors + + # Legacy apitools path + # 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 +1931,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( + project=project_id, dataset_id=dataset_id, table_id=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 @@ -824,20 +1961,74 @@ def _create_table( schema, additional_parameters=None): - valid_tablename = regex.fullmatch(_TABLE_PATTERN, table_id, regex.ASCII) + if regex is not None: + valid_tablename = regex.fullmatch( + _TABLE_PATTERN_REGEX, table_id, regex.ASCII) + else: + valid_tablename = re.fullmatch(_TABLE_PATTERN_RE, table_id, re.ASCII) if not valid_tablename: raise ValueError( 'Invalid BigQuery table name: %s \n' '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: + clean_params = _to_json_compatible(additional_parameters) + if isinstance(clean_params, dict): + for k, v in clean_params.items(): + if k in ('timePartitioning', 'time_partitioning'): + table.time_partitioning = _to_gcp_time_partitioning(v) + elif k in ('rangePartitioning', 'range_partitioning'): + table.range_partitioning = _to_gcp_range_partitioning(v) + elif k in ('clustering', 'clustering_fields'): + table.clustering_fields = _to_gcp_clustering_fields(v) + elif k in ('friendlyName', 'friendly_name'): + table.friendly_name = v + elif k in ('description', ): + table.description = v + elif k in ('labels', ): + table.labels = _extract_dict_labels(v) or {} + elif k in ('encryptionConfiguration', 'encryption_configuration'): + kms = getattr(v, 'kmsKeyName', None) or getattr( + v, 'kms_key_name', None) or ( + v.get('kmsKeyName') or v.get('kms_key_name') + if isinstance(v, dict) else None) + if kms: + table.encryption_configuration = ( + gcp_bigquery.EncryptionConfiguration(kms_key_name=kms)) + else: + try: + setattr(table, k, v) + except Exception: + pass + else: + for k, v in additional_parameters.items(): + try: + setattr(table, k, v) + except Exception: + pass + 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( - tableReference=TableReference( - projectId=project_id, datasetId=dataset_id, tableId=table_id), - schema=schema, - **additional_parameters) - request = bigquery.BigqueryTablesInsertRequest( + table_ref = _to_apitools_table_ref( + TableReference( + project=project_id, dataset_id=dataset_id, table_id=table_id), + default_project=project_id) + table_schema = _to_table_schema(schema) if schema is not None else None + table = apitools_bigquery.Table( + tableReference=table_ref, schema=table_schema, **additional_parameters) + 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 +2047,42 @@ 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(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 +2094,9 @@ def get_or_create_dataset( project_id, dataset_id, location) - dataset_reference = bigquery.DatasetReference( - projectId=project_id, datasetId=dataset_id) - dataset = bigquery.Dataset(datasetReference=dataset_reference) + dataset_reference = _to_apitools_dataset_ref( + dataset_id, default_project=project_id) + dataset = apitools_bigquery.Dataset(datasetReference=dataset_reference) if location is not None: dataset.location = location if labels is not None: @@ -882,7 +2106,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 +2119,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 +2143,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( + project=project_id, dataset_id=dataset_id, table_id=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 +2176,22 @@ def _delete_table(self, project_id, dataset_id, table_id): num_retries=MAX_RETRIES, retry_filter=retry.retry_on_server_errors_and_timeout_filter) def _delete_dataset(self, project_id, dataset_id, delete_contents=True): - request = bigquery.BigqueryDatasetsDeleteRequest( + if self._is_modern_client: + dataset_ref = _to_gcp_dataset_ref(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,14 +2243,51 @@ 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) + temp_ds_id = getattr(temp_table, 'dataset_id', None) or getattr( + temp_table, 'datasetId', None) + temp_proj_id = getattr(temp_table, 'project', None) or getattr( + temp_table, 'projectId', None) or project_id + temp_tbl_id = getattr(temp_table, 'table_id', None) or getattr( + temp_table, 'tableId', None) + + if self._is_modern_client: + dataset_ref = _to_gcp_dataset_ref(temp_ds_id, 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_ds_id) + return + raise + try: + # We do not want to delete temporary datasets configured by the user hence + # we just delete the temporary table in that case. + if not self.is_user_configured_dataset(): + self._delete_dataset(temp_proj_id, temp_ds_id, True) + else: + self._delete_table(temp_proj_id, temp_ds_id, temp_tbl_id) + 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_proj_id, + temp_ds_id) + return + raise + return + + # Fallback for legacy client try: self.client.datasets.Get( - bigquery.BigqueryDatasetsGetRequest( - projectId=project_id, datasetId=temp_table.datasetId)) + apitools_bigquery.BigqueryDatasetsGetRequest( + projectId=project_id, datasetId=temp_ds_id)) except HttpError as exn: if exn.status_code == 404: - _LOGGER.warning( - 'Dataset %s:%s does not exist', project_id, temp_table.datasetId) + _LOGGER.warning('Dataset %s:%s does not exist', project_id, temp_ds_id) return else: raise @@ -991,17 +2295,16 @@ def clean_up_temporary_dataset(self, project_id): # 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) + self._delete_dataset(temp_proj_id, temp_ds_id, True) else: - self._delete_table( - temp_table.projectId, temp_table.datasetId, temp_table.tableId) + self._delete_table(temp_proj_id, temp_ds_id, temp_tbl_id) self.created_temp_dataset = False except HttpError as exn: if exn.status_code == 403: _LOGGER.warning( 'Permission denied to delete temporary dataset %s:%s for clean up', - temp_table.projectId, - temp_table.datasetId) + temp_proj_id, + temp_ds_id) return else: raise @@ -1013,31 +2316,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 +2375,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 +2401,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,17 +2443,54 @@ 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)) + dest_uris = destination if isinstance(destination, list) else [destination] + request = apitools_bigquery.BigqueryJobsInsertRequest( projectId=job_project, - job=bigquery.Job( - configuration=bigquery.JobConfiguration( - extract=bigquery.JobConfigurationExtract( - destinationUris=destination, - sourceTable=table_reference, + job=apitools_bigquery.Job( + configuration=apitools_bigquery.JobConfiguration( + extract=apitools_bigquery.JobConfigurationExtract( + destinationUris=dest_uris, + sourceTable=_to_apitools_table_ref( + table_reference, default_project=job_project), printHeader=include_header, destinationFormat=destination_format, compression=compression, @@ -1155,12 +2524,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 +2542,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 +2583,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 +2598,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 +2624,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 +2655,39 @@ 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 +2706,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 +2716,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 +2733,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 +2837,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 +3088,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 +3113,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 +3133,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 +3213,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 +3223,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 +3256,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 +3272,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 +3283,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 +3339,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 +3351,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 +3363,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..26a47405f810 100644 --- a/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py +++ b/sdks/python/apache_beam/io/gcp/bigquery_tools_test.py @@ -35,12 +35,16 @@ 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 DatasetReference 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 TableReference from apache_beam.io.gcp.bigquery_tools import beam_row_from_dict from apache_beam.io.gcp.bigquery_tools import check_schema_equal from apache_beam.io.gcp.bigquery_tools import generate_bq_job_name @@ -50,6 +54,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 @@ -145,7 +150,10 @@ def test_calling_with_fully_qualified_table_ref( table_id: str, ): parsed_ref = parse_table_reference(fully_qualified_table) - self.assertIsInstance(parsed_ref, bigquery.TableReference) + self.assertIsInstance(parsed_ref, TableReference) + self.assertEqual(parsed_ref.project, project_id) + self.assertEqual(parsed_ref.dataset_id, dataset_id) + self.assertEqual(parsed_ref.table_id, table_id) self.assertEqual(parsed_ref.projectId, project_id) self.assertEqual(parsed_ref.datasetId, dataset_id) self.assertEqual(parsed_ref.tableId, table_id) @@ -155,7 +163,9 @@ def test_calling_with_partially_qualified_table_ref(self): tableId = 'test_table' partially_qualified_table = '{}.{}'.format(datasetId, tableId) parsed_ref = parse_table_reference(partially_qualified_table) - self.assertIsInstance(parsed_ref, bigquery.TableReference) + self.assertIsInstance(parsed_ref, TableReference) + self.assertEqual(parsed_ref.dataset_id, datasetId) + self.assertEqual(parsed_ref.table_id, tableId) self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) @@ -169,11 +179,35 @@ def test_calling_with_all_arguments(self): tableId = 'test_table' parsed_ref = parse_table_reference( tableId, dataset=datasetId, project=projectId) - self.assertIsInstance(parsed_ref, bigquery.TableReference) + self.assertIsInstance(parsed_ref, TableReference) + self.assertEqual(parsed_ref.project, projectId) + self.assertEqual(parsed_ref.dataset_id, datasetId) + self.assertEqual(parsed_ref.table_id, tableId) self.assertEqual(parsed_ref.projectId, projectId) self.assertEqual(parsed_ref.datasetId, datasetId) self.assertEqual(parsed_ref.tableId, tableId) + def test_parse_table_reference_without_regex_package(self): + with mock.patch.object(bigquery_tools, 'regex', None): + parsed_ref = parse_table_reference('my-project:my_dataset.my_table') + self.assertEqual(parsed_ref.project, 'my-project') + self.assertEqual(parsed_ref.dataset_id, 'my_dataset') + self.assertEqual(parsed_ref.table_id, 'my_table') + self.assertEqual(parsed_ref.projectId, 'my-project') + self.assertEqual(parsed_ref.datasetId, 'my_dataset') + self.assertEqual(parsed_ref.tableId, 'my_table') + + parsed_ref2 = parse_table_reference('my_dataset.my_table$20250101') + self.assertIsNone(parsed_ref2.project) + self.assertIsNone(parsed_ref2.projectId) + self.assertEqual(parsed_ref2.dataset_id, 'my_dataset') + self.assertEqual(parsed_ref2.table_id, 'my_table$20250101') + self.assertEqual(parsed_ref2.datasetId, 'my_dataset') + self.assertEqual(parsed_ref2.tableId, 'my_table$20250101') + + with self.assertRaises(ValueError): + parse_table_reference('invalid_table_ref') + @unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') class TestBigQueryWrapper(unittest.TestCase): @@ -565,6 +599,9 @@ def test_insert_rows_sets_metric_on_failure(self): [], ]) wrapper = beam.io.gcp.bigquery_tools.BigQueryWrapper(client) + # Simulate caller (BigQueryWriteFn) retrying across 3 sequential attempts + 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). @@ -1408,6 +1445,256 @@ def test_type_overrides_json_to_dict(self): typehints_dict = get_beam_typehints_from_tableschema(schema, {"JSON": dict}) self.assertEqual(typehints_dict, [("data", Optional[dict])]) + def test_perform_load_job_with_modern_client_additional_parameters(self): + if bigquery_tools.gcp_bigquery is None: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=bigquery_tools.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_start_query_job_with_none_labels(self): + if bigquery_tools.gcp_bigquery is None: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + client = mock.MagicMock(spec=bigquery_tools.gcp_bigquery.Client) + mock_job = mock.MagicMock() + mock_job.job_id = 'query_job_id' + mock_job.project = 'test-project' + mock_job.location = 'US' + 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_create_table_with_additional_parameters_partitioning_and_clustering( + self): + if bigquery_tools.gcp_bigquery is None: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + + mock_client = mock.Mock(spec=bigquery_tools.gcp_bigquery.Client) + mock_client.create_table = mock.Mock(side_effect=lambda table: table) + wrapper = BigQueryWrapper(client=mock_client) + + schema = { + 'fields': [{ + 'name': 'name', 'type': 'STRING', 'mode': 'NULLABLE' + }, { + 'name': 'language', 'type': 'STRING', 'mode': 'NULLABLE' + }] + } + additional_params = { + 'timePartitioning': { + 'type': 'DAY' + }, + 'clustering': { + 'fields': ['language'] + } + } + res = wrapper._create_table( + project_id='test-project', + dataset_id='test_dataset', + table_id='test_table', + schema=schema, + additional_parameters=additional_params) + mock_client.create_table.assert_called_once() + created_table = mock_client.create_table.call_args[0][0] + self.assertEqual(created_table.time_partitioning.type_, 'DAY') + self.assertEqual(created_table.clustering_fields, ['language']) + self.assertEqual(len(created_table.schema), 2) + self.assertEqual(res, created_table) + + +@unittest.skipIf(HttpError is None, 'GCP dependencies are not installed') +class TestBigQueryClientExperimentFallback(unittest.TestCase): + def test_default_client_is_modern(self): + if bigquery_tools.gcp_bigquery is None: + raise unittest.SkipTest('google-cloud-bigquery is not installed') + wrapper = BigQueryWrapper.from_pipeline_options(PipelineOptions([])) + self.assertTrue(wrapper._is_modern_client) + + def test_default_client_is_legacy_when_no_modern_client(self): + if bigquery_tools.gcp_bigquery is not None: + raise unittest.SkipTest('google-cloud-bigquery is installed') + wrapper = BigQueryWrapper.from_pipeline_options(PipelineOptions([])) + self.assertFalse(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 TestTableAndDatasetReferenceCompatibility(unittest.TestCase): + def test_dataset_reference_init_and_properties(self): + ref1 = DatasetReference(project='my-proj', dataset_id='my_ds') + self.assertEqual(ref1.project, 'my-proj') + self.assertEqual(ref1.dataset_id, 'my_ds') + self.assertEqual(ref1.projectId, 'my-proj') + self.assertEqual(ref1.datasetId, 'my_ds') + + # Mutability + ref1.project = 'new-proj' + ref1.dataset_id = 'new_ds' + self.assertEqual(ref1.projectId, 'new-proj') + self.assertEqual(ref1.datasetId, 'new_ds') + + # Legacy kwargs + ref2 = DatasetReference(projectId='legacy-proj', datasetId='legacy_ds') + self.assertEqual(ref2.project, 'legacy-proj') + self.assertEqual(ref2.dataset_id, 'legacy_ds') + + # Equality + ref3 = DatasetReference(project='legacy-proj', dataset_id='legacy_ds') + self.assertEqual(ref2, ref3) + self.assertEqual(hash(ref2), hash(ref3)) + + def test_table_reference_init_and_properties(self): + # Keyword init + ref1 = TableReference( + project='my-proj', dataset_id='my_ds', table_id='my_tbl') + self.assertEqual(ref1.project, 'my-proj') + self.assertEqual(ref1.dataset_id, 'my_ds') + self.assertEqual(ref1.table_id, 'my_tbl') + self.assertEqual(ref1.projectId, 'my-proj') + self.assertEqual(ref1.datasetId, 'my_ds') + self.assertEqual(ref1.tableId, 'my_tbl') + self.assertEqual( + ref1.dataset_ref, + DatasetReference(project='my-proj', dataset_id='my_ds')) + + # Positional init with DatasetReference + ds_ref = DatasetReference(project='my-proj', dataset_id='my_ds') + ref2 = TableReference(ds_ref, 'my_tbl') + self.assertEqual(ref1, ref2) + + # Legacy kwargs + ref3 = TableReference( + projectId='my-proj', datasetId='my_ds', tableId='my_tbl') + self.assertEqual(ref1, ref3) + + # Mutability + ref3.table_id = 'mutated_tbl' + self.assertEqual(ref3.table_id, 'mutated_tbl') + self.assertEqual(ref3.tableId, 'mutated_tbl') + + ref3.tableId = 'mutated_again' + self.assertEqual(ref3.table_id, 'mutated_again') + + def test_conversions_to_gcp_and_apitools(self): + table_ref = TableReference(project='proj', dataset_id='ds', table_id='tbl') + gcp_ref = bigquery_tools._to_gcp_table_ref(table_ref) + if bigquery_tools.gcp_bigquery is not None: + self.assertIsInstance(gcp_ref, bigquery_tools.gcp_bigquery.TableReference) + self.assertEqual(gcp_ref.project, 'proj') + self.assertEqual(gcp_ref.dataset_id, 'ds') + self.assertEqual(gcp_ref.table_id, 'tbl') + + api_ref = bigquery_tools._to_apitools_table_ref(table_ref) + if bigquery_tools.apitools_bigquery is not None and hasattr( + bigquery_tools.apitools_bigquery, 'TableReference'): + self.assertIsInstance( + api_ref, bigquery_tools.apitools_bigquery.TableReference) + self.assertEqual(api_ref.projectId, 'proj') + self.assertEqual(api_ref.datasetId, 'ds') + self.assertEqual(api_ref.tableId, 'tbl') + + def test_legacy_client_fallback_job_insertion_converts_modern_references( + self): + wrapper = BigQueryWrapper(use_legacy_client=True) + wrapper.client = mock.MagicMock() + mock_insert = mock.MagicMock() + wrapper.client.jobs.Insert = mock_insert + mock_job = mock.MagicMock() + mock_job.jobReference = bigquery_tools.JobReference(job_id='test_job') + mock_insert.return_value = mock_job + + table_ref = TableReference( + project='test_proj', dataset_id='test_ds', table_id='test_tbl') + + # Test _insert_load_job with modern table reference on legacy client + wrapper._insert_load_job( + project_id='test_proj', + job_id='load_job_1', + table_reference=table_ref, + source_uris=['gs://bucket/file.csv']) + self.assertTrue(mock_insert.called) + req = mock_insert.call_args[0][0] + if bigquery_tools.apitools_bigquery and hasattr( + bigquery_tools.apitools_bigquery, 'TableReference'): + self.assertIsInstance( + req.job.configuration.load.destinationTable, + bigquery_tools.apitools_bigquery.TableReference) + self.assertEqual( + req.job.configuration.load.destinationTable.tableId, 'test_tbl') + + # Test _insert_copy_job with modern table references on legacy client + mock_insert.reset_mock() + wrapper._insert_copy_job( + project_id='test_proj', + job_id='copy_job_1', + from_table_reference=table_ref, + to_table_reference=TableReference( + project='test_proj', dataset_id='test_ds', table_id='dest_tbl')) + self.assertTrue(mock_insert.called) + req = mock_insert.call_args[0][0] + if bigquery_tools.apitools_bigquery and hasattr( + bigquery_tools.apitools_bigquery, 'TableReference'): + self.assertIsInstance( + req.job.configuration.copy.destinationTable, + bigquery_tools.apitools_bigquery.TableReference) + self.assertIsInstance( + req.job.configuration.copy.sourceTable, + bigquery_tools.apitools_bigquery.TableReference) + if __name__ == '__main__': logging.getLogger().setLevel(logging.INFO) 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 bafe94a1b010..b1a0c1f292c6 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,11 +15,27 @@ # 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 logging 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 26a68d7c9fe3..0a4346ca9fdc 100644 --- a/sdks/python/setup.py +++ b/sdks/python/setup.py @@ -529,7 +529,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.44.0 changed DirectRow mutation storage; native WriteToBigTable