Refactor tracing utils

Signed-off-by: Paul Van Eck <paulvaneck@microsoft.com>
Azure · Apr 26, 2023 · 14c0aa4 · 14c0aa4
1 parent 547f912
commit 14c0aa4
Show file tree

Hide file tree

Showing 12 changed files with 376 additions and 375 deletions.
diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/constants.py
@@ -137,44 +137,6 @@
     "ServiceBusDlqSupplementaryAuthorization"
 )
 
-# Distributed Tracing Constants
-
-TRACE_COMPONENT_PROPERTY = "component"
-TRACE_COMPONENT = "servicebus"
-
-TRACE_NAMESPACE_ATTRIBUTE = "az.namespace"
-TRACE_NAMESPACE = "Microsoft.ServiceBus"
-
-SPAN_NAME_RECEIVE = "ServiceBus.receive"
-SPAN_NAME_RECEIVE_DEFERRED = "ServiceBus.receive_deferred"
-SPAN_NAME_PEEK = "ServiceBus.peek"
-SPAN_NAME_SEND = "ServiceBus.send"
-SPAN_NAME_SCHEDULE = "ServiceBus.schedule"
-SPAN_NAME_MESSAGE = "ServiceBus.message"
-
-SPAN_ENQUEUED_TIME_PROPERTY = "enqueuedTime"
-
-TRACE_ENQUEUED_TIME_PROPERTY = b"x-opt-enqueued-time"
-TRACE_PARENT_PROPERTY = b"traceparent"
-TRACE_STATE_PROPERTY = b"tracestate"
-TRACE_PROPERTY_ENCODING = "ascii"
-
-TRACE_MESSAGING_SYSTEM_ATTRIBUTE = "messaging.system"
-TRACE_MESSAGING_SYSTEM = "servicebus"
-
-TRACE_NET_PEER_NAME_ATTRIBUTE = "net.peer.name"
-TRACE_MESSAGING_DESTINATION_ATTRIBUTE = "messaging.destination.name"
-TRACE_MESSAGING_SOURCE_ATTRIBUTE = "messaging.source.name"
-TRACE_MESSAGING_OPERATION_ATTRIBUTE = "messaging.operation"
-TRACE_MESSAGING_BATCH_COUNT_ATTRIBUTE = "messaging.batch.message_count"
-
-DIAGNOSTIC_ID_PROPERTY = b"Diagnostic-Id"
-
-class TraceOperationTypes(str, Enum, metaclass=CaseInsensitiveEnumMeta):
-    PUBLISH = "publish"
-    RECEIVE = "receive"
-    SETTLE = "settle"
-
 MAX_MESSAGE_LENGTH_BYTES = 1024 * 1024  # Backcompat with uAMQP
 MESSAGE_PROPERTY_MAX_LENGTH = 128
 # .NET TimeSpan.MaxValue: 10675199.02:48:05.4775807

diff --git a/sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py b/sdk/servicebus/azure-servicebus/azure/servicebus/_common/message.py
@@ -34,7 +34,7 @@
     MAX_ABSOLUTE_EXPIRY_TIME,
     MAX_DURATION_VALUE,
     MAX_MESSAGE_LENGTH_BYTES,
-    MESSAGE_STATE_NAME,
+    MESSAGE_STATE_NAME
 )
 from ..amqp import (
     AmqpAnnotatedMessage,
@@ -46,9 +46,9 @@
 from .utils import (
     utc_from_timestamp,
     utc_now,
-    trace_message,
     transform_outbound_messages,
 )
+from .tracing import trace_message
 
 if TYPE_CHECKING:
     try:
@@ -60,7 +60,6 @@
     except ImportError:
         pass
     from .._pyamqp.performatives import TransferFrame
-    from azure.core.tracing import AbstractSpan
     from ..aio._servicebus_receiver_async import (
         ServiceBusReceiver as AsyncServiceBusReceiver,
     )
@@ -659,6 +658,7 @@ def __init__(
         **kwargs: Any
     ) -> None:
         self._amqp_transport = kwargs.pop("amqp_transport", PyamqpTransport)
+        self._tracing_attributes: Dict[str, Union[str, int]] = kwargs.pop("tracing_attributes", {})
 
         self._max_size_in_bytes = max_size_in_bytes or MAX_MESSAGE_LENGTH_BYTES
         self._message = self._amqp_transport.build_batch_message([])
@@ -667,8 +667,6 @@ def __init__(
         self._messages: List[ServiceBusMessage] = []
         self._uamqp_message: Optional[LegacyBatchMessage] = None
 
-        self._tracing_attributes: Dict[str, Union[str, int]] = {}
-
     def __repr__(self) -> str:
         batch_repr = "max_size_in_bytes={}, message_count={}".format(
             self.max_size_in_bytes, self._count