1# Copyright 2023 gRPC authors.
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15import abc
16from typing import Callable, Dict, Iterable, List, Optional
17
18# pytype: disable=pyi-error
19from grpc_observability import _open_telemetry_observability
20from opentelemetry.metrics import MeterProvider
21
22
23class OpenTelemetryLabelInjector(abc.ABC):
24    """
25    An interface that allows you to add additional labels on the calls traced.
26
27    Please note that this class is still work in progress and NOT READY to be used.
28    """
29
30    _labels: List[Dict[str, str]]
31
32    def __init__(self):
33        # Calls Python OTel API to detect resource and get labels, save
34        # those lables to OpenTelemetryLabelInjector.labels.
35        pass
36
37    @abc.abstractmethod
38    def get_labels(self):
39        # Get additional labels for this OpenTelemetryLabelInjector.
40        raise NotImplementedError()
41
42
43class OpenTelemetryPluginOption(abc.ABC):
44    """
45    An interface that allows you to add additional function to OpenTelemetryPlugin.
46
47    Please note that this class is still work in progress and NOT READY to be used.
48    """
49
50    @abc.abstractmethod
51    def is_active_on_method(self, method: str) -> bool:
52        """Determines whether this plugin option is active on a given method.
53
54        Args:
55          method: Required. The RPC method, for example: `/helloworld.Greeter/SayHello`.
56
57        Returns:
58          True if this this plugin option is active on the giving method, false otherwise.
59        """
60        raise NotImplementedError()
61
62    @abc.abstractmethod
63    def is_active_on_server(self, channel_args: List[str]) -> bool:
64        """Determines whether this plugin option is active on a given server.
65
66        Args:
67          channel_args: Required. The channel args used for server.
68          TODO(xuanwn): detail on what channel_args will contain.
69
70        Returns:
71          True if this this plugin option is active on the server, false otherwise.
72        """
73        raise NotImplementedError()
74
75    @abc.abstractmethod
76    def get_label_injector(self) -> Optional[OpenTelemetryLabelInjector]:
77        # Returns the LabelsInjector used by this plugin option, or None.
78        raise NotImplementedError()
79
80
81# pylint: disable=no-self-use
82class OpenTelemetryPlugin:
83    """Describes a Plugin for OpenTelemetry observability."""
84
85    plugin_options: Iterable[OpenTelemetryPluginOption]
86    meter_provider: Optional[MeterProvider]
87    target_attribute_filter: Callable[[str], bool]
88    generic_method_attribute_filter: Callable[[str], bool]
89    _plugin: _open_telemetry_observability._OpenTelemetryPlugin
90
91    def __init__(
92        self,
93        *,
94        plugin_options: Iterable[OpenTelemetryPluginOption] = [],
95        meter_provider: Optional[MeterProvider] = None,
96        target_attribute_filter: Optional[Callable[[str], bool]] = None,
97        generic_method_attribute_filter: Optional[Callable[[str], bool]] = None,
98    ):
99        """
100        Args:
101          plugin_options: An Iterable of OpenTelemetryPluginOption which will be
102        enabled for this OpenTelemetryPlugin.
103          meter_provider: A MeterProvider which will be used to collect telemetry data,
104        or None which means no metrics will be collected.
105          target_attribute_filter: Once provided, this will be called per channel to decide
106        whether to record the target attribute on client or to replace it with "other".
107        This helps reduce the cardinality on metrics in cases where many channels
108        are created with different targets in the same binary (which might happen
109        for example, if the channel target string uses IP addresses directly).
110        Return True means the original target string will be used, False means target string
111        will be replaced with "other".
112          generic_method_attribute_filter: Once provided, this will be called with a generic
113        method type to decide whether to record the method name or to replace it with
114        "other". Note that pre-registered methods will always be recorded no matter what
115        this function returns.
116        Return True means the original method name will be used, False means method name will
117        be replaced with "other".
118        """
119        self.plugin_options = plugin_options
120        self.meter_provider = meter_provider
121        if target_attribute_filter:
122            self.target_attribute_filter = target_attribute_filter
123        else:
124            self.target_attribute_filter = lambda target: True
125        if generic_method_attribute_filter:
126            self.generic_method_attribute_filter = (
127                generic_method_attribute_filter
128            )
129        else:
130            self.generic_method_attribute_filter = lambda method: False
131        self._plugin = _open_telemetry_observability._OpenTelemetryPlugin(self)
132
133    def register_global(self) -> None:
134        """
135        Registers a global plugin that acts on all channels and servers running on the process.
136
137        Raises:
138            RuntimeError: If a global plugin was already registered.
139        """
140        _open_telemetry_observability.start_open_telemetry_observability(
141            plugins=[self._plugin]
142        )
143
144    def deregister_global(self) -> None:
145        """
146        De-register the global plugin that acts on all channels and servers running on the process.
147
148        Raises:
149            RuntimeError: If no global plugin was registered.
150        """
151        _open_telemetry_observability.end_open_telemetry_observability()
152
153    def __enter__(self) -> None:
154        _open_telemetry_observability.start_open_telemetry_observability(
155            plugins=[self._plugin]
156        )
157
158    def __exit__(self, exc_type, exc_val, exc_tb) -> None:
159        _open_telemetry_observability.end_open_telemetry_observability()
160