from ddtrace import tracer
from ddtrace.propagation.http import HTTPPropagator
from ddtrace.data_streams import set_produce_checkpoint

# When publishing:
attributes: dict[str, str] = {}

# Inject trace context into message attributes
current_span = tracer.current_span()
if current_span:
    HTTPPropagator.inject(current_span.context, attributes)

# Set DSM produce checkpoint
set_produce_checkpoint("pubsub", TOPIC, attributes.setdefault)

# Trace the publish operation
with tracer.trace("pubsub.publish", resource=TOPIC) as span:
    span.set_tag("messaging.system", "pubsub")
    span.set_tag("messaging.destination", TOPIC)
    future = publisher.publish(topic_path, data=body, **attributes)
    message_id = future.result()
    span.set_tag("messaging.message_id", message_id)

from ddtrace import tracer
from ddtrace.propagation.http import HTTPPropagator
from ddtrace.data_streams import set_consume_checkpoint

# When receiving:
attributes = dict(message.attributes)

# Extract trace context from message attributes
ctx = HTTPPropagator.extract(attributes)

# Set DSM consume checkpoint
set_consume_checkpoint("pubsub", SUBSCRIPTION, attributes.get)

# Activate the extracted context
if ctx:
    tracer.context_provider.activate(ctx)

try:
    # Trace the consume operation
    with tracer.trace("pubsub.consume", resource=SUBSCRIPTION) as consume_span:
        consume_span.set_tag("messaging.system", "pubsub")
        consume_span.set_tag("messaging.message_id", message.message_id)
        
        payload = json.loads(message.data.decode("utf-8"))
        
        # Trace your processing logic
        with tracer.trace("pubsub.process", resource=SUBSCRIPTION):
            # process message
            pass
            
        message.ack()
finally:
    if ctx:
        tracer.context_provider.activate(None)