google-cloud-spanner 3.55.0__py3-none-any.whl

google/cloud/spanner_v1/keyset.py

@@ -0,0 +1,193 @@
+# Copyright 2016 Google LLC All rights reserved.
+#
+# 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.
+
+"""Wrap representation of Spanner keys / ranges."""
+
+from google.cloud.spanner_v1 import KeyRangePB
+from google.cloud.spanner_v1 import KeySetPB
+
+from google.cloud.spanner_v1._helpers import _make_list_value_pb
+from google.cloud.spanner_v1._helpers import _make_list_value_pbs
+
+
+class KeyRange(object):
+    """Identify range of table rows via start / end points.
+
+    Specify either a `start_open` or `start_closed` key, or defaults to
+    `start_closed = []`.  Specify either an `end_open` or `end_closed` key,
+    or defaults to `end_closed = []`.  However, at least one key has to be
+    specified.  If no keys are specified, ValueError is raised.
+
+    :type start_open: list of scalars
+    :param start_open: keys identifying start of range (this key excluded)
+
+    :type start_closed: list of scalars
+    :param start_closed: keys identifying start of range (this key included)
+
+    :type end_open: list of scalars
+    :param end_open: keys identifying end of range (this key excluded)
+
+    :type end_closed: list of scalars
+    :param end_closed: keys identifying end of range (this key included)
+
+    :raises ValueError: if no keys are specified
+    """
+
+    def __init__(
+        self, start_open=None, start_closed=None, end_open=None, end_closed=None
+    ):
+        if not any([start_open, start_closed, end_open, end_closed]):
+            raise ValueError("Must specify at least a start or end row.")
+
+        if start_open and start_closed:
+            raise ValueError("Specify one of 'start_open' / 'start_closed'.")
+        elif start_open is None and start_closed is None:
+            start_closed = []
+
+        if end_open and end_closed:
+            raise ValueError("Specify one of 'end_open' / 'end_closed'.")
+        elif end_open is None and end_closed is None:
+            end_closed = []
+
+        self.start_open = start_open
+        self.start_closed = start_closed
+        self.end_open = end_open
+        self.end_closed = end_closed
+
+    def _to_pb(self):
+        """Construct a KeyRange protobuf.
+
+        :rtype: :class:`~google.cloud.spanner_v1.types.KeyRange`
+        :returns: protobuf corresponding to this instance.
+        """
+        kwargs = {}
+
+        if self.start_open is not None:
+            kwargs["start_open"] = _make_list_value_pb(self.start_open)
+
+        if self.start_closed is not None:
+            kwargs["start_closed"] = _make_list_value_pb(self.start_closed)
+
+        if self.end_open is not None:
+            kwargs["end_open"] = _make_list_value_pb(self.end_open)
+
+        if self.end_closed is not None:
+            kwargs["end_closed"] = _make_list_value_pb(self.end_closed)
+
+        return KeyRangePB(**kwargs)
+
+    def _to_dict(self):
+        """Return the state of the keyrange as a dict.
+
+        :rtype: dict
+        :returns: state of this instance.
+        """
+        mapping = {}
+
+        if self.start_open:
+            mapping["start_open"] = self.start_open
+
+        if self.start_closed:
+            mapping["start_closed"] = self.start_closed
+
+        if self.end_open:
+            mapping["end_open"] = self.end_open
+
+        if self.end_closed:
+            mapping["end_closed"] = self.end_closed
+
+        return mapping
+
+    def __eq__(self, other):
+        """Compare by serialized state."""
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return self._to_dict() == other._to_dict()
+
+
+class KeySet(object):
+    """Identify table rows via keys / ranges.
+
+    :type keys: list of list of scalars
+    :param keys: keys identifying individual rows within a table.
+
+    :type ranges: list of :class:`KeyRange`
+    :param ranges: ranges identifying rows within a table.
+
+    :type all_: boolean
+    :param all_: if True, identify all rows within a table
+    """
+
+    def __init__(self, keys=(), ranges=(), all_=False):
+        if all_ and (keys or ranges):
+            raise ValueError("'all_' is exclusive of 'keys' / 'ranges'.")
+        self.keys = list(keys)
+        self.ranges = list(ranges)
+        self.all_ = all_
+
+    def _to_pb(self):
+        """Construct a KeySet protobuf.
+
+        :rtype: :class:`~google.cloud.spanner_v1.types.KeySet`
+        :returns: protobuf corresponding to this instance.
+        """
+        if self.all_:
+            return KeySetPB(all_=True)
+        kwargs = {}
+
+        if self.keys:
+            kwargs["keys"] = _make_list_value_pbs(self.keys)
+
+        if self.ranges:
+            kwargs["ranges"] = [krange._to_pb() for krange in self.ranges]
+
+        return KeySetPB(**kwargs)
+
+    def _to_dict(self):
+        """Return the state of the keyset as a dict.
+
+        The result can be used to serialize the instance and reconstitute
+        it later using :meth:`_from_dict`.
+
+        :rtype: dict
+        :returns: state of this instance.
+        """
+        if self.all_:
+            return {"all": True}
+
+        return {
+            "keys": self.keys,
+            "ranges": [keyrange._to_dict() for keyrange in self.ranges],
+        }
+
+    def __eq__(self, other):
+        """Compare by serialized state."""
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+        return self._to_dict() == other._to_dict()
+
+    @classmethod
+    def _from_dict(cls, mapping):
+        """Create an instance from the corresponding state mapping.
+
+        :type mapping: dict
+        :param mapping: the instance state.
+        """
+        if mapping.get("all"):
+            return cls(all_=True)
+
+        r_mappings = mapping.get("ranges", ())
+        ranges = [KeyRange(**r_mapping) for r_mapping in r_mappings]
+
+        return cls(keys=mapping.get("keys", ()), ranges=ranges)

google/cloud/spanner_v1/merged_result_set.py

@@ -0,0 +1,146 @@
+# Copyright 2024 Google LLC All rights reserved.
+#
+# 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.
+from concurrent.futures import ThreadPoolExecutor
+from dataclasses import dataclass
+from queue import Queue
+from typing import Any, TYPE_CHECKING
+from threading import Lock, Event
+
+from google.cloud.spanner_v1._opentelemetry_tracing import trace_call
+from google.cloud.spanner_v1.metrics.metrics_capture import MetricsCapture
+
+if TYPE_CHECKING:
+    from google.cloud.spanner_v1.database import BatchSnapshot
+
+QUEUE_SIZE_PER_WORKER = 32
+MAX_PARALLELISM = 16
+
+
+class PartitionExecutor:
+    """
+    Executor that executes single partition on a separate thread and inserts
+    rows in the queue
+    """
+
+    def __init__(self, batch_snapshot, partition_id, merged_result_set):
+        self._batch_snapshot: BatchSnapshot = batch_snapshot
+        self._partition_id = partition_id
+        self._merged_result_set: MergedResultSet = merged_result_set
+        self._queue: Queue[PartitionExecutorResult] = merged_result_set._queue
+
+    def run(self):
+        observability_options = getattr(
+            self._batch_snapshot, "observability_options", {}
+        )
+        with trace_call(
+            "CloudSpanner.PartitionExecutor.run",
+            observability_options=observability_options,
+        ), MetricsCapture():
+            self.__run()
+
+    def __run(self):
+        results = None
+        try:
+            results = self._batch_snapshot.process_query_batch(self._partition_id)
+            for row in results:
+                if self._merged_result_set._metadata is None:
+                    self._set_metadata(results)
+                self._queue.put(PartitionExecutorResult(data=row))
+            # Special case: The result set did not return any rows.
+            # Push the metadata to the merged result set.
+            if self._merged_result_set._metadata is None:
+                self._set_metadata(results)
+        except Exception as ex:
+            if self._merged_result_set._metadata is None:
+                self._set_metadata(results, True)
+            self._queue.put(PartitionExecutorResult(exception=ex))
+        finally:
+            # Emit a special 'is_last' result to ensure that the MergedResultSet
+            # is not blocked on a queue that never receives any more results.
+            self._queue.put(PartitionExecutorResult(is_last=True))
+
+    def _set_metadata(self, results, is_exception=False):
+        self._merged_result_set.metadata_lock.acquire()
+        try:
+            if not is_exception:
+                self._merged_result_set._metadata = results.metadata
+        finally:
+            self._merged_result_set.metadata_lock.release()
+            self._merged_result_set.metadata_event.set()
+
+
+@dataclass
+class PartitionExecutorResult:
+    data: Any = None
+    exception: Exception = None
+    is_last: bool = False
+
+
+class MergedResultSet:
+    """
+    Executes multiple partitions on different threads and then combines the
+    results from multiple queries using a synchronized queue. The order of the
+    records in the MergedResultSet is not guaranteed.
+    """
+
+    def __init__(self, batch_snapshot, partition_ids, max_parallelism):
+        self._exception = None
+        self._metadata = None
+        self.metadata_event = Event()
+        self.metadata_lock = Lock()
+
+        partition_ids_count = len(partition_ids)
+        self._finished_count_down_latch = partition_ids_count
+        parallelism = min(MAX_PARALLELISM, partition_ids_count)
+        if max_parallelism != 0:
+            parallelism = min(partition_ids_count, max_parallelism)
+        self._queue = Queue(maxsize=QUEUE_SIZE_PER_WORKER * parallelism)
+
+        partition_executors = []
+        for partition_id in partition_ids:
+            partition_executors.append(
+                PartitionExecutor(batch_snapshot, partition_id, self)
+            )
+        executor = ThreadPoolExecutor(max_workers=parallelism)
+        for partition_executor in partition_executors:
+            executor.submit(partition_executor.run)
+        executor.shutdown(False)
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        if self._exception is not None:
+            raise self._exception
+        while True:
+            partition_result = self._queue.get()
+            if partition_result.is_last:
+                self._finished_count_down_latch -= 1
+                if self._finished_count_down_latch == 0:
+                    raise StopIteration
+            elif partition_result.exception is not None:
+                self._exception = partition_result.exception
+                raise self._exception
+            else:
+                return partition_result.data
+
+    @property
+    def metadata(self):
+        self.metadata_event.wait()
+        return self._metadata
+
+    @property
+    def stats(self):
+        # TODO: Implement
+        return None

google/cloud/spanner_v1/metrics/constants.py

@@ -0,0 +1,71 @@
+# Copyright 2025 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.
+
+BUILT_IN_METRICS_METER_NAME = "gax-python"
+NATIVE_METRICS_PREFIX = "spanner.googleapis.com/internal/client"
+SPANNER_RESOURCE_TYPE = "spanner_instance_client"
+SPANNER_SERVICE_NAME = "spanner-python"
+GOOGLE_CLOUD_RESOURCE_KEY = "google-cloud-resource-prefix"
+GOOGLE_CLOUD_REGION_KEY = "cloud.region"
+GOOGLE_CLOUD_REGION_GLOBAL = "global"
+SPANNER_METHOD_PREFIX = "/google.spanner.v1."
+ENABLE_SPANNER_METRICS_ENV_VAR = "SPANNER_ENABLE_BUILTIN_METRICS"
+
+# Monitored resource labels
+MONITORED_RES_LABEL_KEY_PROJECT = "project_id"
+MONITORED_RES_LABEL_KEY_INSTANCE = "instance_id"
+MONITORED_RES_LABEL_KEY_INSTANCE_CONFIG = "instance_config"
+MONITORED_RES_LABEL_KEY_LOCATION = "location"
+MONITORED_RES_LABEL_KEY_CLIENT_HASH = "client_hash"
+MONITORED_RESOURCE_LABELS = [
+    MONITORED_RES_LABEL_KEY_PROJECT,
+    MONITORED_RES_LABEL_KEY_INSTANCE,
+    MONITORED_RES_LABEL_KEY_INSTANCE_CONFIG,
+    MONITORED_RES_LABEL_KEY_LOCATION,
+    MONITORED_RES_LABEL_KEY_CLIENT_HASH,
+]
+
+# Metric labels
+METRIC_LABEL_KEY_CLIENT_UID = "client_uid"
+METRIC_LABEL_KEY_CLIENT_NAME = "client_name"
+METRIC_LABEL_KEY_DATABASE = "database"
+METRIC_LABEL_KEY_METHOD = "method"
+METRIC_LABEL_KEY_STATUS = "status"
+METRIC_LABEL_KEY_DIRECT_PATH_ENABLED = "directpath_enabled"
+METRIC_LABEL_KEY_DIRECT_PATH_USED = "directpath_used"
+METRIC_LABELS = [
+    METRIC_LABEL_KEY_CLIENT_UID,
+    METRIC_LABEL_KEY_CLIENT_NAME,
+    METRIC_LABEL_KEY_DATABASE,
+    METRIC_LABEL_KEY_METHOD,
+    METRIC_LABEL_KEY_STATUS,
+    METRIC_LABEL_KEY_DIRECT_PATH_ENABLED,
+    METRIC_LABEL_KEY_DIRECT_PATH_USED,
+]
+
+# Metric names
+METRIC_NAME_OPERATION_LATENCIES = "operation_latencies"
+METRIC_NAME_ATTEMPT_LATENCIES = "attempt_latencies"
+METRIC_NAME_OPERATION_COUNT = "operation_count"
+METRIC_NAME_ATTEMPT_COUNT = "attempt_count"
+METRIC_NAME_GFE_LATENCY = "gfe_latency"
+METRIC_NAME_GFE_MISSING_HEADER_COUNT = "gfe_missing_header_count"
+METRIC_NAMES = [
+    METRIC_NAME_OPERATION_LATENCIES,
+    METRIC_NAME_ATTEMPT_LATENCIES,
+    METRIC_NAME_OPERATION_COUNT,
+    METRIC_NAME_ATTEMPT_COUNT,
+]
+
+METRIC_EXPORT_INTERVAL_MS = 60000  # 1 Minute

google/cloud/spanner_v1/metrics/metrics_capture.py

@@ -0,0 +1,75 @@
+# Copyright 2025 Google LLC All rights reserved.
+#
+# 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.
+"""
+This module provides functionality for capturing metrics in Cloud Spanner operations.
+
+It includes a context manager class, MetricsCapture, which automatically handles the
+start and completion of metrics tracing for a given operation. This ensures that metrics
+are consistently recorded for Cloud Spanner operations, facilitating observability and
+performance monitoring.
+"""
+
+from .spanner_metrics_tracer_factory import SpannerMetricsTracerFactory
+
+
+class MetricsCapture:
+    """Context manager for capturing metrics in Cloud Spanner operations.
+
+    This class provides a context manager interface to automatically handle
+    the start and completion of metrics tracing for a given operation.
+    """
+
+    def __enter__(self):
+        """Enter the runtime context related to this object.
+
+        This method initializes a new metrics tracer for the operation and
+        records the start of the operation.
+
+        Returns:
+            MetricsCapture: The instance of the context manager.
+        """
+        # Short circuit out if metrics are disabled
+        factory = SpannerMetricsTracerFactory()
+        if not factory.enabled:
+            return self
+
+        # Define a new metrics tracer for the new operation
+        SpannerMetricsTracerFactory.current_metrics_tracer = (
+            factory.create_metrics_tracer()
+        )
+        if SpannerMetricsTracerFactory.current_metrics_tracer:
+            SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_start()
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Exit the runtime context related to this object.
+
+        This method records the completion of the operation. If an exception
+        occurred, it will be propagated after the metrics are recorded.
+
+        Args:
+            exc_type (Type[BaseException]): The exception type.
+            exc_value (BaseException): The exception value.
+            traceback (TracebackType): The traceback object.
+
+        Returns:
+            bool: False to propagate the exception if any occurred.
+        """
+        # Short circuit out if metrics are disable
+        if not SpannerMetricsTracerFactory().enabled:
+            return False
+
+        if SpannerMetricsTracerFactory.current_metrics_tracer:
+            SpannerMetricsTracerFactory.current_metrics_tracer.record_operation_completion()
+        return False  # Propagate the exception if any