google-cloud-pubsub 2.30.0__py3-none-any.whl

google/cloud/pubsub/__init__.py

@@ -0,0 +1,31 @@
+# -*- coding: utf-8 -*-
+#
+# 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
+#
+#     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.
+
+
+from __future__ import absolute_import
+
+from google.cloud.pubsub_v1 import PublisherClient
+from google.cloud.pubsub_v1 import SubscriberClient
+from google.cloud.pubsub_v1 import SchemaServiceClient
+from google.cloud.pubsub_v1 import types
+
+
+__all__ = (
+    "types",
+    "PublisherClient",
+    "SubscriberClient",
+    "SchemaServiceClient",
+)

google/cloud/pubsub_v1/__init__.py

@@ -0,0 +1,35 @@
+# Copyright 2017, 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 __future__ import absolute_import
+
+from google.cloud.pubsub_v1 import types
+from google.cloud.pubsub_v1 import publisher
+from google.cloud.pubsub_v1 import subscriber
+from google.pubsub_v1.services import schema_service
+
+
+class PublisherClient(publisher.Client):
+    __doc__ = publisher.Client.__doc__
+
+
+class SubscriberClient(subscriber.Client):
+    __doc__ = subscriber.Client.__doc__
+
+
+class SchemaServiceClient(schema_service.client.SchemaServiceClient):
+    __doc__ = schema_service.client.SchemaServiceClient.__doc__
+
+
+__all__ = ("types", "PublisherClient", "SubscriberClient", "SchemaServiceClient")

google/cloud/pubsub_v1/exceptions.py

@@ -0,0 +1,20 @@
+# Copyright 2017, 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 __future__ import absolute_import
+
+from concurrent.futures import TimeoutError
+
+
+__all__ = ("TimeoutError",)

google/cloud/pubsub_v1/futures.py

@@ -0,0 +1,56 @@
+# Copyright 2017, 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 __future__ import absolute_import
+
+import concurrent.futures
+from typing import Any, NoReturn, Optional
+
+import google.api_core.future
+
+
+class Future(concurrent.futures.Future, google.api_core.future.Future):
+    """Encapsulation of the asynchronous execution of an action.
+
+    This object is returned from asychronous Pub/Sub calls, and is the
+    interface to determine the status of those calls.
+
+    This object should not be created directly, but is returned by other
+    methods in this library.
+    """
+
+    def running(self) -> bool:
+        """Return ``True`` if the associated Pub/Sub action has not yet completed."""
+        return not self.done()
+
+    def set_running_or_notify_cancel(self) -> NoReturn:
+        raise NotImplementedError(
+            "Only used by executors from `concurrent.futures` package."
+        )
+
+    def set_result(self, result: Any):
+        """Set the return value of work associated with the future.
+
+        Do not use this method, it should only be used internally by the library and its
+        unit tests.
+        """
+        return super().set_result(result=result)
+
+    def set_exception(self, exception: Optional[BaseException]):
+        """Set the result of the future as being the given exception.
+
+        Do not use this method, it should only be used internally by the library and its
+        unit tests.
+        """
+        return super().set_exception(exception=exception)

google/cloud/pubsub_v1/open_telemetry/__init__.py

@@ -0,0 +1,13 @@
+# 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.

google/cloud/pubsub_v1/open_telemetry/context_propagation.py

@@ -0,0 +1,55 @@
+# 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 typing import Optional, List
+
+from opentelemetry.propagators.textmap import Setter, Getter
+
+from google.pubsub_v1 import PubsubMessage
+
+
+class OpenTelemetryContextSetter(Setter):
+    """
+    Used by Open Telemetry for context propagation.
+    """
+
+    def set(self, carrier: PubsubMessage, key: str, value: str) -> None:
+        """
+        Injects trace context into Pub/Sub message attributes with
+        "googclient_" prefix.
+
+        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.attributes["googclient_" + key] = value
+
+
+class OpenTelemetryContextGetter(Getter):
+    """
+    Used by Open Telemetry for context propagation.
+    """
+
+    def get(self, carrier: PubsubMessage, key: str) -> Optional[List[str]]:
+        if ("googclient_" + key) not in carrier.attributes:
+            return None
+        return [carrier.attributes["googclient_" + key]]
+
+    def keys(self, carrier: PubsubMessage) -> List[str]:
+        return list(map(str, carrier.attributes.keys()))

google/cloud/pubsub_v1/open_telemetry/publish_message_wrapper.py

@@ -0,0 +1,142 @@
+# Copyright 2017, 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.
+
+import sys
+from datetime import datetime
+from typing import Optional
+
+from opentelemetry import trace
+from opentelemetry.trace.propagation import set_span_in_context
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+
+from google.pubsub_v1 import types as gapic_types
+from google.cloud.pubsub_v1.open_telemetry.context_propagation import (
+    OpenTelemetryContextSetter,
+)
+
+
+class PublishMessageWrapper:
+    _OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1"
+    _OPEN_TELEMETRY_MESSAGING_SYSTEM: str = "gcp_pubsub"
+    _OPEN_TELEMETRY_PUBLISHER_BATCHING = "publisher batching"
+
+    _PUBLISH_START_EVENT: str = "publish start"
+    _PUBLISH_FLOW_CONTROL: str = "publisher flow control"
+
+    def __init__(self, message: gapic_types.PubsubMessage):
+        self._message: gapic_types.PubsubMessage = message
+        self._create_span: Optional[trace.Span] = None
+        self._flow_control_span: Optional[trace.Span] = None
+        self._batching_span: Optional[trace.Span] = None
+
+    @property
+    def message(self):
+        return self._message
+
+    @message.setter  # type: ignore[no-redef]  # resetting message value is intentional here
+    def message(self, message: gapic_types.PubsubMessage):
+        self._message = message
+
+    @property
+    def create_span(self):
+        return self._create_span
+
+    def __eq__(self, other):  # pragma: NO COVER
+        """Used for pytest asserts to compare two PublishMessageWrapper objects with the same message"""
+        if isinstance(self, other.__class__):
+            return self.message == other.message
+        return False
+
+    def start_create_span(self, topic: str, ordering_key: str) -> None:
+        tracer = trace.get_tracer(self._OPEN_TELEMETRY_TRACER_NAME)
+        assert len(topic.split("/")) == 4
+        topic_short_name = topic.split("/")[3]
+        with tracer.start_as_current_span(
+            name=f"{topic_short_name} create",
+            attributes={
+                "messaging.system": self._OPEN_TELEMETRY_MESSAGING_SYSTEM,
+                "messaging.destination.name": topic_short_name,
+                "code.function": "publish",
+                "messaging.gcp_pubsub.message.ordering_key": ordering_key,
+                "messaging.operation": "create",
+                "gcp.project_id": topic.split("/")[1],
+                "messaging.message.body.size": sys.getsizeof(
+                    self._message.data
+                ),  # sys.getsizeof() used since the attribute expects size of message body in bytes
+            },
+            kind=trace.SpanKind.PRODUCER,
+            end_on_exit=False,
+        ) as create_span:
+            create_span.add_event(
+                name=self._PUBLISH_START_EVENT,
+                attributes={
+                    "timestamp": str(datetime.now()),
+                },
+            )
+            self._create_span = create_span
+            TraceContextTextMapPropagator().inject(
+                carrier=self._message,
+                setter=OpenTelemetryContextSetter(),
+            )
+
+    def end_create_span(self, exc: Optional[BaseException] = None) -> None:
+        assert self._create_span is not None
+        if exc:
+            self._create_span.record_exception(exception=exc)
+            self._create_span.set_status(
+                trace.Status(status_code=trace.StatusCode.ERROR)
+            )
+        self._create_span.end()
+
+    def start_publisher_flow_control_span(self) -> None:
+        tracer = trace.get_tracer(self._OPEN_TELEMETRY_TRACER_NAME)
+        assert self._create_span is not None
+        with tracer.start_as_current_span(
+            name=self._PUBLISH_FLOW_CONTROL,
+            kind=trace.SpanKind.INTERNAL,
+            context=set_span_in_context(self._create_span),
+            end_on_exit=False,
+        ) as flow_control_span:
+            self._flow_control_span = flow_control_span
+
+    def end_publisher_flow_control_span(
+        self, exc: Optional[BaseException] = None
+    ) -> None:
+        assert self._flow_control_span is not None
+        if exc:
+            self._flow_control_span.record_exception(exception=exc)
+            self._flow_control_span.set_status(
+                trace.Status(status_code=trace.StatusCode.ERROR)
+            )
+        self._flow_control_span.end()
+
+    def start_publisher_batching_span(self) -> None:
+        assert self._create_span is not None
+        tracer = trace.get_tracer(self._OPEN_TELEMETRY_TRACER_NAME)
+        with tracer.start_as_current_span(
+            name=self._OPEN_TELEMETRY_PUBLISHER_BATCHING,
+            kind=trace.SpanKind.INTERNAL,
+            context=set_span_in_context(self._create_span),
+            end_on_exit=False,
+        ) as batching_span:
+            self._batching_span = batching_span
+
+    def end_publisher_batching_span(self, exc: Optional[BaseException] = None) -> None:
+        assert self._batching_span is not None
+        if exc:
+            self._batching_span.record_exception(exception=exc)
+            self._batching_span.set_status(
+                trace.Status(status_code=trace.StatusCode.ERROR)
+            )
+        self._batching_span.end()

google/cloud/pubsub_v1/open_telemetry/subscribe_opentelemetry.py

@@ -0,0 +1,280 @@
+# 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 typing import Optional, List
+from datetime import datetime
+
+from opentelemetry import trace, context
+from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator
+from opentelemetry.trace.propagation import set_span_in_context
+
+from google.cloud.pubsub_v1.open_telemetry.context_propagation import (
+    OpenTelemetryContextGetter,
+)
+from google.pubsub_v1.types import PubsubMessage
+
+_OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1"
+_OPEN_TELEMETRY_MESSAGING_SYSTEM: str = "gcp_pubsub"
+
+
+class SubscribeOpenTelemetry:
+    def __init__(self, message: PubsubMessage):
+        self._message: PubsubMessage = message
+
+        # subscribe span will be initialized by the `start_subscribe_span`
+        # method.
+        self._subscribe_span: Optional[trace.Span] = None
+
+        # subscriber concurrency control span will be initialized by the
+        # `start_subscribe_concurrency_control_span` method.
+        self._concurrency_control_span: Optional[trace.Span] = None
+
+        # scheduler span will be initialized by the
+        # `start_subscribe_scheduler_span` method.
+        self._scheduler_span: Optional[trace.Span] = None
+
+        # This will be set by `start_subscribe_span` method and will be used
+        # for other spans, such as process span.
+        self._subscription_id: Optional[str] = None
+
+        # This will be set by `start_process_span` method.
+        self._process_span: Optional[trace.Span] = None
+
+        # This will be set by `start_subscribe_span` method, if a publisher create span
+        # context was extracted from trace propagation. And will be used by spans like
+        # proces span to add links to the publisher create span.
+        self._publisher_create_span_context: Optional[context.Context] = None
+
+        # This will be set by `start_subscribe_span` method and will be used
+        # for other spans, such as modack span.
+        self._project_id: Optional[str] = None
+
+    @property
+    def subscription_id(self) -> Optional[str]:
+        return self._subscription_id
+
+    @property
+    def project_id(self) -> Optional[str]:
+        return self._project_id
+
+    @property
+    def subscribe_span(self) -> Optional[trace.Span]:
+        return self._subscribe_span
+
+    def start_subscribe_span(
+        self,
+        subscription: str,
+        exactly_once_enabled: bool,
+        ack_id: str,
+        delivery_attempt: int,
+    ) -> None:
+        tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+        parent_span_context = TraceContextTextMapPropagator().extract(
+            carrier=self._message,
+            getter=OpenTelemetryContextGetter(),
+        )
+        self._publisher_create_span_context = parent_span_context
+        split_subscription: List[str] = subscription.split("/")
+        assert len(split_subscription) == 4
+        subscription_short_name = split_subscription[3]
+        self._project_id = split_subscription[1]
+        self._subscription_id = subscription_short_name
+        with tracer.start_as_current_span(
+            name=f"{subscription_short_name} subscribe",
+            context=parent_span_context if parent_span_context else None,
+            kind=trace.SpanKind.CONSUMER,
+            attributes={
+                "messaging.system": _OPEN_TELEMETRY_MESSAGING_SYSTEM,
+                "messaging.destination.name": subscription_short_name,
+                "gcp.project_id": subscription.split("/")[1],
+                "messaging.message.id": self._message.message_id,
+                "messaging.message.body.size": len(self._message.data),
+                "messaging.gcp_pubsub.message.ack_id": ack_id,
+                "messaging.gcp_pubsub.message.ordering_key": self._message.ordering_key,
+                "messaging.gcp_pubsub.message.exactly_once_delivery": exactly_once_enabled,
+                "code.function": "_on_response",
+                "messaging.gcp_pubsub.message.delivery_attempt": delivery_attempt,
+            },
+            end_on_exit=False,
+        ) as subscribe_span:
+            self._subscribe_span = subscribe_span
+
+    def add_subscribe_span_event(self, event: str) -> None:
+        assert self._subscribe_span is not None
+        self._subscribe_span.add_event(
+            name=event,
+            attributes={
+                "timestamp": str(datetime.now()),
+            },
+        )
+
+    def end_subscribe_span(self) -> None:
+        assert self._subscribe_span is not None
+        self._subscribe_span.end()
+
+    def set_subscribe_span_result(self, result: str) -> None:
+        assert self._subscribe_span is not None
+        self._subscribe_span.set_attribute(
+            key="messaging.gcp_pubsub.result",
+            value=result,
+        )
+
+    def start_subscribe_concurrency_control_span(self) -> None:
+        assert self._subscribe_span is not None
+        tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+        with tracer.start_as_current_span(
+            name="subscriber concurrency control",
+            kind=trace.SpanKind.INTERNAL,
+            context=set_span_in_context(self._subscribe_span),
+            end_on_exit=False,
+        ) as concurrency_control_span:
+            self._concurrency_control_span = concurrency_control_span
+
+    def end_subscribe_concurrency_control_span(self) -> None:
+        assert self._concurrency_control_span is not None
+        self._concurrency_control_span.end()
+
+    def start_subscribe_scheduler_span(self) -> None:
+        assert self._subscribe_span is not None
+        tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+        with tracer.start_as_current_span(
+            name="subscriber scheduler",
+            kind=trace.SpanKind.INTERNAL,
+            context=set_span_in_context(self._subscribe_span),
+            end_on_exit=False,
+        ) as scheduler_span:
+            self._scheduler_span = scheduler_span
+
+    def end_subscribe_scheduler_span(self) -> None:
+        assert self._scheduler_span is not None
+        self._scheduler_span.end()
+
+    def start_process_span(self) -> None:
+        assert self._subscribe_span is not None
+        tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+        publish_create_span_link: Optional[trace.Link] = None
+        if self._publisher_create_span_context:
+            publish_create_span: trace.Span = trace.get_current_span(
+                self._publisher_create_span_context
+            )
+            span_context: Optional[
+                trace.SpanContext
+            ] = publish_create_span.get_span_context()
+            publish_create_span_link = (
+                trace.Link(span_context) if span_context else None
+            )
+
+        with tracer.start_as_current_span(
+            name=f"{self._subscription_id} process",
+            attributes={
+                "messaging.system": _OPEN_TELEMETRY_MESSAGING_SYSTEM,
+            },
+            kind=trace.SpanKind.INTERNAL,
+            context=set_span_in_context(self._subscribe_span),
+            links=[publish_create_span_link] if publish_create_span_link else None,
+            end_on_exit=False,
+        ) as process_span:
+            self._process_span = process_span
+
+    def end_process_span(self) -> None:
+        assert self._process_span is not None
+        self._process_span.end()
+
+    def add_process_span_event(self, event: str) -> None:
+        assert self._process_span is not None
+        self._process_span.add_event(
+            name=event,
+            attributes={
+                "timestamp": str(datetime.now()),
+            },
+        )
+
+
+def start_modack_span(
+    subscribe_span_links: List[trace.Link],
+    subscription_id: Optional[str],
+    message_count: int,
+    deadline: float,
+    project_id: Optional[str],
+    code_function: str,
+    receipt_modack: bool,
+) -> trace.Span:
+    assert subscription_id is not None
+    assert project_id is not None
+    tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+    with tracer.start_as_current_span(
+        name=f"{subscription_id} modack",
+        attributes={
+            "messaging.system": _OPEN_TELEMETRY_MESSAGING_SYSTEM,
+            "messaging.batch.message_count": message_count,
+            "messaging.gcp_pubsub.message.ack_deadline": deadline,
+            "messaging.destination.name": subscription_id,
+            "gcp.project_id": project_id,
+            "messaging.operation.name": "modack",
+            "code.function": code_function,
+            "messaging.gcp_pubsub.is_receipt_modack": receipt_modack,
+        },
+        links=subscribe_span_links,
+        kind=trace.SpanKind.CLIENT,
+        end_on_exit=False,
+    ) as modack_span:
+        return modack_span
+
+
+def start_ack_span(
+    subscription_id: str,
+    message_count: int,
+    project_id: str,
+    links: List[trace.Link],
+) -> trace.Span:
+    tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+    with tracer.start_as_current_span(
+        name=f"{subscription_id} ack",
+        attributes={
+            "messaging.system": _OPEN_TELEMETRY_MESSAGING_SYSTEM,
+            "messaging.batch.message_count": message_count,
+            "messaging.operation": "ack",
+            "gcp.project_id": project_id,
+            "messaging.destination.name": subscription_id,
+            "code.function": "ack",
+        },
+        kind=trace.SpanKind.CLIENT,
+        links=links,
+        end_on_exit=False,
+    ) as ack_span:
+        return ack_span
+
+
+def start_nack_span(
+    subscription_id: str,
+    message_count: int,
+    project_id: str,
+    links: List[trace.Link],
+) -> trace.Span:
+    tracer = trace.get_tracer(_OPEN_TELEMETRY_TRACER_NAME)
+    with tracer.start_as_current_span(
+        name=f"{subscription_id} nack",
+        attributes={
+            "messaging.system": _OPEN_TELEMETRY_MESSAGING_SYSTEM,
+            "messaging.batch.message_count": message_count,
+            "messaging.operation": "nack",
+            "gcp.project_id": project_id,
+            "messaging.destination.name": subscription_id,
+            "code.function": "modify_ack_deadline",
+        },
+        kind=trace.SpanKind.CLIENT,
+        links=links,
+        end_on_exit=False,
+    ) as nack_span:
+        return nack_span