This is an automated email from the ASF dual-hosted git repository.
jrmccluskey pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/beam.git
The following commit(s) were added to refs/heads/master by this push:
new f915665446e [Gemini] Add compatibility classes between apitools client
and modern bq client (#39998)
f915665446e is described below
commit f915665446e70f227fab99dd1739626615b0e387
Author: Jack McCluskey <[email protected]>
AuthorDate: Tue Sep 29 12:56:42 2026 -0400
[Gemini] Add compatibility classes between apitools client and modern bq
client (#39998)
* [Gemini] Add compatibility classes between apitools client and modern bq
classes
* fix docstring format
* Streamline compat classes
* manage funky apitools import
* Add explicit opt-in via environment variable
---
sdks/python/apache_beam/io/gcp/bigquery_compat.py | 1454 ++++++++++++++++++++
.../apache_beam/io/gcp/bigquery_compat_test.py | 796 +++++++++++
2 files changed, 2250 insertions(+)
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_compat.py
b/sdks/python/apache_beam/io/gcp/bigquery_compat.py
new file mode 100644
index 00000000000..3ad76bbcc1d
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/bigquery_compat.py
@@ -0,0 +1,1454 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Compatibility shims and legacy client emulation for BigQuery.
+
+This module contains temporary compatibility models, monkey patches, and
+helpers designed to ease migration away from the deprecated apitools BigQuery
+client to modern ``google-cloud-bigquery``.
+
+.. warning::
+ 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``.
+
+.. note::
+ This module is intended to be removed in a future Beam release once the
+ apitools client dependency is completely removed.
+
+ **Future Removal Guidance**:
+
+ * Compatibility models (e.g. ``_TableReferenceCompat``,
+ ``_DatasetReferenceCompat``, ``_TableSchemaCompat``) and monkey-patches
+ will be dropped. Code using ``TableReference``, ``DatasetReference``,
+ etc. should import directly from ``google.cloud.bigquery``.
+ * Input normalization helpers (``_to_gcp_table_ref``,
+ ``_to_gcp_dataset_ref``, ``_to_gcp_schema``, ``_extract_dict_labels``,
+ ``_to_table_schema``) that convert string specs or dicts to modern
+ ``google.cloud.bigquery`` instances are actively used across pipeline
+ code paths and should be preserved in ``bigquery_tools.py`` when this
+ compat module is excised.
+"""
+
+# pytype: skip-file
+
+import logging
+import os
+import warnings
+
+from apache_beam.utils.annotations import BeamDeprecationWarning
+
+_LOGGER = logging.getLogger(__name__)
+
+BIGQUERY_COMPAT_ENV_VAR = "BEAM_USE_BIGQUERY_COMPAT_SHIMS"
+
+
+def _check_compat_opt_in():
+ """Checks whether BigQuery compatibility shims are enabled via environment
variable.
+
+ Raises:
+ ImportError: If the environment variable is not set to an enabled value.
+ """
+ val = os.environ.get(BIGQUERY_COMPAT_ENV_VAR, "").strip().lower()
+ if val not in ("1", "true", "yes", "enabled", "on"):
+ 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.")
+
+ 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.",
+ BeamDeprecationWarning,
+ stacklevel=2)
+
+
+_check_compat_opt_in()
+
+try:
+ from google.cloud import bigquery as gcp_bigquery
+ from google.cloud.bigquery import job as gcp_job
+except ImportError:
+ gcp_bigquery = None
+ gcp_job = None
+
+try:
+ from apache_beam.io.gcp.internal.clients import bigquery as apitools_bigquery
+except ImportError:
+ apitools_bigquery = None
+
+try:
+ from apitools.base.protorpclite import messages as _protorpclite_messages
+except ImportError:
+ _protorpclite_messages = None
+
+_LOGGER = logging.getLogger(__name__)
+
+# -----------------------------------------------------------------------------
+# Compatibility Models for TableReference, DatasetReference, Schema, and Jobs.
+#
+# These classes and monkey patches bridge between legacy apitools structures
+# and modern google.cloud.bigquery objects, providing camelCase attribute
access
+# (e.g. projectId, datasetId, tableId, tableReference) for backwards
+# compatibility across pipelines, transforms, and test suites.
+# -----------------------------------------------------------------------------
+
+
+class _DatasetReferenceCompat(object):
+ """Compatibility model for BigQuery DatasetReference when
google-cloud-bigquery is unavailable.
+
+ Supports both camelCase (projectId, datasetId) and snake_case (project,
dataset_id, project_id).
+ """
+ def __init__(
+ self,
+ project=None,
+ dataset_id=None,
+ projectId=None,
+ datasetId=None,
+ project_id=None):
+ p = (
+ projectId if projectId is not None else
+ (project_id if project_id is not None else project))
+ d = datasetId if datasetId is not None else dataset_id
+ self._project = p or ""
+ self._dataset_id = d or ""
+
+ @classmethod
+ def from_string(cls, dataset_ref, default_project=None):
+ 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 or "default"
+ d = dataset_ref
+ return cls(project=p, dataset_id=d)
+
+ @property
+ def projectId(self):
+ return self._project
+
+ @projectId.setter
+ def projectId(self, val):
+ self._project = val
+
+ @property
+ def project(self):
+ return self._project
+
+ @project.setter
+ def project(self, val):
+ self._project = val
+
+ @property
+ def project_id(self):
+ return self._project
+
+ @project_id.setter
+ def project_id(self, val):
+ self._project = val
+
+ @property
+ def datasetId(self):
+ return self._dataset_id
+
+ @datasetId.setter
+ def datasetId(self, val):
+ self._dataset_id = val
+
+ @property
+ def dataset_id(self):
+ return self._dataset_id
+
+ @dataset_id.setter
+ def dataset_id(self, val):
+ self._dataset_id = val
+
+ def __repr__(self):
+ return f"DatasetReference('{self.project}', '{self.dataset_id}')"
+
+ def __eq__(self, other):
+ if other is None:
+ return False
+ if not hasattr(other, "project") and not hasattr(other, "projectId"):
+ return NotImplemented
+ other_p = getattr(other, "projectId", None) or getattr(
+ other, "project", None)
+ other_d = getattr(other, "datasetId", None) or getattr(
+ other, "dataset_id", None)
+ return (self.projectId, self.datasetId) == (other_p, other_d)
+
+ def __hash__(self):
+ return hash((self.projectId, self.datasetId))
+
+
+class _TableReferenceCompat(object):
+ """Compatibility model for BigQuery TableReference when
google-cloud-bigquery is unavailable.
+
+ Supports both camelCase (projectId, datasetId, tableId) and snake_case
+ (project, dataset_id, table_id, project_id).
+ """
+ def __init__(
+ self,
+ dataset_ref=None,
+ table_id=None,
+ projectId=None,
+ datasetId=None,
+ tableId=None,
+ project=None,
+ dataset_id=None,
+ project_id=None):
+ p = (
+ projectId if projectId is not None else
+ (project_id if project_id is not None else project))
+ d = datasetId if datasetId is not None else dataset_id
+ t = tableId if tableId is not None else table_id
+ if p is not None or d is not None or t is not None:
+ self._project = p
+ self._dataset_id = d
+ self._table_id = t
+ elif dataset_ref is not None:
+ self._project = getattr(dataset_ref, "projectId", None) or getattr(
+ dataset_ref, "project", None)
+ self._dataset_id = getattr(dataset_ref, "datasetId", None) or getattr(
+ dataset_ref, "dataset_id", None)
+ self._table_id = table_id or ""
+ else:
+ self._project = None
+ self._dataset_id = None
+ self._table_id = None
+
+ @classmethod
+ def from_string(cls, table_ref, default_project=None):
+ from apache_beam.io.gcp.bigquery_tools import parse_table_reference
+ parsed = parse_table_reference(table_ref, project=default_project)
+ return cls(
+ projectId=parsed.projectId or default_project,
+ datasetId=parsed.datasetId,
+ tableId=parsed.tableId)
+
+ @property
+ def projectId(self):
+ return self._project
+
+ @projectId.setter
+ def projectId(self, val):
+ self._project = val
+
+ @property
+ def project(self):
+ return self._project
+
+ @project.setter
+ def project(self, val):
+ self._project = val
+
+ @property
+ def project_id(self):
+ return self._project
+
+ @project_id.setter
+ def project_id(self, val):
+ self._project = val
+
+ @property
+ def datasetId(self):
+ return self._dataset_id
+
+ @datasetId.setter
+ def datasetId(self, val):
+ self._dataset_id = val
+
+ @property
+ def dataset_id(self):
+ return self._dataset_id
+
+ @dataset_id.setter
+ def dataset_id(self, val):
+ self._dataset_id = val
+
+ @property
+ def tableId(self):
+ return self._table_id
+
+ @tableId.setter
+ def tableId(self, val):
+ self._table_id = val
+
+ @property
+ def table_id(self):
+ return self._table_id
+
+ @table_id.setter
+ def table_id(self, val):
+ self._table_id = val
+
+ @property
+ def dataset_reference(self):
+ return _DatasetReferenceCompat(
+ projectId=self.projectId, datasetId=self.datasetId)
+
+ @property
+ def datasetReference(self):
+ return self.dataset_reference
+
+ def __repr__(self):
+ return (
+ f"TableReference(projectId='{self.projectId}', "
+ f"datasetId='{self.datasetId}', tableId='{self.tableId}')")
+
+ def __eq__(self, other):
+ if other is None:
+ return False
+ if not hasattr(other, "tableId") and not hasattr(other, "table_id"):
+ return NotImplemented
+ other_p = getattr(other, "projectId", None) or getattr(
+ other, "project", None)
+ other_d = getattr(other, "datasetId", None) or getattr(
+ other, "dataset_id", None)
+ other_t = getattr(other, "tableId", None) or getattr(
+ other, "table_id", None)
+ return (self.projectId, self.datasetId,
+ self.tableId) == (other_p, other_d, other_t)
+
+ def __hash__(self):
+ return hash((self.projectId, self.datasetId, self.tableId))
+
+
+class _TableFieldSchemaCompat(object):
+ def __init__(
+ self,
+ name="",
+ type="STRING",
+ mode="NULLABLE",
+ description=None,
+ fields=(),
+ field_type=None,
+ **kwargs):
+ ft = type or field_type or "STRING"
+ self.name = name
+ self.field_type = ft
+ self.mode = mode or "NULLABLE"
+ self.description = description
+ self.fields = list(fields) if fields else []
+
+ @property
+ def type(self):
+ return self.field_type
+
+ @type.setter
+ def type(self, val):
+ self.field_type = val
+
+
+class _TableSchemaCompat(list):
+ def __init__(self, fields=None):
+ if fields:
+ super().__init__(fields)
+ else:
+ super().__init__()
+
+ @property
+ def fields(self):
+ return self
+
+ @fields.setter
+ def fields(self, value):
+ self.clear()
+ if value:
+ self.extend(value)
+
+
+class _TableCellCompat(object):
+ def __init__(self, v=None):
+ self.v = v
+
+
+class _TableRowCompat(object):
+ def __init__(self, f=None):
+ self.f = f or []
+
+
+# Check if an underlying type is available because the internal client silently
+# passes on an import error if any dependencies (mainly apitools) are missing.
+if apitools_bigquery is not None and hasattr(apitools_bigquery,
+ "TableReference"):
+ TableReference = apitools_bigquery.TableReference
+ DatasetReference = apitools_bigquery.DatasetReference
+ TableFieldSchema = apitools_bigquery.TableFieldSchema
+ TableSchema = apitools_bigquery.TableSchema
+ TableRow = apitools_bigquery.TableRow
+ TableCell = apitools_bigquery.TableCell
+ Table = apitools_bigquery.Table
+ Dataset = apitools_bigquery.Dataset
+ Job = apitools_bigquery.Job
+ JobConfiguration = apitools_bigquery.JobConfiguration
+ JobConfigurationLoad = apitools_bigquery.JobConfigurationLoad
+ JobConfigurationQuery = apitools_bigquery.JobConfigurationQuery
+ JobConfigurationExtract = apitools_bigquery.JobConfigurationExtract
+ JobConfigurationTableCopy = apitools_bigquery.JobConfigurationTableCopy
+ JobStatistics = apitools_bigquery.JobStatistics
+ JobStatistics2 = apitools_bigquery.JobStatistics2
+ JobStatistics4 = apitools_bigquery.JobStatistics4
+ ErrorProto = apitools_bigquery.ErrorProto
+else:
+ TableReference = _TableReferenceCompat
+ DatasetReference = _DatasetReferenceCompat
+ TableFieldSchema = _TableFieldSchemaCompat
+ TableSchema = _TableSchemaCompat
+ TableRow = _TableRowCompat
+ TableCell = _TableCellCompat
+ Table = None
+ Dataset = None
+ Job = None
+ JobConfiguration = None
+ JobConfigurationLoad = None
+ JobConfigurationQuery = None
+ JobConfigurationExtract = None
+ JobConfigurationTableCopy = None
+ JobStatistics = None
+ JobStatistics2 = None
+ JobStatistics4 = None
+ ErrorProto = None
+
+
+class JobReference(object):
+ """Compatibility model for BigQuery JobReference.
+
+ Supports both camelCase (jobId, projectId) and snake_case (job_id, project,
project_id)
+ initialization and attribute access.
+ """
+ def __init__(
+ self,
+ jobId=None,
+ projectId=None,
+ location=None,
+ job_id=None,
+ project=None,
+ project_id=None):
+ self.jobId = jobId if jobId is not None else job_id
+ self.projectId = (
+ projectId if projectId is not None else
+ (project if project is not None else project_id))
+ self.location = location
+
+ @property
+ def job_id(self):
+ return self.jobId
+
+ @job_id.setter
+ def job_id(self, val):
+ self.jobId = val
+
+ @property
+ def project(self):
+ return self.projectId
+
+ @project.setter
+ def project(self, val):
+ self.projectId = val
+
+ @property
+ def project_id(self):
+ return self.projectId
+
+ @project_id.setter
+ def project_id(self, val):
+ self.projectId = val
+
+ def __eq__(self, other):
+ if other is None:
+ return False
+ if isinstance(other, JobReference):
+ return (
+ self.jobId == other.jobId and self.projectId == other.projectId and
+ self.location == other.location)
+ if apitools_bigquery and hasattr(apitools_bigquery,
+ "JobReference") and isinstance(
+ other,
apitools_bigquery.JobReference):
+ return (
+ self.jobId == getattr(other, "jobId", None) and
+ self.projectId == getattr(other, "projectId", None) and
+ self.location == getattr(other, "location", None))
+ return NotImplemented
+
+ def __hash__(self):
+ return hash((self.jobId, self.projectId, self.location))
+
+ def __repr__(self):
+ return (
+ f"JobReference(jobId={self.jobId!r}, "
+ f"projectId={self.projectId!r}, "
+ f"location={self.location!r})")
+
+
+_PATCHED_PROTORPCLITE = False
+
+
+def _patch_protorpclite_equality():
+ global _PATCHED_PROTORPCLITE
+ if _PATCHED_PROTORPCLITE:
+ return
+ if _protorpclite_messages is not None and hasattr(_protorpclite_messages,
+ "Message"):
+ _PATCHED_PROTORPCLITE = True
+ _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 _set_table_ref_prop(ref, prop, val):
+ if hasattr(ref, "_properties") and isinstance(ref._properties, dict):
+ ref._properties[prop] = val
+ if prop == "projectId":
+ setattr(ref, "_project", val)
+ elif prop == "datasetId":
+ setattr(ref, "_dataset_id", val)
+ elif prop == "tableId":
+ setattr(ref, "_table_id", val)
+
+
+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
+
+
+class _ClusteringCompat(dict):
+ """Compatibility model for Table clustering.
+
+ Supports dict access, json serialization, and attribute access to fields.
+ Note: Mutations to .fields in-place do not re-sync back to the parent table.
+ """
+ def __init__(self, fields=None):
+ fields_list = list(fields) if fields else []
+ super().__init__(fields=fields_list)
+
+ @property
+ def fields(self):
+ return self.get("fields", [])
+
+ @fields.setter
+ def fields(self, val):
+ self["fields"] = list(val) if val else []
+
+
+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")
+ try:
+ exp_ms = int(tp_exp) if tp_exp is not None else None
+ except (ValueError, TypeError):
+ exp_ms = None
+ if gcp_bigquery is not None and hasattr(gcp_bigquery, "TimePartitioning"):
+ return gcp_bigquery.TimePartitioning(
+ type_=tp_type,
+ field=tp_field,
+ expiration_ms=exp_ms,
+ require_partition_filter=tp_req)
+ 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", "range_")
+ if rp_range is not None and gcp_bigquery is not None and hasattr(
+ gcp_bigquery, "PartitionRange"):
+ if not isinstance(rp_range, getattr(gcp_bigquery, "PartitionRange", ())):
+ start = _extract_field(rp_range, "start")
+ end = _extract_field(rp_range, "end")
+ interval = _extract_field(rp_range, "interval")
+ try:
+ start = int(start) if start is not None else None
+ except (ValueError, TypeError):
+ start = None
+ try:
+ end = int(end) if end is not None else None
+ except (ValueError, TypeError):
+ end = None
+ try:
+ interval = int(interval) if interval is not None else None
+ except (ValueError, TypeError):
+ interval = None
+ rp_range = gcp_bigquery.PartitionRange(
+ start=start, end=end, interval=interval)
+ if gcp_bigquery is not None and hasattr(gcp_bigquery, "RangePartitioning"):
+ return gcp_bigquery.RangePartitioning(field=rp_field, range_=rp_range)
+ return rp
+
+
+def _to_gcp_clustering_fields(clustering):
+ """Extracts clustering field list from list, dict, or apitools Clustering."""
+ 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 isinstance(fields, (list, tuple)):
+ return list(fields)
+ return None
+
+
+# -----------------------------------------------------------------------------
+# Compatibility Monkey-Patching for google.cloud.bigquery Classes
+# -----------------------------------------------------------------------------
+
+_PATCHED_GCP_BIGQUERY = False
+
+
+def _patch_gcp_bigquery():
+ global _PATCHED_GCP_BIGQUERY
+ if not gcp_bigquery or _PATCHED_GCP_BIGQUERY:
+ return
+ _PATCHED_GCP_BIGQUERY = True
+
+ if not hasattr(gcp_bigquery.TableReference, "projectId"):
+ gcp_bigquery.TableReference.projectId = property(
+ lambda self: self.project,
+ lambda self, val: _set_table_ref_prop(self, "projectId", val))
+ gcp_bigquery.TableReference.datasetId = property(
+ lambda self: self.dataset_id,
+ lambda self, val: _set_table_ref_prop(self, "datasetId", val))
+ gcp_bigquery.TableReference.tableId = property(
+ lambda self: self.table_id,
+ lambda self, val: _set_table_ref_prop(self, "tableId", val))
+
+ if not hasattr(gcp_bigquery.DatasetReference, "projectId"):
+ gcp_bigquery.DatasetReference.projectId = property(
+ lambda self: self.project,
+ lambda self, val: setattr(self, "_project", val))
+ gcp_bigquery.DatasetReference.datasetId = property(
+ lambda self: self.dataset_id,
+ lambda self, val: setattr(self, "_dataset_id", val))
+
+ if not hasattr(gcp_bigquery.SchemaField, "type"):
+ gcp_bigquery.SchemaField.type = property(
+ lambda self: self.field_type,
+ lambda self, val: setattr(self, "_field_type", val))
+
+ if not hasattr(gcp_bigquery.Table, "tableReference"):
+ gcp_bigquery.Table.tableReference = property(lambda self: self.reference)
+ gcp_bigquery.Table.numRows = property(lambda self: self.num_rows)
+ gcp_bigquery.Table.numBytes = property(lambda self: self.num_bytes)
+ gcp_bigquery.Table.timePartitioning = property(
+ lambda self: self.time_partitioning, lambda self, val: setattr(
+ self, "time_partitioning", _to_gcp_time_partitioning(val)))
+ gcp_bigquery.Table.rangePartitioning = property(
+ lambda self: self.range_partitioning, lambda self, val: setattr(
+ self, "range_partitioning", _to_gcp_range_partitioning(val)))
+
+ if hasattr(gcp_bigquery.Table, "time_partitioning"):
+ orig_tp_setter = gcp_bigquery.Table.time_partitioning.fset
+ if orig_tp_setter is not None:
+ gcp_bigquery.Table.time_partitioning = property(
+ gcp_bigquery.Table.time_partitioning.fget,
+ lambda self, val: orig_tp_setter(
+ self, _to_gcp_time_partitioning(val)))
+
+ if hasattr(gcp_bigquery.Table, "range_partitioning"):
+ orig_rp_setter = gcp_bigquery.Table.range_partitioning.fset
+ if orig_rp_setter is not None:
+ gcp_bigquery.Table.range_partitioning = property(
+ gcp_bigquery.Table.range_partitioning.fget,
+ lambda self, val: orig_rp_setter(
+ self, _to_gcp_range_partitioning(val)))
+
+ if not hasattr(gcp_bigquery.Table, "clustering"):
+ gcp_bigquery.Table.clustering = property(
+ lambda self: _ClusteringCompat(self.clustering_fields)
+ if self.clustering_fields else None, lambda self, val: setattr(
+ self, "clustering_fields", _to_gcp_clustering_fields(val)))
+
+ if hasattr(gcp_bigquery.Table, "clustering_fields"):
+ orig_cf_setter = gcp_bigquery.Table.clustering_fields.fset
+ if orig_cf_setter is not None:
+ gcp_bigquery.Table.clustering_fields = property(
+ gcp_bigquery.Table.clustering_fields.fget,
+ lambda self, val: orig_cf_setter(
+ self, _to_gcp_clustering_fields(val)))
+
+ if hasattr(gcp_bigquery, "TimePartitioning"):
+ if not hasattr(gcp_bigquery.TimePartitioning, "type"):
+ gcp_bigquery.TimePartitioning.type = property(
+ lambda self: self.type_,
+ lambda self, val: setattr(self, "type_", val))
+ if not hasattr(gcp_bigquery.TimePartitioning, "expirationMs"):
+ gcp_bigquery.TimePartitioning.expirationMs = property(
+ lambda self: self.expiration_ms,
+ lambda self, val: setattr(self, "expiration_ms", val))
+ if not hasattr(gcp_bigquery.TimePartitioning, "requirePartitionFilter"):
+ gcp_bigquery.TimePartitioning.requirePartitionFilter = property(
+ lambda self: self.require_partition_filter,
+ lambda self, val: setattr(self, "require_partition_filter", val))
+
+ if hasattr(gcp_bigquery, "RangePartitioning"):
+ if not hasattr(gcp_bigquery.RangePartitioning, "range"):
+ gcp_bigquery.RangePartitioning.range = property(
+ lambda self: self.range_,
+ lambda self, val: setattr(self, "range_", val))
+
+ if not hasattr(gcp_bigquery.Dataset, "datasetReference"):
+ gcp_bigquery.Dataset.datasetReference = property(
+ lambda self: self.reference)
+ gcp_bigquery.Dataset.defaultTableExpirationMs = property(
+ lambda self: self.default_table_expiration_ms,
+ lambda self, val: setattr(self, "default_table_expiration_ms", val))
+
+ if hasattr(gcp_bigquery, "LoadJobConfig"):
+ if not hasattr(gcp_bigquery.LoadJobConfig, "schemaUpdateOptions"):
+ gcp_bigquery.LoadJobConfig.schemaUpdateOptions = property(
+ lambda self: self.schema_update_options,
+ lambda self, val: setattr(self, "schema_update_options", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "ignoreUnknownValues"):
+ gcp_bigquery.LoadJobConfig.ignoreUnknownValues = property(
+ lambda self: self.ignore_unknown_values,
+ lambda self, val: setattr(self, "ignore_unknown_values", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "maxBadRecords"):
+ gcp_bigquery.LoadJobConfig.maxBadRecords = property(
+ lambda self: self.max_bad_records,
+ lambda self, val: setattr(self, "max_bad_records", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "nullMarker"):
+ gcp_bigquery.LoadJobConfig.nullMarker = property(
+ lambda self: self.null_marker,
+ lambda self, val: setattr(self, "null_marker", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "fieldDelimiter"):
+ gcp_bigquery.LoadJobConfig.fieldDelimiter = property(
+ lambda self: self.field_delimiter,
+ lambda self, val: setattr(self, "field_delimiter", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "skipLeadingRows"):
+ gcp_bigquery.LoadJobConfig.skipLeadingRows = property(
+ lambda self: self.skip_leading_rows,
+ lambda self, val: setattr(self, "skip_leading_rows", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "allowJaggedRows"):
+ gcp_bigquery.LoadJobConfig.allowJaggedRows = property(
+ lambda self: self.allow_jagged_rows,
+ lambda self, val: setattr(self, "allow_jagged_rows", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "allowQuotedNewlines"):
+ gcp_bigquery.LoadJobConfig.allowQuotedNewlines = property(
+ lambda self: self.allow_quoted_newlines,
+ lambda self, val: setattr(self, "allow_quoted_newlines", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "decimalTargetTypes"):
+ gcp_bigquery.LoadJobConfig.decimalTargetTypes = property(
+ lambda self: self.decimal_target_types,
+ lambda self, val: setattr(self, "decimal_target_types", val))
+ if not hasattr(gcp_bigquery.LoadJobConfig, "useAvroLogicalTypes"):
+ gcp_bigquery.LoadJobConfig.useAvroLogicalTypes = property(
+ lambda self: self.use_avro_logical_types,
+ lambda self, val: setattr(self, "use_avro_logical_types", val))
+
+ if hasattr(gcp_bigquery, "QueryJobConfig"):
+ if not hasattr(gcp_bigquery.QueryJobConfig, "schemaUpdateOptions"):
+ gcp_bigquery.QueryJobConfig.schemaUpdateOptions = property(
+ lambda self: self.schema_update_options,
+ lambda self, val: setattr(self, "schema_update_options", val))
+ if not hasattr(gcp_bigquery.QueryJobConfig, "useLegacySql"):
+ gcp_bigquery.QueryJobConfig.useLegacySql = property(
+ lambda self: self.use_legacy_sql,
+ lambda self, val: setattr(self, "use_legacy_sql", val))
+ if not hasattr(gcp_bigquery.QueryJobConfig, "flattenResults"):
+ gcp_bigquery.QueryJobConfig.flattenResults = property(
+ lambda self: self.flatten_results,
+ lambda self, val: setattr(self, "flatten_results", val))
+ if not hasattr(gcp_bigquery.QueryJobConfig, "allowLargeResults"):
+ gcp_bigquery.QueryJobConfig.allowLargeResults = property(
+ lambda self: self.allow_large_results,
+ lambda self, val: setattr(self, "allow_large_results", val))
+ if not hasattr(gcp_bigquery.QueryJobConfig, "maximumBytesBilled"):
+ gcp_bigquery.QueryJobConfig.maximumBytesBilled = property(
+ lambda self: self.maximum_bytes_billed,
+ lambda self, val: setattr(self, "maximum_bytes_billed", val))
+
+ if hasattr(gcp_bigquery, "Table") and hasattr(gcp_bigquery.Table, "labels"):
+ _orig_tbl_labels_setter = gcp_bigquery.Table.labels.fset
+ if _orig_tbl_labels_setter:
+
+ def _safe_tbl_labels_setter(self, value):
+ if value is None:
+ value = {}
+ elif not isinstance(value, dict) and hasattr(value,
+ "additionalProperties"):
+ from apitools.base.py import encoding
+ value = encoding.MessageToDict(value)
+ _orig_tbl_labels_setter(self, value)
+
+ gcp_bigquery.Table.labels = gcp_bigquery.Table.labels.setter(
+ _safe_tbl_labels_setter)
+
+ if hasattr(gcp_bigquery, "Dataset") and hasattr(gcp_bigquery.Dataset,
+ "labels"):
+ _orig_ds_labels_setter = gcp_bigquery.Dataset.labels.fset
+ if _orig_ds_labels_setter:
+
+ def _safe_ds_labels_setter(self, value):
+ if value is None:
+ value = {}
+ elif not isinstance(value, dict) and hasattr(value,
+ "additionalProperties"):
+ from apitools.base.py import encoding
+ value = encoding.MessageToDict(value)
+ _orig_ds_labels_setter(self, value)
+
+ gcp_bigquery.Dataset.labels = gcp_bigquery.Dataset.labels.setter(
+ _safe_ds_labels_setter)
+
+ try:
+ from google.cloud.bigquery.job.base import _JobConfig as _GcpJobConfig
+ if hasattr(_GcpJobConfig, "labels") and hasattr(_GcpJobConfig.labels,
+ "fset"):
+ _orig_job_labels_setter = _GcpJobConfig.labels.fset
+ if _orig_job_labels_setter:
+
+ def _safe_job_labels_setter(self, value):
+ if value is None:
+ value = {}
+ elif not isinstance(value, dict) and hasattr(value,
+ "additionalProperties"):
+ from apitools.base.py import encoding
+ value = encoding.MessageToDict(value)
+ _orig_job_labels_setter(self, value)
+
+ _GcpJobConfig.labels = _GcpJobConfig.labels.setter(
+ _safe_job_labels_setter)
+ except ImportError:
+ pass
+
+ if hasattr(gcp_job,
+ "_AsyncJob") and not hasattr(gcp_job._AsyncJob, "jobReference"):
+ gcp_job._AsyncJob.jobReference = property(
+ lambda self: JobReference(
+ job_id=self.job_id, project=self.project, location=self.location))
+ gcp_job._AsyncJob.status = property(lambda self: _JobStatusCompat(self))
+ gcp_job._AsyncJob.statistics = property(lambda self: _JobStatsCompat(self))
+
+ if not hasattr(gcp_bigquery.Client, "tables"):
+ gcp_bigquery.Client.tables = property(
+ lambda self: _ClientTablesCompat(self))
+ if not hasattr(gcp_bigquery.Client, "datasets"):
+ gcp_bigquery.Client.datasets = property(
+ lambda self: _ClientDatasetsCompat(self))
+ if not hasattr(gcp_bigquery.Client, "jobs"):
+ gcp_bigquery.Client.jobs = property(lambda self: _ClientJobsCompat(self))
+
+
+# -----------------------------------------------------------------------------
+# Compatibility Helpers
+# -----------------------------------------------------------------------------
+
+
+def _to_json_compatible(obj):
+ """Converts an object or nested structure to JSON/API-compatible
dicts/types."""
+ if obj is None:
+ return None
+ if isinstance(obj, (str, int, float, bool)):
+ return obj
+ if isinstance(obj, (list, tuple, set)):
+ return [_to_json_compatible(item) for item in obj]
+ if isinstance(obj, dict):
+ return {k: _to_json_compatible(v) for k, v in obj.items()}
+ if hasattr(obj, "to_api_repr") and callable(obj.to_api_repr):
+ return obj.to_api_repr()
+ if _protorpclite_messages is not None and hasattr(
+ _protorpclite_messages, "Message") and isinstance(
+ obj, _protorpclite_messages.Message):
+ try:
+ from apitools.base.py import encoding
+ return encoding.MessageToDict(obj)
+ except Exception:
+ pass
+ return obj
+
+
+def _extract_dict_labels(labels):
+ """Converts labels to a non-empty dictionary or returns None."""
+ if not labels:
+ return None
+ labels = _to_json_compatible(labels)
+ if isinstance(labels, dict) and labels:
+ return labels
+ return None
+
+
+def _to_gcp_table_ref(table_ref, default_project=None):
+ """Converts a TableReference or string into a
google.cloud.bigquery.TableReference."""
+ if table_ref is None:
+ return None
+ if gcp_bigquery is not None and isinstance(
+ table_ref, getattr(gcp_bigquery, "TableReference", ())):
+ return table_ref
+ if isinstance(table_ref, str):
+ from apache_beam.io.gcp.bigquery_tools import parse_table_reference
+ table_ref = parse_table_reference(table_ref, project=default_project)
+ proj = getattr(table_ref, "projectId", None) or getattr(
+ table_ref, "project", None) or getattr(
+ table_ref, "project_id", None) or default_project or "default"
+ dataset_id = getattr(table_ref, "datasetId", None) or getattr(
+ table_ref, "dataset_id", None) or getattr(table_ref, "dataset", None)
+ table_id = getattr(table_ref, "tableId", None) or getattr(
+ table_ref, "table_id", None) or getattr(table_ref, "table", None)
+ if dataset_id and table_id:
+ if gcp_bigquery is not None and hasattr(
+ gcp_bigquery, "TableReference") and hasattr(gcp_bigquery,
+ "DatasetReference"):
+ return gcp_bigquery.TableReference(
+ gcp_bigquery.DatasetReference(proj, dataset_id), table_id)
+ return _TableReferenceCompat(
+ projectId=proj, datasetId=dataset_id, tableId=table_id)
+ return table_ref
+
+
+def _to_gcp_dataset_ref(dataset_ref, project=None):
+ """Converts a DatasetReference or string into a
google.cloud.bigquery.DatasetReference."""
+ if dataset_ref is None:
+ return None
+ if gcp_bigquery is not None and isinstance(
+ dataset_ref, getattr(gcp_bigquery, "DatasetReference", ())):
+ return dataset_ref
+ if isinstance(dataset_ref, str):
+ 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 _DatasetReferenceCompat(projectId=proj, datasetId=ds_id)
+ proj = project or "default"
+ if gcp_bigquery is not None and hasattr(gcp_bigquery, "DatasetReference"):
+ return gcp_bigquery.DatasetReference(proj, dataset_ref)
+ return _DatasetReferenceCompat(projectId=proj, datasetId=dataset_ref)
+ if hasattr(dataset_ref, "projectId") or hasattr(dataset_ref, "project"):
+ proj = getattr(dataset_ref, "projectId", None) or getattr(
+ dataset_ref, "project", None) or getattr(
+ dataset_ref, "project_id", None) or project or "default"
+ ds_id = getattr(dataset_ref, "datasetId", None) or getattr(
+ dataset_ref, "dataset_id", None)
+ if gcp_bigquery is not None and hasattr(gcp_bigquery, "DatasetReference"):
+ return gcp_bigquery.DatasetReference(proj, ds_id)
+ return _DatasetReferenceCompat(projectId=proj, datasetId=ds_id)
+ return dataset_ref
+
+
+def _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"):
+ from apache_beam.io.gcp.bigquery_tools import get_dict_table_schema
+ 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):
+ from apache_beam.io.gcp.bigquery_tools import get_dict_table_schema
+ return _to_gcp_schema(get_dict_table_schema(schema))
+ return schema
+
+
+def _to_table_schema(schema):
+ """Converts a list of google.cloud.bigquery.SchemaField, dict, or
TableSchema into a TableSchema."""
+ if schema is None:
+ return TableSchema()
+ if isinstance(schema, TableSchema):
+ return schema
+ if isinstance(schema, dict):
+ return _to_table_schema(schema.get("fields", []))
+ if hasattr(schema, "fields") and not isinstance(schema, (list, tuple)):
+ return _to_table_schema(schema.fields)
+
+ def _to_field_schema(f):
+ if isinstance(f, TableFieldSchema):
+ return f
+ if isinstance(f, dict):
+ f_dict = f
+ elif hasattr(f, "to_api_repr"):
+ f_dict = f.to_api_repr()
+ else:
+ f_dict = None
+
+ if f_dict is not None:
+ name = f_dict.get("name", "")
+ field_type = f_dict.get("type") or f_dict.get("type_") or "STRING"
+ mode = f_dict.get("mode", "NULLABLE")
+ description = f_dict.get("description", None)
+ sub_fields = [_to_field_schema(sf) for sf in f_dict.get("fields", [])]
+ return TableFieldSchema(
+ name=name,
+ type=field_type,
+ mode=mode,
+ description=description,
+ fields=sub_fields)
+
+ name = getattr(f, "name", "")
+ field_type = getattr(f, "field_type", None) or getattr(
+ f, "type", None) or "STRING"
+ mode = getattr(f, "mode", "NULLABLE")
+ description = getattr(f, "description", None)
+ sub = getattr(f, "fields", ())
+ sub_fields = [_to_field_schema(sf) for sf in sub] if sub else ()
+ return TableFieldSchema(
+ name=name,
+ type=field_type,
+ mode=mode,
+ description=description,
+ fields=sub_fields)
+
+ if isinstance(schema, (list, tuple)):
+ return TableSchema(fields=[_to_field_schema(f) for f in schema])
+ return TableSchema()
+
+
+# -----------------------------------------------------------------------------
+# Emulated Client Compatibility Interfaces
+# -----------------------------------------------------------------------------
+
+
+class _JobStatusCompat:
+ def __init__(self, job):
+ self._job = job
+
+ @property
+ def state(self):
+ return self._job.state
+
+ @property
+ def errorResult(self):
+ return self._job.error_result
+
+ @property
+ def errors(self):
+ return self._job.errors
+
+
+class _JobStatsCompat:
+ def __init__(self, job):
+ self._job = job
+
+ @property
+ def query(self):
+ return self
+
+ @property
+ def totalBytesBilled(self):
+ return getattr(self._job, "total_bytes_billed", None)
+
+ @property
+ def totalBytesProcessed(self):
+ return getattr(self._job, "total_bytes_processed", None)
+
+ @property
+ def referencedTables(self):
+ tables = getattr(self._job, "referenced_tables", None)
+ if tables is not None:
+ return [
+ TableReference(
+ projectId=t.project, datasetId=t.dataset_id, tableId=t.table_id)
+ for t in tables
+ ]
+ return None
+
+
+class _ClientTablesCompat:
+ def __init__(self, client):
+ self._client = client
+
+ def Get(self, request):
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ tbl_id = getattr(request, "tableId", None)
+ if ds_id and tbl_id:
+ table_ref = gcp_bigquery.TableReference(
+ gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default",
+ ds_id),
+ tbl_id)
+ else:
+ t_ref = getattr(request, "tableReference", None) or getattr(
+ request, "tableId", None) or request
+ table_ref = _to_gcp_table_ref(
+ t_ref, default_project=proj or getattr(self._client, "project",
None))
+ return self._client.get_table(table_ref)
+
+ def Insert(self, request):
+ table = getattr(request, "table", None)
+ if table is not None:
+ t_ref = getattr(table, "tableReference", None)
+ proj = getattr(t_ref, "projectId", None) or getattr(
+ request, "projectId", None)
+ ds_id = getattr(t_ref, "datasetId", None) or getattr(
+ request, "datasetId", None)
+ tbl_id = getattr(t_ref, "tableId", None)
+ schema = getattr(table, "schema", None)
+ else:
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ tbl_id = getattr(request, "tableId", None)
+ schema = getattr(request, "schema", None)
+ gcp_tbl_ref = gcp_bigquery.TableReference(
+ gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default",
ds_id),
+ tbl_id)
+ gcp_table = gcp_bigquery.Table(gcp_tbl_ref, schema=_to_gcp_schema(schema))
+ if table is not None:
+ tp = _extract_field(table, "timePartitioning", "time_partitioning")
+ if tp is not None:
+ gcp_table.time_partitioning = _to_gcp_time_partitioning(tp)
+ rp = _extract_field(table, "rangePartitioning", "range_partitioning")
+ if rp is not None:
+ gcp_table.range_partitioning = _to_gcp_range_partitioning(rp)
+ clustering = _extract_field(table, "clustering", "clustering_fields")
+ if clustering is not None:
+ cl_fields = _to_gcp_clustering_fields(clustering)
+ if cl_fields:
+ gcp_table.clustering_fields = cl_fields
+ desc = _extract_field(table, "description")
+ if desc is not None:
+ gcp_table.description = desc
+ fname = _extract_field(table, "friendlyName", "friendly_name")
+ if fname is not None:
+ gcp_table.friendly_name = fname
+ dict_labels = _extract_dict_labels(getattr(table, "labels", None))
+ if dict_labels:
+ gcp_table.labels = dict_labels
+ enc = _extract_field(
+ table, "encryptionConfiguration", "encryption_configuration")
+ kms = _extract_field(
+ enc, "kmsKeyName", "kms_key_name") if enc is not None else None
+ if kms:
+ gcp_table.encryption_configuration = (
+ gcp_bigquery.EncryptionConfiguration(kms_key_name=kms))
+ return self._client.create_table(gcp_table, exists_ok=True)
+
+ def Delete(self, request):
+ t_ref = getattr(request, "tableReference", None)
+ proj = getattr(t_ref, "projectId", None) or getattr(
+ request, "projectId", None)
+ ds_id = getattr(t_ref, "datasetId", None) or getattr(
+ request, "datasetId", None)
+ tbl_id = getattr(t_ref, "tableId", None) or getattr(
+ request, "tableId", None)
+ gcp_tbl_ref = gcp_bigquery.TableReference(
+ gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default",
ds_id),
+ tbl_id)
+ return self._client.delete_table(gcp_tbl_ref, not_found_ok=True)
+
+ def List(self, request):
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ ds_ref = gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default", ds_id)
+ return self._client.list_tables(ds_ref)
+
+ def Patch(self, request):
+ table = getattr(request, "table", None)
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ tbl_id = getattr(request, "tableId", None)
+ gcp_tbl_ref = gcp_bigquery.TableReference(
+ gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default",
ds_id),
+ tbl_id)
+ gcp_table = gcp_bigquery.Table(gcp_tbl_ref)
+ if table and getattr(table, "schema", None):
+ gcp_table.schema = _to_gcp_schema(table.schema)
+ return self._client.update_table(gcp_table, ["schema"])
+
+ def Update(self, request):
+ return self.Patch(request)
+
+
+class _ClientDatasetsCompat:
+ def __init__(self, client):
+ self._client = client
+
+ def Get(self, request):
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ ds_ref = gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default", ds_id)
+ return self._client.get_dataset(ds_ref)
+
+ def Insert(self, request):
+ dataset = getattr(request, "dataset", None)
+ ds_ref_raw = getattr(dataset, "datasetReference", None) if dataset else
None
+ proj = getattr(ds_ref_raw, "projectId", None) or getattr(
+ request, "projectId", None)
+ ds_id = getattr(ds_ref_raw, "datasetId", None) or getattr(
+ request, "datasetId", None)
+ ds_ref = gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default", ds_id)
+ gcp_ds = gcp_bigquery.Dataset(ds_ref)
+ if dataset:
+ if getattr(dataset, "location", None):
+ gcp_ds.location = dataset.location
+ if getattr(dataset, "defaultTableExpirationMs", None):
+ gcp_ds.default_table_expiration_ms = dataset.defaultTableExpirationMs
+ return self._client.create_dataset(gcp_ds, exists_ok=True)
+
+ def Delete(self, request):
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ delete_contents = getattr(request, "deleteContents", True)
+ ds_ref = gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default", ds_id)
+ return self._client.delete_dataset(
+ ds_ref, delete_contents=delete_contents, not_found_ok=True)
+
+ def List(self, request):
+ proj = getattr(request, "projectId", None) or getattr(
+ self._client, "project", None)
+ return self._client.list_datasets(project=proj)
+
+ def Patch(self, request):
+ dataset = getattr(request, "dataset", None)
+ proj = getattr(request, "projectId", None)
+ ds_id = getattr(request, "datasetId", None)
+ ds_ref = gcp_bigquery.DatasetReference(
+ proj or getattr(self._client, "project", None) or "default", ds_id)
+ gcp_ds = gcp_bigquery.Dataset(ds_ref)
+ fields_to_update = []
+ if dataset:
+ if getattr(dataset, "defaultTableExpirationMs", None):
+ gcp_ds.default_table_expiration_ms = dataset.defaultTableExpirationMs
+ fields_to_update.append("default_table_expiration_ms")
+ return self._client.update_dataset(gcp_ds, fields_to_update)
+
+ def Update(self, request):
+ return self.Patch(request)
+
+
+class _ClientJobsCompat:
+ def __init__(self, client):
+ self._client = client
+
+ def Get(self, request):
+ proj = getattr(request, "projectId", None)
+ job_id = getattr(request, "jobId", None)
+ loc = getattr(request, "location", None)
+ return self._client.get_job(job_id, project=proj, location=loc)
+
+ def GetQueryResults(self, request):
+ proj = getattr(request, "projectId", None)
+ job_id = getattr(request, "jobId", None)
+ loc = getattr(request, "location", None)
+ page_token = getattr(request, "pageToken", None)
+ max_results = getattr(request, "maxResults", None)
+ job = self._client.get_job(job_id, project=proj, location=loc)
+ if page_token is not None:
+ return self._client.list_rows(
+ job, page_token=page_token, max_results=max_results)
+ return job.result(max_results=max_results)
+
+ def Insert(self, request, upload=None):
+ job_obj = getattr(request, "job", None)
+ job_ref = (
+ getattr(job_obj, "jobReference", None) if job_obj else getattr(
+ request, "jobReference", None))
+ job_id = getattr(job_ref, "jobId", None) or getattr(job_ref, "job_id",
None)
+ proj = (
+ getattr(request, "projectId", None) or
+ getattr(job_ref, "projectId", None) or
+ getattr(job_ref, "project", None))
+ config = getattr(job_obj, "configuration", None) if job_obj else None
+ dict_labels = _extract_dict_labels(
+ getattr(config, "labels", None)) if config else None
+ if config and getattr(config, "query", None):
+ q = config.query
+ dest = None
+ if getattr(q, "destinationTable", None):
+ dest = _to_gcp_table_ref(q.destinationTable, default_project=proj)
+ job_config = gcp_bigquery.QueryJobConfig(
+ dry_run=getattr(q, "dryRun", False),
+ use_legacy_sql=getattr(q, "useLegacySql", False) if getattr(
+ q, "useLegacySql", None) is not None else False,
+ flatten_results=getattr(q, "flattenResults", None),
+ priority=getattr(q, "priority", "INTERACTIVE"),
+ destination=dest,
+ )
+ if dict_labels:
+ job_config.labels = dict_labels
+ kms = getattr(
+ getattr(q, "destinationEncryptionConfiguration", None),
+ "kmsKeyName",
+ None)
+ if kms:
+ job_config.destination_encryption_configuration = (
+ gcp_bigquery.EncryptionConfiguration(kms_key_name=kms))
+ return self._client.query(
+ q.query,
+ job_config=job_config,
+ job_id=job_id,
+ project=proj,
+ job_retry=None,
+ )
+ elif config and getattr(config, "load", None):
+ ld = config.load
+ dest = _to_gcp_table_ref(
+ getattr(ld, "destinationTable", None), default_project=proj)
+ uris = list(getattr(ld, "sourceUris", []))
+ if uris:
+ job_config = gcp_bigquery.LoadJobConfig()
+ if dict_labels:
+ job_config.labels = dict_labels
+ if getattr(ld, "schema", None):
+ job_config.schema = _to_gcp_schema(ld.schema)
+ if getattr(ld, "writeDisposition", None):
+ job_config.write_disposition = ld.writeDisposition
+ if getattr(ld, "createDisposition", None):
+ job_config.create_disposition = ld.createDisposition
+ return self._client.load_table_from_uri(
+ uris, dest, job_config=job_config, job_id=job_id, project=proj)
+ elif config and getattr(config, "copy", None):
+ cp = config.copy
+ sources = [
+ _to_gcp_table_ref(s, default_project=proj)
+ for s in getattr(cp, "sourceTables", [])
+ ]
+ dest = _to_gcp_table_ref(
+ getattr(cp, "destinationTable", None), default_project=proj)
+ job_config = gcp_bigquery.CopyJobConfig()
+ if dict_labels:
+ job_config.labels = dict_labels
+ if getattr(cp, "writeDisposition", None):
+ job_config.write_disposition = cp.writeDisposition
+ if getattr(cp, "createDisposition", None):
+ job_config.create_disposition = cp.createDisposition
+ return self._client.copy_table(
+ sources, dest, job_config=job_config, job_id=job_id, project=proj)
+ elif config and getattr(config, "extract", None):
+ ex = config.extract
+ src = _to_gcp_table_ref(
+ getattr(ex, "sourceTable", None), default_project=proj)
+ uris = list(getattr(ex, "destinationUris", []))
+ job_config = gcp_bigquery.ExtractJobConfig()
+ if dict_labels:
+ job_config.labels = dict_labels
+ if getattr(ex, "destinationFormat", None):
+ job_config.destination_format = ex.destinationFormat
+ return self._client.extract_table(
+ src, uris, job_config=job_config, job_id=job_id, project=proj)
+
+ return self._client.get_job(job_id, project=proj)
+
+
+# Automatically execute compatibility setup and monkey-patching upon import
+_patch_protorpclite_equality()
+_patch_gcp_bigquery()
diff --git a/sdks/python/apache_beam/io/gcp/bigquery_compat_test.py
b/sdks/python/apache_beam/io/gcp/bigquery_compat_test.py
new file mode 100644
index 00000000000..ca77abc8dcb
--- /dev/null
+++ b/sdks/python/apache_beam/io/gcp/bigquery_compat_test.py
@@ -0,0 +1,796 @@
+#
+# Licensed to the Apache Software Foundation (ASF) under one or more
+# contributor license agreements. See the NOTICE file distributed with
+# this work for additional information regarding copyright ownership.
+# The ASF licenses this file to You under the Apache License, Version 2.0
+# (the "License"); you may not use this file except in compliance with
+# the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Unit tests for BigQuery compatibility shims and legacy client emulation.
+
+These tests verify the compatibility models, monkey-patches, and emulated
+interfaces in ``apache_beam.io.gcp.bigquery_compat``.
+
+.. note::
+ This test suite is intended to be completely removed along with
+ ``bigquery_compat.py`` in a future Beam release when apitools is removed.
+"""
+
+# pytype: skip-file
+
+import importlib
+import logging
+import os
+import pickle
+import unittest
+import warnings
+from unittest import mock
+
+import apache_beam as beam
+from apache_beam.utils.annotations import BeamDeprecationWarning
+
+# Opt in to BigQuery compatibility shims so the compatibility test suite can
run.
+os.environ["BEAM_USE_BIGQUERY_COMPAT_SHIMS"] = "1"
+
+try:
+ from apache_beam.io.gcp import bigquery_compat
+ from apache_beam.io.gcp.internal.clients import bigquery as apitools_bigquery
+except ImportError:
+ bigquery_compat = None
+ apitools_bigquery = None
+
+try:
+ from apitools.base.py.exceptions import HttpError
+except ImportError:
+ HttpError = None
+
+try:
+ from google.cloud import bigquery as gcp_bigquery
+except ImportError:
+ gcp_bigquery = None
+
+
[email protected](bigquery_compat is None, 'GCP dependencies are not installed')
+class TestBigQueryCompatOptIn(unittest.TestCase):
+ """Tests for explicit opt-in environment variable check and warnings."""
+ def test_check_compat_opt_in_disabled(self):
+ # Disabled when unset or empty
+ with mock.patch.dict(os.environ, {"BEAM_USE_BIGQUERY_COMPAT_SHIMS": ""}):
+ with self.assertRaises(ImportError) as ctx:
+ bigquery_compat._check_compat_opt_in()
+ self.assertIn("BEAM_USE_BIGQUERY_COMPAT_SHIMS", str(ctx.exception))
+ self.assertIn("disabled by default", str(ctx.exception))
+
+ # Disabled when 0 or false or no
+ for disabled_val in ("0", "false", "False", "no", "NO"):
+ with mock.patch.dict(os.environ,
+ {"BEAM_USE_BIGQUERY_COMPAT_SHIMS": disabled_val}):
+ with self.assertRaises(ImportError):
+ bigquery_compat._check_compat_opt_in()
+
+ def test_check_compat_opt_in_enabled_warning(self):
+ for enabled_val in ("1", "true", "True", "yes", "enabled", "on"):
+ with mock.patch.dict(os.environ,
+ {"BEAM_USE_BIGQUERY_COMPAT_SHIMS": enabled_val}):
+ with warnings.catch_warnings(record=True) as recorded:
+ warnings.simplefilter("always")
+ bigquery_compat._check_compat_opt_in()
+ self.assertTrue(
+ any(
+ issubclass(w.category, BeamDeprecationWarning)
+ for w in recorded),
+ f"Expected BeamDeprecationWarning for {enabled_val}")
+ self.assertTrue(
+ any("deprecated" in str(w.message).lower() for w in recorded),
+ f"Expected deprecation message for {enabled_val}")
+
+ def test_reload_without_env_var_raises_import_error(self):
+ try:
+ with mock.patch.dict(os.environ, {"BEAM_USE_BIGQUERY_COMPAT_SHIMS": ""}):
+ with self.assertRaises(ImportError):
+ importlib.reload(bigquery_compat)
+ finally:
+ with mock.patch.dict(os.environ, {"BEAM_USE_BIGQUERY_COMPAT_SHIMS":
"1"}):
+ importlib.reload(bigquery_compat)
+
+
[email protected](bigquery_compat is None, 'GCP dependencies are not installed')
+class TestJobReferenceCompatibility(unittest.TestCase):
+ def test_init_camel_case(self):
+ ref = bigquery_compat.JobReference(
+ jobId="test-job", projectId="test-proj", location="US")
+ self.assertEqual(ref.jobId, "test-job")
+ self.assertEqual(ref.job_id, "test-job")
+ self.assertEqual(ref.projectId, "test-proj")
+ self.assertEqual(ref.project, "test-proj")
+ self.assertEqual(ref.project_id, "test-proj")
+ self.assertEqual(ref.location, "US")
+
+ def test_init_snake_case(self):
+ ref = bigquery_compat.JobReference(
+ job_id="test-job", project="test-proj", location="EU")
+ self.assertEqual(ref.jobId, "test-job")
+ self.assertEqual(ref.job_id, "test-job")
+ self.assertEqual(ref.projectId, "test-proj")
+ self.assertEqual(ref.project, "test-proj")
+ self.assertEqual(ref.project_id, "test-proj")
+ self.assertEqual(ref.location, "EU")
+
+ def test_init_with_project_id(self):
+ ref = bigquery_compat.JobReference(
+ job_id="test-job", project_id="test-proj", location="EU")
+ self.assertEqual(ref.jobId, "test-job")
+ self.assertEqual(ref.job_id, "test-job")
+ self.assertEqual(ref.projectId, "test-proj")
+ self.assertEqual(ref.project, "test-proj")
+ self.assertEqual(ref.project_id, "test-proj")
+
+ def test_attribute_setters(self):
+ ref = bigquery_compat.JobReference()
+ ref.job_id = "j1"
+ ref.project = "p1"
+ self.assertEqual(ref.jobId, "j1")
+ self.assertEqual(ref.job_id, "j1")
+ self.assertEqual(ref.projectId, "p1")
+ self.assertEqual(ref.project, "p1")
+ self.assertEqual(ref.project_id, "p1")
+ ref.project_id = "p2"
+ self.assertEqual(ref.projectId, "p2")
+ self.assertEqual(ref.project, "p2")
+ self.assertEqual(ref.project_id, "p2")
+
+ def test_equality_with_custom_and_apitools(self):
+ ref1 = bigquery_compat.JobReference(
+ job_id="j1", project="p1", location="US")
+ ref2 = bigquery_compat.JobReference(
+ jobId="j1", projectId="p1", location="US")
+ ref3 = bigquery_compat.JobReference(
+ jobId="j2", projectId="p1", location="US")
+ self.assertEqual(ref1, ref2)
+ self.assertNotEqual(ref1, ref3)
+ self.assertEqual(hash(ref1), hash(ref2))
+
+ if apitools_bigquery is not None and hasattr(apitools_bigquery,
+ "JobReference"):
+ ap_ref = apitools_bigquery.JobReference(
+ jobId="j1", projectId="p1", location="US")
+ self.assertEqual(ref1, ap_ref)
+ self.assertEqual(ap_ref, ref1)
+
+ def test_equality_type_safety(self):
+ empty_ref = bigquery_compat.JobReference()
+ self.assertNotEqual(empty_ref, None)
+ self.assertNotEqual(empty_ref, 123)
+ self.assertNotEqual(empty_ref, "")
+ self.assertNotEqual(empty_ref, {})
+
+ ref = bigquery_compat.JobReference(project="p1")
+ table_ref = bigquery_compat.TableReference(
+ projectId="p1", datasetId="d1", tableId="t1")
+ self.assertNotEqual(ref, table_ref)
+ self.assertNotEqual(empty_ref, table_ref)
+
+ def test_pickle_and_coder_roundtrip(self):
+ ref = bigquery_compat.JobReference(job_id="j1", project="p1",
location="US")
+ pickled = pickle.dumps(ref)
+ unpickled = pickle.loads(pickled)
+ self.assertEqual(ref, unpickled)
+ self.assertEqual(unpickled.jobId, "j1")
+ self.assertEqual(unpickled.projectId, "p1")
+ self.assertEqual(unpickled.location, "US")
+
+ coder = beam.coders.FastPrimitivesCoder()
+ encoded = coder.encode(ref)
+ decoded = coder.decode(encoded)
+ self.assertEqual(ref, decoded)
+
+
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestTableAndDatasetReferenceCompatibility(unittest.TestCase):
+ @unittest.skipIf(gcp_bigquery is None, "google-cloud-bigquery not installed")
+ def test_table_reference_property_mutability(self):
+ ds = gcp_bigquery.DatasetReference("p1", "d1")
+ table = gcp_bigquery.TableReference(ds, "t1")
+ table.tableId = "t2"
+ self.assertEqual(table.tableId, "t2")
+ self.assertEqual(table.table_id, "t2")
+ table.datasetId = "d2"
+ self.assertEqual(table.datasetId, "d2")
+ self.assertEqual(table.dataset_id, "d2")
+ table.projectId = "p2"
+ self.assertEqual(table.projectId, "p2")
+ self.assertEqual(table.project, "p2")
+
+ def test_table_reference_compat_model(self):
+ table = bigquery_compat._TableReferenceCompat(
+ projectId="p1", datasetId="d1", tableId="t1")
+ self.assertEqual(table.projectId, "p1")
+ self.assertEqual(table.project, "p1")
+ self.assertEqual(table.project_id, "p1")
+ self.assertEqual(table.datasetId, "d1")
+ self.assertEqual(table.dataset_id, "d1")
+ self.assertEqual(table.tableId, "t1")
+ self.assertEqual(table.table_id, "t1")
+ table.tableId = "t2"
+ self.assertEqual(table.tableId, "t2")
+ self.assertEqual(table.table_id, "t2")
+ table.datasetId = "d2"
+ self.assertEqual(table.datasetId, "d2")
+ self.assertEqual(table.dataset_id, "d2")
+ table.projectId = "p2"
+ self.assertEqual(table.projectId, "p2")
+ self.assertEqual(table.project, "p2")
+ self.assertEqual(table.project_id, "p2")
+
+ def test_dataset_reference_compat_model(self):
+ ds = bigquery_compat._DatasetReferenceCompat(projectId="p1",
datasetId="d1")
+ self.assertEqual(ds.projectId, "p1")
+ self.assertEqual(ds.project, "p1")
+ self.assertEqual(ds.project_id, "p1")
+ self.assertEqual(ds.datasetId, "d1")
+ self.assertEqual(ds.dataset_id, "d1")
+ ds.projectId = "p2"
+ self.assertEqual(ds.projectId, "p2")
+ self.assertEqual(ds.project, "p2")
+ ds.datasetId = "d2"
+ self.assertEqual(ds.datasetId, "d2")
+ self.assertEqual(ds.dataset_id, "d2")
+
+ def test_to_gcp_dataset_ref_colon_format(self):
+ ds_ref = bigquery_compat._to_gcp_dataset_ref("my-project:my_dataset")
+ self.assertEqual(ds_ref.project, "my-project")
+ self.assertEqual(ds_ref.dataset_id, "my_dataset")
+ self.assertEqual(ds_ref.projectId, "my-project")
+ self.assertEqual(ds_ref.datasetId, "my_dataset")
+
+ def test_to_gcp_dataset_ref_domain_scoped(self):
+ ds_ref1 = bigquery_compat._to_gcp_dataset_ref(
+ "google.com:clouddfe:my_dataset")
+ self.assertEqual(ds_ref1.project, "google.com:clouddfe")
+ self.assertEqual(ds_ref1.dataset_id, "my_dataset")
+
+ ds_ref2 = bigquery_compat._to_gcp_dataset_ref(
+ "google.com:clouddfe.my_dataset")
+ self.assertEqual(ds_ref2.project, "google.com:clouddfe")
+ self.assertEqual(ds_ref2.dataset_id, "my_dataset")
+
+ def test_table_reference_from_string_default_project(self):
+ t_ref = bigquery_compat._TableReferenceCompat.from_string(
+ "my_ds.my_tbl", default_project="default-proj")
+ self.assertEqual(t_ref.projectId, "default-proj")
+ self.assertEqual(t_ref.datasetId, "my_ds")
+ self.assertEqual(t_ref.tableId, "my_tbl")
+
+ def test_dataset_reference_from_string_domain_scoped(self):
+ ds_ref = bigquery_compat._DatasetReferenceCompat.from_string(
+ "google.com:clouddfe:my_dataset")
+ self.assertEqual(ds_ref.projectId, "google.com:clouddfe")
+ self.assertEqual(ds_ref.datasetId, "my_dataset")
+
+ def test_table_field_schema_compat(self):
+ f = bigquery_compat._TableFieldSchemaCompat(
+ name="age", type="INTEGER", mode="REQUIRED")
+ self.assertEqual(f.name, "age")
+ self.assertEqual(f.type, "INTEGER")
+ self.assertEqual(f.mode, "REQUIRED")
+
+ def test_table_schema_compat(self):
+ f1 = bigquery_compat._TableFieldSchemaCompat(name="id", type="INTEGER")
+ f2 = bigquery_compat._TableFieldSchemaCompat(name="val", type="STRING")
+ s = bigquery_compat._TableSchemaCompat([f1, f2])
+ self.assertEqual(len(s.fields), 2)
+ self.assertEqual(s.fields[0].name, "id")
+ self.assertEqual(s.fields[1].name, "val")
+
+
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestJobConfigCompatibility(unittest.TestCase):
+ def test_load_job_config_camel_case_properties(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ config = gcp_bigquery.LoadJobConfig(
+ schemaUpdateOptions=["ALLOW_FIELD_ADDITION", "ALLOW_FIELD_RELAXATION"],
+ ignoreUnknownValues=True,
+ maxBadRecords=10,
+ nullMarker="\\N",
+ fieldDelimiter="\t",
+ skipLeadingRows=1,
+ allowJaggedRows=True,
+ allowQuotedNewlines=True,
+ decimalTargetTypes=["NUMERIC"],
+ useAvroLogicalTypes=True,
+ )
+ self.assertEqual(
+ config.schemaUpdateOptions,
+ ["ALLOW_FIELD_ADDITION", "ALLOW_FIELD_RELAXATION"])
+ self.assertEqual(
+ config.schema_update_options,
+ ["ALLOW_FIELD_ADDITION", "ALLOW_FIELD_RELAXATION"])
+ self.assertTrue(config.ignoreUnknownValues)
+ self.assertTrue(config.ignore_unknown_values)
+ self.assertEqual(config.maxBadRecords, 10)
+ self.assertEqual(config.max_bad_records, 10)
+ self.assertEqual(config.nullMarker, "\\N")
+ self.assertEqual(config.null_marker, "\\N")
+ self.assertEqual(config.fieldDelimiter, "\t")
+ self.assertEqual(config.field_delimiter, "\t")
+ self.assertEqual(config.skipLeadingRows, 1)
+ self.assertEqual(config.skip_leading_rows, 1)
+ self.assertTrue(config.allowJaggedRows)
+ self.assertTrue(config.allow_jagged_rows)
+ self.assertTrue(config.allowQuotedNewlines)
+ self.assertTrue(config.allow_quoted_newlines)
+ self.assertEqual(set(config.decimalTargetTypes), {"NUMERIC"})
+ self.assertEqual(set(config.decimal_target_types), {"NUMERIC"})
+ self.assertTrue(config.useAvroLogicalTypes)
+ self.assertTrue(config.use_avro_logical_types)
+
+ def test_query_job_config_camel_case_properties(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ config = gcp_bigquery.QueryJobConfig(
+ schemaUpdateOptions=["ALLOW_FIELD_ADDITION"],
+ useLegacySql=True,
+ flattenResults=False,
+ allowLargeResults=True,
+ maximumBytesBilled=1000000,
+ )
+ self.assertEqual(config.schemaUpdateOptions, ["ALLOW_FIELD_ADDITION"])
+ self.assertEqual(config.schema_update_options, ["ALLOW_FIELD_ADDITION"])
+ self.assertTrue(config.useLegacySql)
+ self.assertTrue(config.use_legacy_sql)
+ self.assertFalse(config.flattenResults)
+ self.assertFalse(config.flatten_results)
+ self.assertTrue(config.allowLargeResults)
+ self.assertTrue(config.allow_large_results)
+ self.assertEqual(config.maximumBytesBilled, 1000000)
+ self.assertEqual(config.maximum_bytes_billed, 1000000)
+
+ def test_insert_load_job_with_none_labels(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.MagicMock(spec=gcp_bigquery.Client)
+ mock_job = mock.MagicMock()
+ mock_job.job_id = "load_job_id"
+ mock_job.project = "test-project"
+ mock_job.location = "US"
+ client.load_table_from_uri.return_value = mock_job
+
+ jobs_compat = bigquery_compat._ClientJobsCompat(client)
+ request = mock.MagicMock()
+ request.projectId = "test-project"
+ request.job = mock.MagicMock()
+ request.job.jobReference =
bigquery_compat.JobReference(jobId="load_job_id")
+ request.job.configuration = mock.MagicMock(spec=["load", "labels"])
+ request.job.configuration.load.destinationTable = (
+ bigquery_compat.TableReference(
+ projectId="test-project",
+ datasetId="test_dataset",
+ tableId="test_table"))
+ request.job.configuration.load.sourceUris = ["gs://test-bucket/test.csv"]
+ request.job.configuration.load.schema = None
+ request.job.configuration.load.writeDisposition = None
+ request.job.configuration.load.createDisposition = None
+ request.job.configuration.labels = None
+
+ jobs_compat.Insert(request)
+ client.load_table_from_uri.assert_called_once()
+ called_config = client.load_table_from_uri.call_args.kwargs["job_config"]
+ self.assertEqual(called_config.labels, {})
+
+ def test_insert_copy_job_with_none_labels(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.MagicMock(spec=gcp_bigquery.Client)
+ mock_job = mock.MagicMock()
+ mock_job.job_id = "copy_job_id"
+ mock_job.project = "test-project"
+ mock_job.location = "US"
+ client.copy_table.return_value = mock_job
+
+ jobs_compat = bigquery_compat._ClientJobsCompat(client)
+ request = mock.MagicMock()
+ request.projectId = "test-project"
+ request.job = mock.MagicMock()
+ request.job.jobReference =
bigquery_compat.JobReference(jobId="copy_job_id")
+ request.job.configuration = mock.MagicMock(spec=["copy", "labels"])
+ request.job.configuration.copy.sourceTables = [
+ bigquery_compat.TableReference(
+ projectId="test-project",
+ datasetId="test_dataset",
+ tableId="src_table")
+ ]
+ request.job.configuration.copy.destinationTable = (
+ bigquery_compat.TableReference(
+ projectId="test-project",
+ datasetId="test_dataset",
+ tableId="dst_table"))
+ request.job.configuration.copy.writeDisposition = None
+ request.job.configuration.copy.createDisposition = None
+ request.job.configuration.labels = None
+
+ jobs_compat.Insert(request)
+ client.copy_table.assert_called_once()
+ called_config = client.copy_table.call_args.kwargs["job_config"]
+ self.assertEqual(called_config.labels, {})
+
+ def test_perform_extract_job_with_none_labels(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.MagicMock(spec=gcp_bigquery.Client)
+ mock_job = mock.MagicMock()
+ mock_job.job_id = "extract_job_id"
+ mock_job.project = "test-project"
+ mock_job.location = "US"
+ client.extract_table.return_value = mock_job
+
+ jobs_compat = bigquery_compat._ClientJobsCompat(client)
+ request = mock.MagicMock()
+ request.projectId = "test-project"
+ request.job = mock.MagicMock()
+ request.job.jobReference = bigquery_compat.JobReference(
+ jobId="extract_job_id")
+ request.job.configuration = mock.MagicMock(spec=["extract", "labels"])
+ request.job.configuration.extract.sourceTable = (
+ bigquery_compat.TableReference(
+ projectId="test-project",
+ datasetId="test_dataset",
+ tableId="src_table"))
+ request.job.configuration.extract.destinationUris = [
+ "gs://test-bucket/output.csv"
+ ]
+ request.job.configuration.extract.destinationFormat = "CSV"
+ request.job.configuration.labels = None
+
+ jobs_compat.Insert(request)
+ client.extract_table.assert_called_once()
+ called_config = client.extract_table.call_args.kwargs["job_config"]
+ self.assertEqual(called_config.labels, {})
+
+ def test_labels_setter_clears_on_none(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ # Table labels
+ table = gcp_bigquery.Table("test-project.dataset.table")
+ table.labels = {"initial": "label"}
+ self.assertEqual(table.labels, {"initial": "label"})
+ table.labels = None
+ self.assertEqual(table.labels, {})
+
+ # Dataset labels
+ ds = gcp_bigquery.Dataset("test-project.dataset")
+ ds.labels = {"initial": "label"}
+ self.assertEqual(ds.labels, {"initial": "label"})
+ ds.labels = None
+ self.assertEqual(ds.labels, {})
+
+ # Job configs labels
+ for config_cls in (
+ gcp_bigquery.QueryJobConfig,
+ gcp_bigquery.LoadJobConfig,
+ gcp_bigquery.CopyJobConfig,
+ gcp_bigquery.ExtractJobConfig,
+ ):
+ cfg = config_cls(labels={"initial": "label"})
+ self.assertEqual(cfg.labels, {"initial": "label"})
+ cfg.labels = None
+ self.assertEqual(cfg.labels, {})
+
+ def test_table_partitioning_property_setters(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ table = gcp_bigquery.Table("test-project.dataset.table")
+ tp = gcp_bigquery.TimePartitioning(type_="DAY")
+ table.timePartitioning = tp
+ self.assertEqual(table.timePartitioning, tp)
+ self.assertEqual(table.time_partitioning, tp)
+
+ rp = gcp_bigquery.RangePartitioning(field="id")
+ table.rangePartitioning = rp
+ self.assertEqual(table.rangePartitioning, rp)
+ self.assertEqual(table.range_partitioning, rp)
+
+
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestSchemaConversionCompatibility(unittest.TestCase):
+ def test_to_table_schema_nested_records(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ schema = [
+ gcp_bigquery.SchemaField(
+ "person",
+ "RECORD",
+ mode="NULLABLE",
+ fields=[
+ gcp_bigquery.SchemaField("name", "STRING", mode="REQUIRED"),
+ gcp_bigquery.SchemaField("age", "INTEGER", mode="NULLABLE"),
+ gcp_bigquery.SchemaField(
+ "address",
+ "RECORD",
+ fields=[
+ gcp_bigquery.SchemaField(
+ "city", "STRING", mode="NULLABLE"),
+ ]),
+ ]),
+ ]
+ table_schema = bigquery_compat._to_table_schema(schema)
+ self.assertEqual(len(table_schema.fields), 1)
+ person = table_schema.fields[0]
+ self.assertEqual(person.name, "person")
+ self.assertEqual(person.type, "RECORD")
+ self.assertEqual(len(person.fields), 3)
+ self.assertEqual(person.fields[0].name, "name")
+ self.assertEqual(person.fields[0].type, "STRING")
+ self.assertEqual(person.fields[1].name, "age")
+ self.assertEqual(person.fields[1].type, "INTEGER")
+ self.assertEqual(person.fields[2].name, "address")
+ self.assertEqual(person.fields[2].type, "RECORD")
+ self.assertEqual(len(person.fields[2].fields), 1)
+ self.assertEqual(person.fields[2].fields[0].name, "city")
+ self.assertEqual(person.fields[2].fields[0].type, "STRING")
+
+ def test_to_table_schema_dict(self):
+ dict_schema = {
+ "fields": [
+ {
+ "name": "id", "type": "INTEGER", "mode": "REQUIRED"
+ },
+ {
+ "name": "val", "type": "STRING", "mode": "NULLABLE"
+ },
+ ]
+ }
+ table_schema = bigquery_compat._to_table_schema(dict_schema)
+ self.assertEqual(len(table_schema.fields), 2)
+ self.assertEqual(table_schema.fields[0].name, "id")
+ self.assertEqual(table_schema.fields[0].type, "INTEGER")
+ self.assertEqual(table_schema.fields[1].name, "val")
+ self.assertEqual(table_schema.fields[1].type, "STRING")
+
+
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestClientCompatibility(unittest.TestCase):
+ def test_job_stats_referenced_tables(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ mock_job = mock.MagicMock(spec=gcp_bigquery.QueryJob)
+ mock_job.job_id = "query_job_id"
+ mock_job.project = "test-project"
+ mock_job.location = "US"
+ t1 = gcp_bigquery.TableReference.from_string("test-project.dataset.table1")
+ t2 = gcp_bigquery.TableReference.from_string("test-project.dataset.table2")
+ mock_job.referenced_tables = [t1, t2]
+
+ stats = bigquery_compat._JobStatsCompat(mock_job)
+ ref_tables = stats.referencedTables
+ self.assertEqual(len(ref_tables), 2)
+ self.assertEqual(ref_tables[0].projectId, "test-project")
+ self.assertEqual(ref_tables[0].datasetId, "dataset")
+ self.assertEqual(ref_tables[0].tableId, "table1")
+ self.assertEqual(ref_tables[1].projectId, "test-project")
+ self.assertEqual(ref_tables[1].datasetId, "dataset")
+ self.assertEqual(ref_tables[1].tableId, "table2")
+
+ def test_client_tables_compat_insert_labels_and_metadata(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.MagicMock(spec=gcp_bigquery.Client)
+ client.project = "test-project"
+ created_table_mock = mock.MagicMock(spec=gcp_bigquery.Table)
+ client.create_table.return_value = created_table_mock
+
+ tables_compat = bigquery_compat._ClientTablesCompat(client)
+ request = mock.MagicMock()
+ request.projectId = "test-project"
+ request.datasetId = "test_dataset"
+ request.table = mock.MagicMock()
+ request.table.tableReference = mock.MagicMock()
+ request.table.tableReference.projectId = "test-project"
+ request.table.tableReference.datasetId = "test_dataset"
+ request.table.tableReference.tableId = "test_table"
+ request.table.schema = None
+ request.table.labels = {"env": "test", "tier": "frontend"}
+ request.table.friendlyName = "My Test Table"
+ request.table.description = "A test table description"
+ request.table.timePartitioning = None
+ request.table.rangePartitioning = None
+ request.table.clustering = None
+ request.table.encryptionConfiguration = None
+
+ tables_compat.Insert(request)
+ client.create_table.assert_called_once()
+ passed_table = client.create_table.call_args.args[0]
+ self.assertEqual(passed_table.labels, {"env": "test", "tier": "frontend"})
+ self.assertEqual(passed_table.friendly_name, "My Test Table")
+ self.assertEqual(passed_table.description, "A test table description")
+
+
[email protected](
+ HttpError is None or gcp_bigquery is None,
+ 'GCP dependencies are not installed')
+class TestTablePartitioningAndClusteringCompatibility(unittest.TestCase):
+ def test_time_partitioning_from_dict_and_camel_case(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ table = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table.timePartitioning = {"type": "DAY"}
+ self.assertIsNotNone(table.time_partitioning)
+ self.assertEqual(table.time_partitioning.type_, "DAY")
+ self.assertEqual(table.timePartitioning.type, "DAY")
+
+ table2 = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table2.time_partitioning = {
+ "type": "HOUR", "field": "ts", "expirationMs": 86400000
+ }
+ self.assertEqual(table2.time_partitioning.type_, "HOUR")
+ self.assertEqual(table2.time_partitioning.field, "ts")
+ self.assertEqual(table2.time_partitioning.expiration_ms, 86400000)
+
+ def test_range_partitioning_from_dict_and_camel_case(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ table = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table.rangePartitioning = {
+ "field": "id", "range": {
+ "start": 0, "end": 100, "interval": 10
+ }
+ }
+ self.assertIsNotNone(table.range_partitioning)
+ self.assertEqual(table.range_partitioning.field, "id")
+ self.assertEqual(table.range_partitioning.range_.start, 0)
+ self.assertEqual(table.range_partitioning.range_.end, 100)
+ self.assertEqual(table.range_partitioning.range_.interval, 10)
+
+ def test_clustering_from_dict_and_list(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ table = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table.clustering = {"fields": ["language", "country"]}
+ self.assertEqual(table.clustering_fields, ["language", "country"])
+ self.assertEqual(table.clustering.fields, ["language", "country"])
+ self.assertEqual(table.clustering["fields"], ["language", "country"])
+ self.assertEqual(table.clustering.get("fields"), ["language", "country"])
+
+ table2 = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table2.clustering = ["language"]
+ self.assertEqual(table2.clustering_fields, ["language"])
+ self.assertEqual(table2.clustering.fields, ["language"])
+
+ def test_client_tables_compat_insert_with_partitioning_and_clustering(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.Mock()
+ tables_compat = bigquery_compat._ClientTablesCompat(client)
+
+ request = mock.Mock()
+ request.projectId = "my-proj"
+ request.datasetId = "my-ds"
+ request.table = mock.Mock()
+ request.table.tableReference = mock.Mock()
+ request.table.tableReference.projectId = "my-proj"
+ request.table.tableReference.datasetId = "my-ds"
+ request.table.tableReference.tableId = "my-tbl"
+ request.table.schema = None
+ request.table.timePartitioning = {"type": "DAY"}
+ request.table.rangePartitioning = None
+ request.table.clustering = {"fields": ["language"]}
+ request.table.description = None
+ request.table.friendlyName = None
+ request.table.labels = None
+ request.table.encryptionConfiguration = None
+
+ tables_compat.Insert(request)
+ client.create_table.assert_called_once()
+ passed_table = client.create_table.call_args.args[0]
+ self.assertEqual(passed_table.time_partitioning.type_, "DAY")
+ self.assertEqual(passed_table.clustering_fields, ["language"])
+
+ def test_range_partitioning_from_apitools_with_zero_start(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ # If apitools RangePartitioning model is available, test with it directly
+ if apitools_bigquery is not None and hasattr(apitools_bigquery,
+ "RangePartitioning"):
+ rp = apitools_bigquery.RangePartitioning(
+ field="id",
+ range=apitools_bigquery.RangePartitioning.RangeValue(
+ start=0, end=100, interval=10))
+ gcp_rp = bigquery_compat._to_gcp_range_partitioning(rp)
+ self.assertEqual(gcp_rp.field, "id")
+ self.assertEqual(gcp_rp.range_.start, 0)
+ self.assertEqual(gcp_rp.range_.end, 100)
+ self.assertEqual(gcp_rp.range_.interval, 10)
+
+ # Also test with object having start=0
+ class RangeObj:
+ def __init__(self):
+ self.start = 0
+ self.end = 50
+ self.interval = 5
+
+ class RPObj:
+ def __init__(self):
+ self.field = "num"
+ self.range = RangeObj()
+
+ gcp_rp2 = bigquery_compat._to_gcp_range_partitioning(RPObj())
+ self.assertEqual(gcp_rp2.field, "num")
+ self.assertEqual(gcp_rp2.range_.start, 0)
+ self.assertEqual(gcp_rp2.range_.end, 50)
+ self.assertEqual(gcp_rp2.range_.interval, 5)
+
+ def test_time_partitioning_require_filter_false_and_str(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ # String input
+ table = gcp_bigquery.Table("my-proj.my_ds.my_tbl")
+ table.timePartitioning = "DAY"
+ self.assertEqual(table.time_partitioning.type_, "DAY")
+
+ # requirePartitionFilter = False and expirationMs = 0 explicitly preserved
+ table2 = gcp_bigquery.Table("my-proj.my_ds.my_tbl2")
+ table2.timePartitioning = {
+ "type": "HOUR",
+ "requirePartitionFilter": False,
+ "expirationMs": 0,
+ }
+ self.assertEqual(table2.time_partitioning.type_, "HOUR")
+ self.assertIs(table2.time_partitioning.require_partition_filter, False)
+ self.assertEqual(table2.time_partitioning.expiration_ms, 0)
+
+ def test_client_tables_compat_insert_with_bare_mock(self):
+ if gcp_bigquery is None:
+ raise unittest.SkipTest("google-cloud-bigquery is not installed")
+
+ client = mock.Mock()
+ tables_compat = bigquery_compat._ClientTablesCompat(client)
+
+ # Bare mock where request.table has unconfigured attributes returning Mocks
+ req = mock.Mock()
+ req.table = mock.Mock()
+ req.table.tableReference = mock.Mock()
+ req.table.tableReference.projectId = "my-proj"
+ req.table.tableReference.datasetId = "my-ds"
+ req.table.tableReference.tableId = "my-tbl"
+ req.table.schema = None
+
+ tables_compat.Insert(req)
+ client.create_table.assert_called_once()
+
+
+if __name__ == "__main__":
+ logging.getLogger().setLevel(logging.INFO)
+ unittest.main()
