interceptors.py

# Copyright 2023     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 collections import defaultdict
import threading

from grpc_interceptor import ClientInterceptor
from google.api_core.exceptions import Aborted


class MethodCountInterceptor(ClientInterceptor):
    """Test interceptor that counts number of times a method is being called."""

    def __init__(self):
        self._counts = defaultdict(int)

    def intercept(self, method, request_or_iterator, call_details):
        """Count number of times a method is being called."""
        self._counts[call_details.method] += 1
        return method(request_or_iterator, call_details)

    def reset(self):
        self._counts = defaultdict(int)


class MethodAbortInterceptor(ClientInterceptor):
    """Test interceptor that throws Aborted exception for a specific method."""

    def __init__(self):
        self._method_to_abort = None
        self._count = 0
        self._max_raise_count = 1
        self._connection = None

    def intercept(self, method, request_or_iterator, call_details):
        if (
            self._count < self._max_raise_count
            and call_details.method == self._method_to_abort
        ):
            self._count += 1
            if self._connection is not None:
                self._connection._transaction.rollback()
            raise Aborted("Thrown from ClientInterceptor for testing")
        return method(request_or_iterator, call_details)

    def set_method_to_abort(self, method_to_abort, connection=None, max_raise_count=1):
        self._method_to_abort = method_to_abort
        self._count = 0
        self._max_raise_count = max_raise_count
        self._connection = connection

    def reset(self):
        """Reset the interceptor to the original state."""
        self._method_to_abort = None
        self._count = 0
        self._connection = None


X_GOOG_REQUEST_ID = "x-goog-spanner-request-id"


class XGoogRequestIDHeaderInterceptor(ClientInterceptor):
    # TODO:(@odeke-em): delete this guard when PR #1367 is merged.
    X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED = False

    def __init__(self):
        self._unary_req_segments = []
        self._stream_req_segments = []
        self.__lock = threading.Lock()

    def intercept(self, method, request_or_iterator, call_details):
        metadata = call_details.metadata
        x_goog_request_id = None
        for key, value in metadata:
            if key == X_GOOG_REQUEST_ID:
                x_goog_request_id = value
                break

        if self.X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED and not x_goog_request_id:
            raise Exception(
                f"Missing {X_GOOG_REQUEST_ID} header in {call_details.method}"
            )

        response_or_iterator = method(request_or_iterator, call_details)
        streaming = getattr(response_or_iterator, "__iter__", None) is not None

        if self.X_GOOG_REQUEST_ID_FUNCTIONALITY_MERGED:
            with self.__lock:
                if streaming:
                    self._stream_req_segments.append(
                        (call_details.method, parse_request_id(x_goog_request_id))
                    )
                else:
                    self._unary_req_segments.append(
                        (call_details.method, parse_request_id(x_goog_request_id))
                    )

        return response_or_iterator

    @property
    def unary_request_ids(self):
        return self._unary_req_segments

    @property
    def stream_request_ids(self):
        return self._stream_req_segments

    def reset(self):
        self._stream_req_segments.clear()
        self._unary_req_segments.clear()


def parse_request_id(request_id_str):
    splits = request_id_str.split(".")
    version, rand_process_id, client_id, channel_id, nth_request, nth_attempt = list(
        map(lambda v: int(v), splits)
    )
    return (
        version,
        rand_process_id,
        client_id,
        channel_id,
        nth_request,
        nth_attempt,
    )