_helpers.py

# 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.

"""Helper functions for Cloud Spanner."""

import datetime
import decimal
import math
import time
import base64
import threading
import logging
import uuid
from contextlib import contextmanager

from google.protobuf.struct_pb2 import ListValue
from google.protobuf.struct_pb2 import Value
from google.protobuf.message import Message
from google.protobuf.internal.enum_type_wrapper import EnumTypeWrapper

from google.api_core import datetime_helpers
from google.api_core.exceptions import Aborted
from google.cloud._helpers import _date_from_iso8601_date
from google.cloud.spanner_v1.types import ExecuteSqlRequest
from google.cloud.spanner_v1.types import TransactionOptions
from google.cloud.spanner_v1.types import ClientContext
from google.cloud.spanner_v1.types import RequestOptions
from google.cloud.spanner_v1.data_types import JsonObject, Interval
from google.cloud.spanner_v1.request_id_header import (
    with_request_id,
    with_request_id_metadata_only,
)
from google.cloud.spanner_v1.types import TypeCode
from google.cloud.spanner_v1.exceptions import wrap_with_request_id

from google.rpc.error_details_pb2 import RetryInfo

try:
    from opentelemetry.propagate import inject
    from opentelemetry.propagators.textmap import Setter
    from opentelemetry.semconv.resource import ResourceAttributes
    from opentelemetry.resourcedetector import gcp_resource_detector
    from opentelemetry.resourcedetector.gcp_resource_detector import (
        GoogleCloudResourceDetector,
    )

    # Overwrite the requests timeout for the detector.
    # This is necessary as the client will wait the full timeout if the
    # code is not run in a GCP environment, with the location endpoints available.
    gcp_resource_detector._TIMEOUT_SEC = 0.2

    HAS_OPENTELEMETRY_INSTALLED = True
except ImportError:
    HAS_OPENTELEMETRY_INSTALLED = False
from typing import List, Tuple
import random

# Validation error messages
NUMERIC_MAX_SCALE_ERR_MSG = (
    "Max scale for a numeric is 9. The requested numeric has scale {}"
)
NUMERIC_MAX_PRECISION_ERR_MSG = (
    "Max precision for the whole component of a numeric is 29. The requested "
    + "numeric has a whole component with precision {}"
)

GOOGLE_CLOUD_REGION_GLOBAL = "global"

log = logging.getLogger(__name__)

_cloud_region: str = None


if HAS_OPENTELEMETRY_INSTALLED:

    class OpenTelemetryContextSetter(Setter):
        """
        Used by Open Telemetry for context propagation.
        """

        def set(self, carrier: List[Tuple[str, str]], key: str, value: str) -> None:
            """
            Injects trace context into Spanner metadata

            Args:
                carrier(PubsubMessage): The Pub/Sub message which is the carrier of Open Telemetry
                data.
                key(str): The key for which the Open Telemetry context data needs to be set.
                value(str): The Open Telemetry context value to be set.

            Returns:
                None
            """
            carrier.append((key, value))


def _get_cloud_region() -> str:
    """Get the location of the resource, caching the result.

    Returns:
        str: The location of the resource. If OpenTelemetry is not installed, returns a global region.
    """
    global _cloud_region
    if _cloud_region is not None:
        return _cloud_region

    try:
        detector = GoogleCloudResourceDetector()
        resources = detector.detect()
        if ResourceAttributes.CLOUD_REGION in resources.attributes:
            _cloud_region = resources.attributes[ResourceAttributes.CLOUD_REGION]
        else:
            _cloud_region = GOOGLE_CLOUD_REGION_GLOBAL
    except Exception as e:
        log.warning(
            "Failed to detect GCP resource location for Spanner metrics, defaulting to 'global'. Error: %s",
            e,
        )
        _cloud_region = GOOGLE_CLOUD_REGION_GLOBAL

    return _cloud_region


def _try_to_coerce_bytes(bytestring):
    """Try to coerce a byte string into the right thing based on Python
    version and whether or not it is base64 encoded.

    Return a text string or raise ValueError.
    """
    # Attempt to coerce using google.protobuf.Value, which will expect
    # something that is utf-8 (and base64 consistently is).
    try:
        Value(string_value=bytestring)
        return bytestring
    except ValueError:
        raise ValueError(
            "Received a bytes that is not base64 encoded. "
            "Ensure that you either send a Unicode string or a "
            "base64-encoded bytes."
        )


def _merge_query_options(base, merge):
    """Merge higher precedence QueryOptions with current QueryOptions.

    :type base:
        :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
        or :class:`dict` or None
    :param base: The current QueryOptions that is intended for use.

    :type merge:
        :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
        or :class:`dict` or None
    :param merge:
        The QueryOptions that have a higher priority than base. These options
        should overwrite the fields in base.

    :rtype:
        :class:`~google.cloud.spanner_v1.types.ExecuteSqlRequest.QueryOptions`
        or None
    :returns:
        QueryOptions object formed by merging the two given QueryOptions.
        If the resultant object only has empty fields, returns None.
    """
    combined = base or ExecuteSqlRequest.QueryOptions()
    if isinstance(combined, dict):
        combined = ExecuteSqlRequest.QueryOptions(
            optimizer_version=combined.get("optimizer_version", ""),
            optimizer_statistics_package=combined.get(
                "optimizer_statistics_package", ""
            ),
        )
    merge = merge or ExecuteSqlRequest.QueryOptions()
    if isinstance(merge, dict):
        merge = ExecuteSqlRequest.QueryOptions(
            optimizer_version=merge.get("optimizer_version", ""),
            optimizer_statistics_package=merge.get("optimizer_statistics_package", ""),
        )
    type(combined).pb(combined).MergeFrom(type(merge).pb(merge))
    if not combined.optimizer_version and not combined.optimizer_statistics_package:
        return None
    return combined


def _merge_client_context(base, merge):
    """Merge higher precedence ClientContext with current ClientContext.

    :type base: :class:`~google.cloud.spanner_v1.types.ClientContext`
        or :class:`dict` or None
    :param base: The current ClientContext that is intended for use.

    :type merge: :class:`~google.cloud.spanner_v1.types.ClientContext`
        or :class:`dict` or None
    :param merge:
        The ClientContext that has a higher priority than base. These options
        should overwrite the fields in base.

    :rtype: :class:`~google.cloud.spanner_v1.types.ClientContext`
        or None
    :returns:
        ClientContext object formed by merging the two given ClientContexts.
    """
    if base is None and merge is None:
        return None

    # Avoid in-place modification of base
    combined_pb = ClientContext()._pb
    if base:
        base_pb = ClientContext(base)._pb if isinstance(base, dict) else base._pb
        combined_pb.MergeFrom(base_pb)
    if merge:
        merge_pb = ClientContext(merge)._pb if isinstance(merge, dict) else merge._pb
        combined_pb.MergeFrom(merge_pb)

    combined = ClientContext(combined_pb)

    if not combined.secure_context:
        return None
    return combined


def _validate_client_context(client_context):
    """Validate and convert client_context.

    :type client_context: :class:`~google.cloud.spanner_v1.types.ClientContext`
        or :class:`dict`
    :param client_context: (Optional) Client context to use.

    :rtype: :class:`~google.cloud.spanner_v1.types.ClientContext`
    :returns: Validated ClientContext object or None.
    :raises TypeError: if client_context is not a ClientContext or a dict.
    """
    if client_context is not None:
        if isinstance(client_context, dict):
            client_context = ClientContext(client_context)
        elif not isinstance(client_context, ClientContext):
            raise TypeError("client_context must be a ClientContext or a dict")
    return client_context


def _merge_request_options(request_options, client_context):
    """Merge RequestOptions and ClientContext.

    :type request_options: :class:`~google.cloud.spanner_v1.types.RequestOptions`
        or :class:`dict` or None
    :param request_options: The current RequestOptions that is intended for use.

    :type client_context: :class:`~google.cloud.spanner_v1.types.ClientContext`
        or :class:`dict` or None
    :param client_context:
        The ClientContext to merge into request_options.

    :rtype: :class:`~google.cloud.spanner_v1.types.RequestOptions`
        or None
    :returns:
        RequestOptions object formed by merging the given ClientContext.
    """
    if request_options is None and client_context is None:
        return None

    if request_options is None:
        request_options = RequestOptions()
    elif isinstance(request_options, dict):
        request_options = RequestOptions(request_options)

    if client_context:
        request_options.client_context = _merge_client_context(
            client_context, request_options.client_context
        )

    return request_options


def _assert_numeric_precision_and_scale(value):
    """
    Asserts that input numeric field is within Spanner supported range.

    Spanner supports fixed 38 digits of precision and 9 digits of scale.
    This number can be optionally prefixed with a plus or minus sign.
    Read more here: https://cloud.google.com/spanner/docs/data-types#numeric_type

    :type value: decimal.Decimal
    :param value: The value to check for Cloud Spanner compatibility.

    :raises NotSupportedError: If value is not within supported precision or scale of Spanner.
    """
    scale = value.as_tuple().exponent
    precision = len(value.as_tuple().digits)

    if scale < -9:
        raise ValueError(NUMERIC_MAX_SCALE_ERR_MSG.format(abs(scale)))
    if precision + scale > 29:
        raise ValueError(NUMERIC_MAX_PRECISION_ERR_MSG.format(precision + scale))


def _datetime_to_rfc3339(value):
    """Format the provided datatime in the RFC 3339 format.

    :type value: datetime.datetime
    :param value: value to format

    :rtype: str
    :returns: RFC 3339 formatted datetime string
    """
    # Convert to UTC and then drop the timezone so we can append "Z" in lieu of
    # allowing isoformat to append the "+00:00" zone offset.
    value = value.astimezone(datetime.timezone.utc).replace(tzinfo=None)
    return value.isoformat(sep="T", timespec="microseconds") + "Z"


def _datetime_to_rfc3339_nanoseconds(value):
    """Format the provided datatime in the RFC 3339 format.

    :type value: datetime_helpers.DatetimeWithNanoseconds
    :param value: value to format

    :rtype: str
    :returns: RFC 3339 formatted datetime string
    """

    if value.nanosecond == 0:
        return _datetime_to_rfc3339(value)
    nanos = str(value.nanosecond).rjust(9, "0").rstrip("0")
    # Convert to UTC and then drop the timezone so we can append "Z" in lieu of
    # allowing isoformat to append the "+00:00" zone offset.
    value = value.astimezone(datetime.timezone.utc).replace(tzinfo=None)
    return "{}.{}Z".format(value.isoformat(sep="T", timespec="seconds"), nanos)


def _make_value_pb(value):
    """Helper for :func:`_make_list_value_pbs`.

    :type value: scalar value
    :param value: value to convert

    :rtype: :class:`~google.protobuf.struct_pb2.Value`
    :returns: value protobufs
    :raises ValueError: if value is not of a known scalar type.
    """
    if value is None:
        return Value(null_value="NULL_VALUE")
    if isinstance(value, (list, tuple)):
        return Value(list_value=_make_list_value_pb(value))
    if isinstance(value, bool):
        return Value(bool_value=value)
    if isinstance(value, int):
        return Value(string_value=str(value))
    if isinstance(value, float):
        if math.isnan(value):
            return Value(string_value="NaN")
        if math.isinf(value):
            if value > 0:
                return Value(string_value="Infinity")
            else:
                return Value(string_value="-Infinity")
        return Value(number_value=value)
    if isinstance(value, datetime_helpers.DatetimeWithNanoseconds):
        return Value(string_value=_datetime_to_rfc3339_nanoseconds(value))
    if isinstance(value, datetime.datetime):
        return Value(string_value=_datetime_to_rfc3339(value))
    if isinstance(value, datetime.date):
        return Value(string_value=value.isoformat())
    if isinstance(value, bytes):
        value = _try_to_coerce_bytes(value)
        return Value(string_value=value)
    if isinstance(value, str):
        return Value(string_value=value)
    if isinstance(value, ListValue):
        return Value(list_value=value)
    if isinstance(value, decimal.Decimal):
        _assert_numeric_precision_and_scale(value)
        return Value(string_value=str(value))
    if isinstance(value, JsonObject):
        value = value.serialize()
        if value is None:
            return Value(null_value="NULL_VALUE")
        else:
            return Value(string_value=value)
    if isinstance(value, Message):
        value = value.SerializeToString()
        if value is None:
            return Value(null_value="NULL_VALUE")
        else:
            return Value(string_value=base64.b64encode(value))
    if isinstance(value, Interval):
        return Value(string_value=str(value))
    if isinstance(value, uuid.UUID):
        return Value(string_value=str(value))

    raise ValueError("Unknown type: %s" % (value,))


def _make_list_value_pb(values):
    """Construct of ListValue protobufs.

    :type values: list of scalar
    :param values: Row data

    :rtype: :class:`~google.protobuf.struct_pb2.ListValue`
    :returns: protobuf
    """
    return ListValue(values=[_make_value_pb(value) for value in values])


def _make_list_value_pbs(values):
    """Construct a sequence of ListValue protobufs.

    :type values: list of list of scalar
    :param values: Row data

    :rtype: list of :class:`~google.protobuf.struct_pb2.ListValue`
    :returns: sequence of protobufs
    """
    return [_make_list_value_pb(row) for row in values]


def _parse_value_pb(value_pb, field_type, field_name, column_info=None):
    """Convert a Value protobuf to cell data.

    :type value_pb: :class:`~google.protobuf.struct_pb2.Value`
    :param value_pb: protobuf to convert

    :type field_type: :class:`~google.cloud.spanner_v1.types.Type`
    :param field_type: type code for the value

    :type field_name: str
    :param field_name: column name

    :type column_info: dict
    :param column_info: (Optional) dict of column name and column information.
            An object where column names as keys and custom objects as corresponding
            values for deserialization. It's specifically useful for data types like
            protobuf where deserialization logic is on user-specific code. When provided,
            the custom object enables deserialization of backend-received column data.
            If not provided, data remains serialized as bytes for Proto Messages and
            integer for Proto Enums.

    :rtype: varies on field_type
    :returns: value extracted from value_pb
    :raises ValueError: if unknown type is passed
    """
    decoder = _get_type_decoder(field_type, field_name, column_info)
    return _parse_nullable(value_pb, decoder)


def _get_type_decoder(field_type, field_name, column_info=None):
    """Returns a function that converts a Value protobuf to cell data.

    :type field_type: :class:`~google.cloud.spanner_v1.types.Type`
    :param field_type: type code for the value

    :type field_name: str
    :param field_name: column name

    :type column_info: dict
    :param column_info: (Optional) dict of column name and column information.
            An object where column names as keys and custom objects as corresponding
            values for deserialization. It's specifically useful for data types like
            protobuf where deserialization logic is on user-specific code. When provided,
            the custom object enables deserialization of backend-received column data.
            If not provided, data remains serialized as bytes for Proto Messages and
            integer for Proto Enums.

    :rtype: a function that takes a single protobuf value as an input argument
    :returns: a function that can be used to extract a value from a protobuf value
    :raises ValueError: if unknown type is passed
    """

    type_code = field_type.code
    if type_code == TypeCode.STRING:
        return _parse_string
    elif type_code == TypeCode.BYTES:
        return _parse_bytes
    elif type_code == TypeCode.BOOL:
        return _parse_bool
    elif type_code == TypeCode.INT64:
        return _parse_int64
    elif type_code == TypeCode.FLOAT64:
        return _parse_float
    elif type_code == TypeCode.FLOAT32:
        return _parse_float
    elif type_code == TypeCode.DATE:
        return _parse_date
    elif type_code == TypeCode.TIMESTAMP:
        return _parse_timestamp
    elif type_code == TypeCode.NUMERIC:
        return _parse_numeric
    elif type_code == TypeCode.JSON:
        return _parse_json
    elif type_code == TypeCode.UUID:
        return _parse_uuid
    elif type_code == TypeCode.PROTO:
        return lambda value_pb: _parse_proto(value_pb, column_info, field_name)
    elif type_code == TypeCode.ENUM:
        return lambda value_pb: _parse_proto_enum(value_pb, column_info, field_name)
    elif type_code == TypeCode.ARRAY:
        element_decoder = _get_type_decoder(
            field_type.array_element_type, field_name, column_info
        )
        return lambda value_pb: _parse_array(value_pb, element_decoder)
    elif type_code == TypeCode.STRUCT:
        element_decoders = [
            _get_type_decoder(item_field.type_, field_name, column_info)
            for item_field in field_type.struct_type.fields
        ]
        return lambda value_pb: _parse_struct(value_pb, element_decoders)
    elif type_code == TypeCode.INTERVAL:
        return _parse_interval
    else:
        raise ValueError("Unknown type: %s" % (field_type,))


def _parse_list_value_pbs(rows, row_type):
    """Convert a list of ListValue protobufs into a list of list of cell data.

    :type rows: list of :class:`~google.protobuf.struct_pb2.ListValue`
    :param rows: row data returned from a read/query

    :type row_type: :class:`~google.cloud.spanner_v1.types.StructType`
    :param row_type: row schema specification

    :rtype: list of list of cell data
    :returns: data for the rows, coerced into appropriate types
    """
    result = []
    for row in rows:
        row_data = []
        for value_pb, field in zip(row.values, row_type.fields):
            row_data.append(_parse_value_pb(value_pb, field.type_, field.name))
        result.append(row_data)
    return result


def _parse_string(value_pb) -> str:
    return value_pb.string_value


def _parse_bytes(value_pb):
    return value_pb.string_value.encode("utf8")


def _parse_bool(value_pb) -> bool:
    return value_pb.bool_value


def _parse_int64(value_pb) -> int:
    return int(value_pb.string_value)


def _parse_float(value_pb) -> float:
    if value_pb.HasField("string_value"):
        return float(value_pb.string_value)
    else:
        return value_pb.number_value


def _parse_date(value_pb):
    return _date_from_iso8601_date(value_pb.string_value)


def _parse_timestamp(value_pb):
    DatetimeWithNanoseconds = datetime_helpers.DatetimeWithNanoseconds
    return DatetimeWithNanoseconds.from_rfc3339(value_pb.string_value)


def _parse_numeric(value_pb):
    return decimal.Decimal(value_pb.string_value)


def _parse_json(value_pb):
    return JsonObject.from_str(value_pb.string_value)


def _parse_uuid(value_pb):
    return uuid.UUID(value_pb.string_value)


def _parse_proto(value_pb, column_info, field_name):
    bytes_value = base64.b64decode(value_pb.string_value)
    if column_info is not None and column_info.get(field_name) is not None:
        default_proto_message = column_info.get(field_name)
        if isinstance(default_proto_message, Message):
            proto_message = type(default_proto_message)()
            proto_message.ParseFromString(bytes_value)
            return proto_message
    return bytes_value


def _parse_proto_enum(value_pb, column_info, field_name):
    int_value = int(value_pb.string_value)
    if column_info is not None and column_info.get(field_name) is not None:
        proto_enum = column_info.get(field_name)
        if isinstance(proto_enum, EnumTypeWrapper):
            return proto_enum.Name(int_value)
    return int_value


def _parse_array(value_pb, element_decoder) -> []:
    return [
        _parse_nullable(item_pb, element_decoder)
        for item_pb in value_pb.list_value.values
    ]


def _parse_struct(value_pb, element_decoders):
    return [
        _parse_nullable(item_pb, element_decoders[i])
        for (i, item_pb) in enumerate(value_pb.list_value.values)
    ]


def _parse_nullable(value_pb, decoder):
    if value_pb.HasField("null_value"):
        return None
    else:
        return decoder(value_pb)


def _parse_interval(value_pb):
    """Parse a Value protobuf containing an interval."""
    if hasattr(value_pb, "string_value"):
        return Interval.from_str(value_pb.string_value)
    return Interval.from_str(value_pb)


class _SessionWrapper(object):
    """Base class for objects wrapping a session.

    :type session: :class:`~google.cloud.spanner_v1.session.Session`
    :param session: the session used to perform the commit
    """

    def __init__(self, session):
        self._session = session


def _metadata_with_prefix(prefix, **kw):
    """Create RPC metadata containing a prefix.

    Args:
        prefix (str): appropriate resource path.

    Returns:
        List[Tuple[str, str]]: RPC metadata with supplied prefix
    """
    return [("google-cloud-resource-prefix", prefix)]


def _retry_on_aborted_exception(
    func,
    deadline,
    default_retry_delay=None,
):
    """
    Handles retry logic for Aborted exceptions, considering the deadline.
    """
    attempts = 0
    while True:
        try:
            attempts += 1
            return func()
        except Aborted as exc:
            _delay_until_retry(
                exc,
                deadline=deadline,
                attempts=attempts,
                default_retry_delay=default_retry_delay,
            )
            continue


def _retry(
    func,
    retry_count=5,
    delay=2,
    allowed_exceptions=None,
    before_next_retry=None,
):
    """
    Retry a function with a specified number of retries, delay between retries, and list of allowed exceptions.

    Args:
        func: The function to be retried.
        retry_count: The maximum number of times to retry the function.
        delay: The delay in seconds between retries.
        allowed_exceptions: A tuple of exceptions that are allowed to occur without triggering a retry.
                            Passing allowed_exceptions as None will lead to retrying for all exceptions.

    Returns:
        The result of the function if it is successful, or raises the last exception if all retries fail.
    """
    retries = 0
    while retries <= retry_count:
        if retries > 0 and before_next_retry:
            before_next_retry(retries, delay)

        try:
            return func()
        except Exception as exc:
            is_allowed = (
                allowed_exceptions is None or exc.__class__ in allowed_exceptions
            )

            if is_allowed and retries < retry_count:
                if (
                    allowed_exceptions is not None
                    and allowed_exceptions[exc.__class__] is not None
                ):
                    allowed_exceptions[exc.__class__](exc)
                time.sleep(delay)
                delay = delay * 2
                retries = retries + 1
            else:
                raise exc


def _check_rst_stream_error(exc):
    resumable_error = (
        any(
            resumable_message in exc.message
            for resumable_message in (
                "RST_STREAM",
                "Received unexpected EOS on DATA frame from server",
            )
        ),
    )
    if not resumable_error:
        raise


def _metadata_with_leader_aware_routing(value, **kw):
    """Create RPC metadata containing a leader aware routing header

    Args:
        value (bool): header value

    Returns:
        List[Tuple[str, str]]: RPC metadata with leader aware routing header
    """
    return ("x-goog-spanner-route-to-leader", str(value).lower())


def _metadata_with_span_context(metadata: List[Tuple[str, str]], **kw) -> None:
    """
    Appends metadata with end to end tracing header and OpenTelemetry span context .

    Args:
        metadata (list[tuple[str, str]]): The metadata carrier where the OpenTelemetry context
                                          should be injected.
    Returns:
        None
    """
    if HAS_OPENTELEMETRY_INSTALLED and metadata is not None:
        metadata.append(("x-goog-spanner-end-to-end-tracing", "true"))
        inject(setter=OpenTelemetryContextSetter(), carrier=metadata)


def _delay_until_retry(exc, deadline, attempts, default_retry_delay=None):
    """Helper for :meth:`Session.run_in_transaction`.

    Detect retryable abort, and impose server-supplied delay.

    :type exc: :class:`google.api_core.exceptions.Aborted`
    :param exc: exception for aborted transaction

    :type deadline: float
    :param deadline: maximum timestamp to continue retrying the transaction.

    :type attempts: int
    :param attempts: number of call retries
    """

    cause = exc.errors[0]
    now = time.time()
    if now >= deadline:
        raise

    delay = _get_retry_delay(cause, attempts, default_retry_delay=default_retry_delay)
    if delay is not None:
        if now + delay > deadline:
            raise

        time.sleep(delay)


def _get_retry_delay(cause, attempts, default_retry_delay=None):
    """Helper for :func:`_delay_until_retry`.

    :type exc: :class:`grpc.Call`
    :param exc: exception for aborted transaction

    :rtype: float
    :returns: seconds to wait before retrying the transaction.

    :type attempts: int
    :param attempts: number of call retries
    """
    if hasattr(cause, "trailing_metadata"):
        metadata = dict(cause.trailing_metadata())
    else:
        metadata = {}
    retry_info_pb = metadata.get("google.rpc.retryinfo-bin")
    if retry_info_pb is not None:
        retry_info = RetryInfo()
        retry_info.ParseFromString(retry_info_pb)
        nanos = retry_info.retry_delay.nanos
        return retry_info.retry_delay.seconds + nanos / 1.0e9
    if default_retry_delay is not None:
        return default_retry_delay

    return 2**attempts + random.random()


class AtomicCounter:
    def __init__(self, start_value=0):
        self.__lock = threading.Lock()
        self.__value = start_value

    @property
    def value(self):
        with self.__lock:
            return self.__value

    def increment(self, n=1):
        with self.__lock:
            self.__value += n
            return self.__value

    def __iadd__(self, n):
        """
        Defines the inplace += operator result.
        """
        with self.__lock:
            self.__value += n
            return self

    def __add__(self, n):
        """
        Defines the result of invoking: value = AtomicCounter + addable
        """
        with self.__lock:
            n += self.__value
            return n

    def __radd__(self, n):
        """
        Defines the result of invoking: value = addable + AtomicCounter
        """
        return self.__add__(n)

    def reset(self):
        with self.__lock:
            self.__value = 0


def _metadata_with_request_id(*args, **kwargs):
    """Return metadata with request ID header.

    This function returns only the metadata list (not a tuple),
    maintaining backward compatibility with existing code.

    Args:
        *args: Arguments to pass to with_request_id
        **kwargs: Keyword arguments to pass to with_request_id

    Returns:
        list: gRPC metadata with request ID header
    """
    return with_request_id_metadata_only(*args, **kwargs)


def _metadata_with_request_id_and_req_id(*args, **kwargs):
    """Return both metadata and request ID string.

    This is used when we need to augment errors with the request ID.

    Args:
        *args: Arguments to pass to with_request_id
        **kwargs: Keyword arguments to pass to with_request_id

    Returns:
        tuple: (metadata, request_id)
    """
    return with_request_id(*args, **kwargs)


def _augment_error_with_request_id(error, request_id=None):
    """Augment an error with request ID information.

    Args:
        error: The error to augment (typically GoogleAPICallError)
        request_id (str): The request ID to include

    Returns:
        The augmented error with request ID information
    """
    return wrap_with_request_id(error, request_id)


@contextmanager
def _augment_errors_with_request_id(request_id):
    """Context manager to augment exceptions with request ID.

    Args:
        request_id (str): The request ID to include in exceptions

    Yields:
        None
    """
    try:
        yield
    except Exception as exc:
        augmented = _augment_error_with_request_id(exc, request_id)
        # Use exception chaining to preserve the original exception
        raise augmented from exc


def _merge_Transaction_Options(
    defaultTransactionOptions: TransactionOptions,
    mergeTransactionOptions: TransactionOptions,
) -> TransactionOptions:
    """Merges two TransactionOptions objects.

    - Values from `mergeTransactionOptions` take precedence if set.
    - Values from `defaultTransactionOptions` are used only if missing.

    Args:
        defaultTransactionOptions (TransactionOptions): The default transaction options (fallback values).
        mergeTransactionOptions (TransactionOptions): The main transaction options (overrides when set).

    Returns:
        TransactionOptions: A merged TransactionOptions object.
    """

    if defaultTransactionOptions is None:
        return mergeTransactionOptions

    if mergeTransactionOptions is None:
        return defaultTransactionOptions

    merged_pb = TransactionOptions()._pb  # Create a new protobuf object

    # Merge defaultTransactionOptions first
    merged_pb.MergeFrom(defaultTransactionOptions._pb)

    # Merge transactionOptions, ensuring it overrides default values
    merged_pb.MergeFrom(mergeTransactionOptions._pb)

    # Convert protobuf object back into a TransactionOptions instance
    return TransactionOptions(merged_pb)


def _create_experimental_host_transport(
    transport_factory,
    experimental_host,
    use_plain_text,
    ca_certificate,
    client_certificate,
    client_key,
    interceptors=None,
):
    """Creates an experimental host transport for Spanner.

    Args:
        transport_factory (type): The transport class to instantiate (e.g.
            `SpannerGrpcTransport`).
        experimental_host (str): The endpoint for the experimental host.
        use_plain_text (bool): Whether to use a plain text (insecure) connection.
        ca_certificate (str): Path to the CA certificate file for TLS.
        client_certificate (str): Path to the client certificate file for mTLS.
        client_key (str): Path to the client key file for mTLS.
        interceptors (list): Optional list of interceptors to add to the channel.

    Returns:
        object: An instance of the transport class created by `transport_factory`.

    Raises:
        ValueError: If TLS/mTLS configuration is invalid.
    """
    import grpc
    from google.auth.credentials import AnonymousCredentials

    channel = None
    if use_plain_text:
        channel = grpc.insecure_channel(target=experimental_host)
    elif ca_certificate:
        with open(ca_certificate, "rb") as f:
            ca_cert = f.read()
        if client_certificate and client_key:
            with open(client_certificate, "rb") as f:
                client_cert = f.read()
            with open(client_key, "rb") as f:
                private_key = f.read()
            ssl_creds = grpc.ssl_channel_credentials(
                root_certificates=ca_cert,
                private_key=private_key,
                certificate_chain=client_cert,
            )
        elif client_certificate or client_key:
            raise ValueError(
                "Both client_certificate and client_key must be provided for mTLS connection"
            )
        else:
            ssl_creds = grpc.ssl_channel_credentials(root_certificates=ca_cert)
        channel = grpc.secure_channel(experimental_host, ssl_creds)
    else:
        raise ValueError(
            "TLS/mTLS connection requires ca_certificate to be set for experimental_host"
        )
    if interceptors is not None:
        channel = grpc.intercept_channel(channel, *interceptors)
    return transport_factory(channel=channel, credentials=AnonymousCredentials())