google-cloud-bigquery 3.31.0__py3-none-any.whl

google/cloud/bigquery/model.py

@@ -0,0 +1,517 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 Google LLC
+#
+# Licensed 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
+#
+#     https://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.
+
+"""Define resources for the BigQuery ML Models API."""
+
+from __future__ import annotations  # type: ignore
+
+import copy
+import datetime
+import typing
+from typing import Any, Dict, Optional, Sequence, Union
+
+import google.cloud._helpers  # type: ignore
+from google.cloud.bigquery import _helpers
+from google.cloud.bigquery import standard_sql
+from google.cloud.bigquery.encryption_configuration import EncryptionConfiguration
+
+
+class Model:
+    """Model represents a machine learning model resource.
+
+    See
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/models
+
+    Args:
+        model_ref:
+            A pointer to a model. If ``model_ref`` is a string, it must
+            included a project ID, dataset ID, and model ID, each separated
+            by ``.``.
+    """
+
+    _PROPERTY_TO_API_FIELD = {
+        "expires": "expirationTime",
+        "friendly_name": "friendlyName",
+        # Even though it's not necessary for field mapping to map when the
+        # property name equals the resource name, we add these here so that we
+        # have an exhaustive list of all mutable properties.
+        "labels": "labels",
+        "description": "description",
+        "encryption_configuration": "encryptionConfiguration",
+    }
+
+    def __init__(self, model_ref: Union["ModelReference", str, None]):
+        # Use _properties on read-write properties to match the REST API
+        # semantics. The BigQuery API makes a distinction between an unset
+        # value, a null value, and a default value (0 or ""), but the protocol
+        # buffer classes do not.
+        self._properties: Dict[str, Any] = {}
+
+        if isinstance(model_ref, str):
+            model_ref = ModelReference.from_string(model_ref)
+
+        if model_ref:
+            self._properties["modelReference"] = model_ref.to_api_repr()
+
+    @property
+    def reference(self) -> Optional["ModelReference"]:
+        """A model reference pointing to this model.
+
+        Read-only.
+        """
+        resource = self._properties.get("modelReference")
+        if resource is None:
+            return None
+        else:
+            return ModelReference.from_api_repr(resource)
+
+    @property
+    def project(self) -> Optional[str]:
+        """Project bound to the model."""
+        ref = self.reference
+        return ref.project if ref is not None else None
+
+    @property
+    def dataset_id(self) -> Optional[str]:
+        """ID of dataset containing the model."""
+        ref = self.reference
+        return ref.dataset_id if ref is not None else None
+
+    @property
+    def model_id(self) -> Optional[str]:
+        """The model ID."""
+        ref = self.reference
+        return ref.model_id if ref is not None else None
+
+    @property
+    def path(self) -> Optional[str]:
+        """URL path for the model's APIs."""
+        ref = self.reference
+        return ref.path if ref is not None else None
+
+    @property
+    def location(self) -> Optional[str]:
+        """The geographic location where the model resides.
+
+        This value is inherited from the dataset.
+
+        Read-only.
+        """
+        return typing.cast(Optional[str], self._properties.get("location"))
+
+    @property
+    def etag(self) -> Optional[str]:
+        """ETag for the model resource (:data:`None` until set from the server).
+
+        Read-only.
+        """
+        return typing.cast(Optional[str], self._properties.get("etag"))
+
+    @property
+    def created(self) -> Optional[datetime.datetime]:
+        """Datetime at which the model was created (:data:`None` until set from the server).
+
+        Read-only.
+        """
+        value = typing.cast(Optional[float], self._properties.get("creationTime"))
+        if value is None:
+            return None
+        else:
+            # value will be in milliseconds.
+            return google.cloud._helpers._datetime_from_microseconds(
+                1000.0 * float(value)
+            )
+
+    @property
+    def modified(self) -> Optional[datetime.datetime]:
+        """Datetime at which the model was last modified (:data:`None` until set from the server).
+
+        Read-only.
+        """
+        value = typing.cast(Optional[float], self._properties.get("lastModifiedTime"))
+        if value is None:
+            return None
+        else:
+            # value will be in milliseconds.
+            return google.cloud._helpers._datetime_from_microseconds(
+                1000.0 * float(value)
+            )
+
+    @property
+    def model_type(self) -> str:
+        """Type of the model resource.
+
+        Read-only.
+        """
+        return typing.cast(
+            str, self._properties.get("modelType", "MODEL_TYPE_UNSPECIFIED")
+        )
+
+    @property
+    def training_runs(self) -> Sequence[Dict[str, Any]]:
+        """Information for all training runs in increasing order of start time.
+
+        Dictionaries are in REST API format. See:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/models#trainingrun
+
+        Read-only.
+        """
+        return typing.cast(
+            Sequence[Dict[str, Any]], self._properties.get("trainingRuns", [])
+        )
+
+    @property
+    def feature_columns(self) -> Sequence[standard_sql.StandardSqlField]:
+        """Input feature columns that were used to train this model.
+
+        Read-only.
+        """
+        resource: Sequence[Dict[str, Any]] = typing.cast(
+            Sequence[Dict[str, Any]], self._properties.get("featureColumns", [])
+        )
+        return [
+            standard_sql.StandardSqlField.from_api_repr(column) for column in resource
+        ]
+
+    @property
+    def transform_columns(self) -> Sequence[TransformColumn]:
+        """The input feature columns that were used to train this model.
+        The output transform columns used to train this model.
+
+        See REST API:
+        https://cloud.google.com/bigquery/docs/reference/rest/v2/models#transformcolumn
+
+        Read-only.
+        """
+        resources: Sequence[Dict[str, Any]] = typing.cast(
+            Sequence[Dict[str, Any]], self._properties.get("transformColumns", [])
+        )
+        return [TransformColumn(resource) for resource in resources]
+
+    @property
+    def label_columns(self) -> Sequence[standard_sql.StandardSqlField]:
+        """Label columns that were used to train this model.
+
+        The output of the model will have a ``predicted_`` prefix to these columns.
+
+        Read-only.
+        """
+        resource: Sequence[Dict[str, Any]] = typing.cast(
+            Sequence[Dict[str, Any]], self._properties.get("labelColumns", [])
+        )
+        return [
+            standard_sql.StandardSqlField.from_api_repr(column) for column in resource
+        ]
+
+    @property
+    def best_trial_id(self) -> Optional[int]:
+        """The best trial_id across all training runs.
+
+        .. deprecated::
+            This property is deprecated!
+
+        Read-only.
+        """
+        value = typing.cast(Optional[int], self._properties.get("bestTrialId"))
+        if value is not None:
+            value = int(value)
+        return value
+
+    @property
+    def expires(self) -> Optional[datetime.datetime]:
+        """The datetime when this model expires.
+
+        If not present, the model will persist indefinitely. Expired models will be
+        deleted and their storage reclaimed.
+        """
+        value = typing.cast(Optional[float], self._properties.get("expirationTime"))
+        if value is None:
+            return None
+        else:
+            # value will be in milliseconds.
+            return google.cloud._helpers._datetime_from_microseconds(
+                1000.0 * float(value)
+            )
+
+    @expires.setter
+    def expires(self, value: Optional[datetime.datetime]):
+        if value is None:
+            value_to_store: Optional[str] = None
+        else:
+            value_to_store = str(google.cloud._helpers._millis_from_datetime(value))
+        # TODO: Consider using typing.TypedDict when only Python 3.8+ is supported.
+        self._properties["expirationTime"] = value_to_store  # type: ignore
+
+    @property
+    def description(self) -> Optional[str]:
+        """Description of the model (defaults to :data:`None`)."""
+        return typing.cast(Optional[str], self._properties.get("description"))
+
+    @description.setter
+    def description(self, value: Optional[str]):
+        # TODO: Consider using typing.TypedDict when only Python 3.8+ is supported.
+        self._properties["description"] = value  # type: ignore
+
+    @property
+    def friendly_name(self) -> Optional[str]:
+        """Title of the table (defaults to :data:`None`)."""
+        return typing.cast(Optional[str], self._properties.get("friendlyName"))
+
+    @friendly_name.setter
+    def friendly_name(self, value: Optional[str]):
+        # TODO: Consider using typing.TypedDict when only Python 3.8+ is supported.
+        self._properties["friendlyName"] = value  # type: ignore
+
+    @property
+    def labels(self) -> Dict[str, str]:
+        """Labels for the table.
+
+        This method always returns a dict. To change a model's labels, modify the dict,
+        then call ``Client.update_model``. To delete a label, set its value to
+        :data:`None` before updating.
+        """
+        return self._properties.setdefault("labels", {})
+
+    @labels.setter
+    def labels(self, value: Optional[Dict[str, str]]):
+        if value is None:
+            value = {}
+        self._properties["labels"] = value
+
+    @property
+    def encryption_configuration(self) -> Optional[EncryptionConfiguration]:
+        """Custom encryption configuration for the model.
+
+        Custom encryption configuration (e.g., Cloud KMS keys) or :data:`None`
+        if using default encryption.
+
+        See `protecting data with Cloud KMS keys
+        <https://cloud.google.com/bigquery/docs/customer-managed-encryption>`_
+        in the BigQuery documentation.
+        """
+        prop = self._properties.get("encryptionConfiguration")
+        if prop:
+            prop = EncryptionConfiguration.from_api_repr(prop)
+        return typing.cast(Optional[EncryptionConfiguration], prop)
+
+    @encryption_configuration.setter
+    def encryption_configuration(self, value: Optional[EncryptionConfiguration]):
+        api_repr = value.to_api_repr() if value else value
+        self._properties["encryptionConfiguration"] = api_repr
+
+    @classmethod
+    def from_api_repr(cls, resource: Dict[str, Any]) -> "Model":
+        """Factory: construct a model resource given its API representation
+
+        Args:
+            resource:
+                Model resource representation from the API
+
+        Returns:
+            Model parsed from ``resource``.
+        """
+        this = cls(None)
+        resource = copy.deepcopy(resource)
+        this._properties = resource
+        return this
+
+    def _build_resource(self, filter_fields):
+        """Generate a resource for ``update``."""
+        return _helpers._build_resource_from_properties(self, filter_fields)
+
+    def __repr__(self):
+        return f"Model(reference={self.reference!r})"
+
+    def to_api_repr(self) -> Dict[str, Any]:
+        """Construct the API resource representation of this model.
+
+        Returns:
+            Model reference represented as an API resource
+        """
+        return copy.deepcopy(self._properties)
+
+
+class ModelReference:
+    """ModelReferences are pointers to models.
+
+    See
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/models#modelreference
+    """
+
+    def __init__(self):
+        self._properties = {}
+
+    @property
+    def project(self):
+        """str: Project bound to the model"""
+        return self._properties.get("projectId")
+
+    @property
+    def dataset_id(self):
+        """str: ID of dataset containing the model."""
+        return self._properties.get("datasetId")
+
+    @property
+    def model_id(self):
+        """str: The model ID."""
+        return self._properties.get("modelId")
+
+    @property
+    def path(self) -> str:
+        """URL path for the model's APIs."""
+        return f"/projects/{self.project}/datasets/{self.dataset_id}/models/{self.model_id}"
+
+    @classmethod
+    def from_api_repr(cls, resource: Dict[str, Any]) -> "ModelReference":
+        """Factory: construct a model reference given its API representation.
+
+        Args:
+            resource:
+                Model reference representation returned from the API
+
+        Returns:
+            Model reference parsed from ``resource``.
+        """
+        ref = cls()
+        ref._properties = resource
+        return ref
+
+    @classmethod
+    def from_string(
+        cls, model_id: str, default_project: Optional[str] = None
+    ) -> "ModelReference":
+        """Construct a model reference from model ID string.
+
+        Args:
+            model_id:
+                A model ID in standard SQL format. If ``default_project``
+                is not specified, this must included a project ID, dataset
+                ID, and model ID, each separated by ``.``.
+            default_project:
+                The project ID to use when ``model_id`` does not include
+                a project ID.
+
+        Returns:
+            Model reference parsed from ``model_id``.
+
+        Raises:
+            ValueError:
+                If ``model_id`` is not a fully-qualified table ID in
+                standard SQL format.
+        """
+        proj, dset, model = _helpers._parse_3_part_id(
+            model_id, default_project=default_project, property_name="model_id"
+        )
+        return cls.from_api_repr(
+            {"projectId": proj, "datasetId": dset, "modelId": model}
+        )
+
+    def to_api_repr(self) -> Dict[str, Any]:
+        """Construct the API resource representation of this model reference.
+
+        Returns:
+            Model reference represented as an API resource.
+        """
+        return copy.deepcopy(self._properties)
+
+    def _key(self):
+        """Unique key for this model.
+
+        This is used for hashing a ModelReference.
+        """
+        return self.project, self.dataset_id, self.model_id
+
+    def __eq__(self, other):
+        if not isinstance(other, ModelReference):
+            return NotImplemented
+        return self._properties == other._properties
+
+    def __ne__(self, other):
+        return not self == other
+
+    def __hash__(self):
+        return hash(self._key())
+
+    def __repr__(self):
+        return "ModelReference(project_id='{}', dataset_id='{}', model_id='{}')".format(
+            self.project, self.dataset_id, self.model_id
+        )
+
+
+class TransformColumn:
+    """TransformColumn represents a transform column feature.
+
+    See
+    https://cloud.google.com/bigquery/docs/reference/rest/v2/models#transformcolumn
+
+    Args:
+        resource:
+            A dictionary representing a transform column feature.
+    """
+
+    def __init__(self, resource: Dict[str, Any]):
+        self._properties = resource
+
+    @property
+    def name(self) -> Optional[str]:
+        """Name of the column."""
+        return self._properties.get("name")
+
+    @property
+    def type_(self) -> Optional[standard_sql.StandardSqlDataType]:
+        """Data type of the column after the transform.
+
+        Returns:
+            Optional[google.cloud.bigquery.standard_sql.StandardSqlDataType]:
+                Data type of the column.
+        """
+        type_json = self._properties.get("type")
+        if type_json is None:
+            return None
+        return standard_sql.StandardSqlDataType.from_api_repr(type_json)
+
+    @property
+    def transform_sql(self) -> Optional[str]:
+        """The SQL expression used in the column transform."""
+        return self._properties.get("transformSql")
+
+    @classmethod
+    def from_api_repr(cls, resource: Dict[str, Any]) -> "TransformColumn":
+        """Constructs a transform column feature given its API representation
+
+        Args:
+            resource:
+                Transform column feature representation from the API
+
+        Returns:
+            Transform column feature parsed from ``resource``.
+        """
+        this = cls({})
+        resource = copy.deepcopy(resource)
+        this._properties = resource
+        return this
+
+
+def _model_arg_to_model_ref(value, default_project=None):
+    """Helper to convert a string or Model to ModelReference.
+
+    This function keeps ModelReference and other kinds of objects unchanged.
+    """
+    if isinstance(value, str):
+        return ModelReference.from_string(value, default_project=default_project)
+    if isinstance(value, Model):
+        return value.reference
+    return value

google/cloud/bigquery/opentelemetry_tracing.py

@@ -0,0 +1,164 @@
+# Copyright 2020 Google LLC
+#
+# Licensed 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.
+
+import logging
+from contextlib import contextmanager
+from google.api_core.exceptions import GoogleAPICallError  # type: ignore
+
+logger = logging.getLogger(__name__)
+try:
+    from opentelemetry import trace  # type: ignore
+    from opentelemetry.instrumentation.utils import http_status_to_status_code  # type: ignore
+    from opentelemetry.trace.status import Status  # type: ignore
+
+    HAS_OPENTELEMETRY = True
+    _warned_telemetry = True
+
+except ImportError:
+    HAS_OPENTELEMETRY = False
+    _warned_telemetry = False
+
+_default_attributes = {
+    "db.system": "BigQuery"
+}  # static, default values assigned to all spans
+
+
+@contextmanager
+def create_span(name, attributes=None, client=None, job_ref=None):
+    """Creates a ContextManager for a Span to be exported to the configured exporter.
+    If no configuration exists yields None.
+
+        Args:
+            name (str): Name that will be set for the span being created
+            attributes (Optional[dict]):
+                Additional attributes that pertain to
+                the specific API call (i.e. not a default attribute)
+            client (Optional[google.cloud.bigquery.client.Client]):
+                Pass in a Client object to extract any attributes that may be
+                relevant to it and add them to the created spans.
+            job_ref (Optional[google.cloud.bigquery.job._AsyncJob])
+                Pass in a _AsyncJob object to extract any attributes that may be
+                relevant to it and add them to the created spans.
+
+        Yields:
+            opentelemetry.trace.Span: Yields the newly created Span.
+
+        Raises:
+            google.api_core.exceptions.GoogleAPICallError:
+                Raised if a span could not be yielded or issue with call to
+                OpenTelemetry.
+    """
+    global _warned_telemetry
+    final_attributes = _get_final_span_attributes(attributes, client, job_ref)
+    if not HAS_OPENTELEMETRY:
+        if not _warned_telemetry:
+            logger.debug(
+                "This service is instrumented using OpenTelemetry. "
+                "OpenTelemetry or one of its components could not be imported; "
+                "please add compatible versions of opentelemetry-api and "
+                "opentelemetry-instrumentation packages in order to get BigQuery "
+                "Tracing data."
+            )
+            _warned_telemetry = True
+
+        yield None
+        return
+    tracer = trace.get_tracer(__name__)
+
+    # yield new span value
+    with tracer.start_as_current_span(name=name, attributes=final_attributes) as span:
+        try:
+            yield span
+        except GoogleAPICallError as error:
+            if error.code is not None:
+                span.set_status(Status(http_status_to_status_code(error.code)))
+            raise
+
+
+def _get_final_span_attributes(attributes=None, client=None, job_ref=None):
+    """Compiles attributes from: client, job_ref, user-provided attributes.
+
+    Attributes from all of these sources are merged together. Note the
+    attributes are added sequentially based on perceived order of precedence:
+    i.e. attributes added last may overwrite attributes added earlier.
+
+    Args:
+        attributes (Optional[dict]):
+            Additional attributes that pertain to
+            the specific API call (i.e. not a default attribute)
+
+        client (Optional[google.cloud.bigquery.client.Client]):
+            Pass in a Client object to extract any attributes that may be
+            relevant to it and add them to the final_attributes
+
+        job_ref (Optional[google.cloud.bigquery.job._AsyncJob])
+            Pass in a _AsyncJob object to extract any attributes that may be
+            relevant to it and add them to the final_attributes.
+
+    Returns: dict
+    """
+
+    collected_attributes = _default_attributes.copy()
+
+    if client:
+        collected_attributes.update(_set_client_attributes(client))
+    if job_ref:
+        collected_attributes.update(_set_job_attributes(job_ref))
+    if attributes:
+        collected_attributes.update(attributes)
+
+    final_attributes = {k: v for k, v in collected_attributes.items() if v is not None}
+    return final_attributes
+
+
+def _set_client_attributes(client):
+    return {"db.name": client.project, "location": client.location}
+
+
+def _set_job_attributes(job_ref):
+    job_attributes = {
+        "db.name": job_ref.project,
+        "job_id": job_ref.job_id,
+        "state": job_ref.state,
+    }
+
+    job_attributes["hasErrors"] = job_ref.error_result is not None
+
+    if job_ref.created is not None:
+        job_attributes["timeCreated"] = job_ref.created.isoformat()
+
+    if job_ref.started is not None:
+        job_attributes["timeStarted"] = job_ref.started.isoformat()
+
+    if job_ref.ended is not None:
+        job_attributes["timeEnded"] = job_ref.ended.isoformat()
+
+    if job_ref.location is not None:
+        job_attributes["location"] = job_ref.location
+
+    if job_ref.parent_job_id is not None:
+        job_attributes["parent_job_id"] = job_ref.parent_job_id
+
+    if job_ref.num_child_jobs is not None:
+        job_attributes["num_child_jobs"] = job_ref.num_child_jobs
+
+    total_bytes_billed = getattr(job_ref, "total_bytes_billed", None)
+    if total_bytes_billed is not None:
+        job_attributes["total_bytes_billed"] = total_bytes_billed
+
+    total_bytes_processed = getattr(job_ref, "total_bytes_processed", None)
+    if total_bytes_processed is not None:
+        job_attributes["total_bytes_processed"] = total_bytes_processed
+
+    return job_attributes

google/cloud/bigquery/py.typed

@@ -0,0 +1,2 @@
+# Marker file for PEP 561.
+# The google-cloud-bigquery package uses inline types.