Address review comments

Author: mukund-ananthu

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

@@ -1,4 +1,4 @@
-# Copyright 2017, Google LLC All rights reserved.
+# 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.
@@ -14,17 +14,26 @@
 
 from opentelemetry.propagators.textmap import Setter
 
-from google.pubsub_v1 import types as gapic_types
+from google.pubsub_v1 import PubsubMessage
 
 
 class OpenTelemetryContextSetter(Setter):
  """
  Used by Open Telemetry for context propagation.
  """
 
- def set(self, carrier: gapic_types.PubsubMessage, key: str, value: str) -> None:
+ 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

google/cloud/pubsub_v1/open_telemetry/publish_message_wrapper.py

@@ -14,7 +14,6 @@
 
 import sys
 from datetime import datetime
-import warnings
 from typing import Optional
 
 from opentelemetry import trace
@@ -28,7 +27,7 @@
 
 
 class PublishMessageWrapper:
- _OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1.publisher"
+ _OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1"
  _OPEN_TELEMETRY_MESSAGING_SYSTEM: str = "gcp_pubsub"
  _OPEN_TELEMETRY_PUBLISHER_BATCHING = "publisher batching"
 
@@ -58,6 +57,7 @@ def __eq__(self, other): # pragma: NO COVER
 
  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",
@@ -68,7 +68,9 @@ def start_create_span(self, topic: str, ordering_key: str) -> None:
  "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),
+ "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,
@@ -86,12 +88,7 @@ def start_create_span(self, topic: str, ordering_key: str) -> None:
  )
 
  def end_create_span(self, exc: Optional[BaseException] = None) -> None:
- if self._create_span is None: # pragma: NO COVER
- warnings.warn(
- message="publish create span is None. Hence, not ending it",
- category=RuntimeWarning,
- )
- return
+ assert self._create_span is not None
  if exc:
  self._create_span.record_exception(exception=exc)
  self._create_span.set_status(
@@ -101,12 +98,7 @@ def end_create_span(self, exc: Optional[BaseException] = None) -> None:
 
  def start_publisher_flow_control_span(self) -> None:
  tracer = trace.get_tracer(self._OPEN_TELEMETRY_TRACER_NAME)
- if self._create_span is None: # pragma: NO COVER
- warnings.warn(
- message="publish create span is None. Hence, not starting publish flow control span",
- category=RuntimeWarning,
- )
- return
+ assert self._create_span is not None
  with tracer.start_as_current_span(
  name=self._PUBLISH_FLOW_CONTROL,
  kind=trace.SpanKind.INTERNAL,
@@ -118,12 +110,7 @@ def start_publisher_flow_control_span(self) -> None:
  def end_publisher_flow_control_span(
  self, exc: Optional[BaseException] = None
  ) -> None:
- if self._flow_control_span is None: # pragma: NO COVER
- warnings.warn(
- message="publish flow control span is None. Hence, not ending it",
- category=RuntimeWarning,
- )
- return
+ assert self._flow_control_span is not None
  if exc:
  self._flow_control_span.record_exception(exception=exc)
  self._flow_control_span.set_status(
@@ -132,12 +119,7 @@ def end_publisher_flow_control_span(
  self._flow_control_span.end()
 
  def start_publisher_batching_span(self) -> None:
- if self._create_span is None: # pragma: NO COVER
- warnings.warn(
- message="publish create span is None. Hence, not starting publisher batching span",
- category=RuntimeWarning,
- )
- return
+ 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,
@@ -148,12 +130,7 @@ def start_publisher_batching_span(self) -> None:
  self._batching_span = batching_span
 
  def end_publisher_batching_span(self, exc: Optional[BaseException] = None) -> None:
- if self._batching_span is None: # pragma: NO COVER
- warnings.warn(
- message="publisher batching span is None. Hence, not ending it",
- category=RuntimeWarning,
- )
- return
+ assert self._batching_span is not None
  if exc:
  self._batching_span.record_exception(exception=exc)
  self._batching_span.set_status(

google/cloud/pubsub_v1/publisher/_batch/thread.py

@@ -91,7 +91,7 @@ class Batch(base.Batch):
  timeout is used.
  """
 
- _OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1.publisher"
+ _OPEN_TELEMETRY_TRACER_NAME: str = "google.cloud.pubsub_v1"
  _OPEN_TELEMETRY_MESSAGING_SYSTEM: str = "gcp_pubsub"
 
  def __init__(
@@ -246,7 +246,7 @@ def _start_publish_rpc_span(self) -> None:
  for wrapper in self._message_wrappers:
  span = wrapper.create_span
  # Add links only for sampled spans.
- if span.is_recording():
+ if span.get_span_context().trace_flags.sampled:
  links.append(trace.Link(span.get_span_context()))
  topic_short_name = self._topic.split("/")[3]
  with tracer.start_as_current_span(
@@ -266,7 +266,7 @@ def _start_publish_rpc_span(self) -> None:
  ctx = rpc_span.get_span_context()
  for wrapper in self._message_wrappers:
  span = wrapper.create_span
- if span.is_recording():
+ if span.get_span_context().trace_flags.sampled:
  span.add_link(ctx)
  self._rpc_span = rpc_span
 

setup.py

@@ -47,6 +47,7 @@
  "grpcio-status >= 1.33.2",
  "opentelemetry-api",
  "opentelemetry-sdk",
+ "opentelemetry-semantic-conventions",
 ]
 extras = {"libcst": "libcst >= 0.3.10"}
 url = "https://github.com/googleapis/python-pubsub"

tests/unit/pubsub_v1/publisher/batch/test_thread.py

@@ -26,6 +26,7 @@
 import pytest
 
 from opentelemetry import trace
+from opentelemetry.trace import SpanContext
 
 import google.api_core.exceptions
 from google.api_core import gapic_v1
@@ -830,21 +831,29 @@ def test_opentelemetry_commit_sampling(span_exporter):
  message=gapic_types.PubsubMessage(data=b"foo"),
  )
  message1.start_create_span(topic=TOPIC, ordering_key=None)
- message1.create_span.is_recording = mock.Mock(return_value=False)
 
  message2 = PublishMessageWrapper(
  message=gapic_types.PubsubMessage(data=b"bar"),
  )
  message2.start_create_span(topic=TOPIC, ordering_key=None)
 
+ # Mock the 'get_span_context' method to return a mock SpanContext
+ mock_span_context = mock.Mock(spec=SpanContext)
+ mock_span_context.trace_flags.sampled = False
+
  batch.publish(message1)
  batch.publish(message2)
 
  publish_response = gapic_types.PublishResponse(message_ids=["a", "b"])
+
+ # Patch the 'create_span' method to return the mock SpanContext
  with mock.patch.object(
- type(batch.client), "_gapic_publish", return_value=publish_response
+ message1.create_span, "get_span_context", return_value=mock_span_context
  ):
- batch._commit()
+ with mock.patch.object(
+ type(batch.client), "_gapic_publish", return_value=publish_response
+ ):
+ batch._commit()
 
  spans = span_exporter.get_finished_spans()