deepset-ai--haystack/test/tracing/utils.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import contextlib
import dataclasses
import uuid
from collections.abc import Iterator
from typing import Any

from haystack.tracing import Span, Tracer


@dataclasses.dataclass
class SpyingSpan(Span):
    operation_name: str
    parent_span: Span | None = None

    tags: dict[str, Any] = dataclasses.field(default_factory=dict)

    trace_id: str | None = dataclasses.field(default_factory=lambda: str(uuid.uuid4()))
    span_id: str | None = dataclasses.field(default_factory=lambda: str(uuid.uuid4()))

    def set_tag(self, key: str, value: Any) -> None:
        self.tags[key] = value

    def get_correlation_data_for_logs(self) -> dict[str, Any]:
        return {"trace_id": self.trace_id, "span_id": self.span_id}

    def set_content_tag(self, key: str, value: Any) -> None:
        """
        Set a content tag, but only if content tracing is enabled in the tracer.
        """
        self.set_tag(key, value)


class SpyingTracer(Tracer):
    def current_span(self) -> Span | None:
        return self.spans[-1] if self.spans else None

    def __init__(self) -> None:
        self.spans: list[SpyingSpan] = []

    @contextlib.contextmanager
    def trace(
        self, operation_name: str, tags: dict[str, Any] | None = None, parent_span: Span | None = None
    ) -> Iterator[Span]:
        new_span = SpyingSpan(operation_name, parent_span)
        for key, value in (tags or {}).items():
            new_span.set_tag(key, value)

        self.spans.append(new_span)

        yield new_span