Attaching a Unit to Traces

from unomiq.sdk import Unomiq

Unomiq.init(
    app_name="my-service",
    resolve_unit=lambda: "my-unit-id",
    resolve_parent_unit=lambda: "my-org-id",
    resource_attributes={"service.name": "my-service"},
)

from flask import request
from unomiq.sdk import Unomiq

FALLBACK_UNIT = "default-unit"

def resolve_unit_from_request():
    try:
        # From a path parameter
        customer_id = (request.view_args or {}).get("customer_id")
        if customer_id is not None:
            return str(customer_id)

        # Or from a request header
        return request.headers.get("X-Customer-Id") or FALLBACK_UNIT
    except RuntimeError:
        return FALLBACK_UNIT

Unomiq.init(
    app_name="my-service",
    resolve_unit=resolve_unit_from_request,
    resolve_parent_unit=lambda: "my-org-id",
    resource_attributes={"service.name": "my-service"},
)

from opentelemetry import trace
from unomiq.sdk import set_unit, set_parent_unit

@app.route('/customers/<int:customer_id>', methods=['GET'])
def get_customer(customer_id):
    span = trace.get_current_span()
    set_unit(span, str(customer_id))
    set_parent_unit(span, "my-org-id")
    # ... rest of handler

Unomiq.init(
    app_name="my-service",
    resource_attributes={"service.name": "my-service"},
)

from opentelemetry import trace

tracer = trace.get_tracer(__name__)

with tracer.start_as_current_span("process-request") as span:
    span.set_attribute("unomiq.unit", user_id)
    span.set_attribute("unomiq.parent_unit", tenant_id)  # optional: parent entity
    # ... your code here

from opentelemetry.sdk.trace import SpanProcessor


class UnitSpanProcessor(SpanProcessor):
    """Attaches unomiq.unit to every span from a context resolver."""

    def __init__(self, resolve_unit):
        """
        Args:
            resolve_unit: A callable that returns the current unit value,
                          or None if no unit is available.
        """
        self._resolve_unit = resolve_unit

    def on_start(self, span, parent_context=None):
        unit = self._resolve_unit()
        if unit:
            span.set_attribute("unomiq.unit", unit)

    def on_end(self, span):
        pass

    def shutdown(self):
        pass

    def force_flush(self, timeout_millis=None):
        pass

from flask import request, Flask
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider

app = Flask(__name__)


def resolve_unit_from_request():
    """Extract the unit from the current request context."""
    try:
        # Option 1: From a request header
        return request.headers.get("X-Customer-Id")

        # Option 2: From a JWT claim (after auth middleware)
        # return g.current_user.get("tenant_id")

        # Option 3: From a path parameter
        # return request.view_args.get("customer_id")
    except RuntimeError:
        # Outside of a request context
        return None


# Register the processor with the tracer provider
provider = TracerProvider()
provider.add_span_processor(UnitSpanProcessor(resolve_unit_from_request))
trace.set_tracer_provider(provider)

from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider

# Use threading.local or contextvars to store the unit per-request
import contextvars

_current_unit = contextvars.ContextVar("current_unit", default=None)


# In your Django middleware:
class UnitMiddleware:
    def __init__(self, get_response):
        self.get_response = get_response

    def __call__(self, request):
        _current_unit.set(request.headers.get("X-Customer-Id"))
        return self.get_response(request)


# Register the processor
provider = TracerProvider()
provider.add_span_processor(UnitSpanProcessor(lambda: _current_unit.get()))
trace.set_tracer_provider(provider)