damccorm commented on code in PR #39998: URL: https://github.com/apache/beam/pull/39998#discussion_r3980502954
########## sdks/python/apache_beam/io/gcp/bigquery_compat.py: ########## @@ -0,0 +1,1413 @@ +# +# 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``. + +.. 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 + +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 [] + + +if apitools_bigquery is not None and hasattr(apitools_bigquery, Review Comment: > The intent is to throw a warning on import of the generated client and types (this is shown in the master PR and will be done in a follow-up) but if we wanted to make a really explicit opt-in here we could require some sort of arg? Yeah, I think we should do this (and still warn if they set it). This could probably just be an environment variable the user adds directly before the import since actually plumbing through an arg seems hard. My thought here, though, is basically: 1. This module is clearly internal. I think we're within our "rights" to make a breaking change 2. This module will go away, and the sooner folks start to move off of it the better. 3. At the same time, it is helpful to provide a smooth upgrade path so that folks can address this independently of the beam version bump 4. Warnings are easy to miss/ignore, while a forced action with an exception at least makes it clear that (a) this is potentially unsafe, and (b) this will go away. As much as possible, I want to avoid users depending on this without knowing (since I don't think it is guaranteed to be a perfect 1:1 replacement) -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
