o11y/tests/fixtures/traces.py

import datetime
import hashlib
import json
import secrets
import uuid
from abc import ABC
from collections.abc import Callable, Generator
from enum import Enum
from typing import Any, Literal
from urllib.parse import urlparse

import numpy as np
import pytest

from fixtures import types
from fixtures.fingerprint import LogsOrTracesFingerprint
from fixtures.time import parse_duration, parse_timestamp


class TracesKind(Enum):
    SPAN_KIND_UNSPECIFIED = 0
    SPAN_KIND_INTERNAL = 1
    SPAN_KIND_SERVER = 2
    SPAN_KIND_CLIENT = 3
    SPAN_KIND_PRODUCER = 4
    SPAN_KIND_CONSUMER = 5

    @classmethod
    def from_value(cls, value: int) -> "TracesKind":
        return cls(value)


class TracesStatusCode(Enum):
    STATUS_CODE_UNSET = 0
    STATUS_CODE_OK = 1
    STATUS_CODE_ERROR = 2

    @classmethod
    def from_value(cls, value: int) -> "TracesStatusCode":
        return cls(value)


class TracesRefType(Enum):
    REF_TYPE_CHILD_OF = "CHILD_OF"
    REF_TYPE_FOLLOWS_FROM = "FOLLOWS_FROM"


class TraceIdGenerator(ABC):
    @staticmethod
    def trace_id() -> str:
        return secrets.token_hex(16)

    @staticmethod
    def span_id() -> str:
        return secrets.token_hex(8)


class TracesResource(ABC):
    labels: str
    fingerprint: str
    seen_at_ts_bucket_start: np.int64

    def __init__(
        self,
        labels: dict[str, str],
        fingerprint: str,
        seen_at_ts_bucket_start: np.int64,
    ) -> None:
        self.labels = json.dumps(labels, separators=(",", ":"))  # clickhouse treats {"a": "b"} differently from {"a":"b"}. In the first case it is not able to run json functions
        self.fingerprint = fingerprint
        self.seen_at_ts_bucket_start = seen_at_ts_bucket_start

    def np_arr(self) -> np.array:
        return np.array([self.labels, self.fingerprint, self.seen_at_ts_bucket_start])


class TracesResourceOrAttributeKeys(ABC):
    name: str
    datatype: str
    tag_type: str
    is_column: bool

    def __init__(self, name: str, datatype: str, tag_type: str, is_column: bool = False) -> None:
        self.name = name
        self.datatype = datatype
        self.tag_type = tag_type
        self.is_column = is_column

    def np_arr(self) -> np.array:
        return np.array([self.name, self.tag_type, self.datatype, self.is_column], dtype=object)


class TracesTagAttributes(ABC):
    unix_milli: np.int64
    tag_key: str
    tag_type: str
    tag_data_type: str
    string_value: str
    number_value: np.float64

    def __init__(
        self,
        timestamp: datetime.datetime,
        tag_key: str,
        tag_type: str,
        tag_data_type: str,
        string_value: str | None,
        number_value: np.float64,
    ) -> None:
        self.unix_milli = np.int64(int(timestamp.timestamp() * 1e3))
        self.tag_key = tag_key
        self.tag_type = tag_type
        self.tag_data_type = tag_data_type
        self.string_value = string_value or ""
        self.number_value = number_value

    def np_arr(self) -> np.array:
        return np.array(
            [
                self.unix_milli,
                self.tag_key,
                self.tag_type,
                self.tag_data_type,
                self.string_value,
                self.number_value,
            ]
        )


class TracesEvent(ABC):
    name: str
    time_unix_nano: np.uint64
    attribute_map: dict[str, str]

    def __init__(
        self,
        name: str,
        timestamp: datetime.datetime,
        attribute_map: dict[str, str] = {},
    ) -> None:
        self.name = name
        self.time_unix_nano = np.uint64(int(timestamp.timestamp() * 1e9))
        self.attribute_map = attribute_map

    def np_arr(self) -> np.array:
        return np.array([self.name, self.time_unix_nano, json.dumps(self.attribute_map)])


class TracesErrorEvent(ABC):
    event: TracesEvent
    error_id: str
    error_group_id: str

    def __init__(
        self,
        event: TracesEvent,
        error_id: str = "",
        error_group_id: str = "",
    ) -> None:
        self.event = event
        self.error_id = error_id
        self.error_group_id = error_group_id

    def np_arr(self) -> np.array:
        return np.array(
            [
                self.event.time_unix_nano,
                self.error_id,
                self.error_group_id,
                self.event.name,
                json.dumps(self.event.attribute_map),
            ]
        )


class TracesSpanAttribute(ABC):
    key: str
    tag_type: str
    data_type: str
    string_value: str
    number_value: np.float64
    is_column: bool

    def __init__(
        self,
        key: str,
        tag_type: str,
        data_type: str,
        string_value: str = "",
        number_value: np.float64 = None,
        is_column: bool = False,
    ) -> None:
        self.key = key
        self.tag_type = tag_type
        self.data_type = data_type
        self.string_value = string_value
        self.number_value = number_value
        self.is_column = is_column


class TracesLink(ABC):
    trace_id: str
    span_id: str
    ref_type: TracesRefType

    def __init__(self, trace_id: str, span_id: str, ref_type: TracesRefType) -> None:
        self.trace_id = trace_id
        self.span_id = span_id
        self.ref_type = ref_type

    def __dict__(self) -> dict[str, Any]:
        return {
            "traceId": self.trace_id,
            "spanId": self.span_id,
            "refType": self.ref_type.value,
        }


class Traces(ABC):
    ts_bucket_start: np.uint64
    resource_fingerprint: str
    timestamp: np.datetime64
    trace_id: str
    span_id: str
    trace_state: str
    parent_span_id: str
    flags: np.uint32
    name: str
    kind: np.int8
    kind_string: str
    duration_nano: np.uint64
    status_code: np.int16
    status_message: str
    status_code_string: str
    attribute_string: dict[str, str]
    attributes_number: dict[str, np.float64]
    attributes_bool: dict[str, bool]
    resources_string: dict[str, str]
    resource_json: dict[str, str]
    events: list[str]
    links: str
    response_status_code: str
    external_http_url: str
    http_url: str
    external_http_method: str
    http_method: str
    http_host: str
    db_name: str
    db_operation: str
    has_error: bool
    is_remote: str

    resource: list[TracesResource]
    tag_attributes: list[TracesTagAttributes]
    resource_keys: list[TracesResourceOrAttributeKeys]
    attribute_keys: list[TracesResourceOrAttributeKeys]
    span_attributes: list[TracesSpanAttribute]
    error_events: list[TracesErrorEvent]

    def __init__(
        self,
        timestamp: datetime.datetime | None = None,
        duration: datetime.timedelta = datetime.timedelta(seconds=1),
        trace_id: str = "",
        span_id: str = "",
        parent_span_id: str = "",
        name: str = "default span",
        kind: TracesKind = TracesKind.SPAN_KIND_INTERNAL,
        status_code: TracesStatusCode = TracesStatusCode.STATUS_CODE_UNSET,
        status_message: str = "",
        resources: dict[str, Any] = {},
        attributes: dict[str, Any] = {},
        events: list[TracesEvent] = [],
        links: list[TracesLink] = [],
        trace_state: str = "",
        flags: np.uint32 = 0,
        resource_write_mode: Literal["legacy_only", "dual_write"] = "dual_write",
    ) -> None:
        if timestamp is None:
            timestamp = datetime.datetime.now()
        self.tag_attributes = []
        self.attribute_keys = []
        self.resource_keys = []
        self.span_attributes = []
        self.error_events = []

        # Calculate ts_bucket_start (30mins bucket)
        # Round down to nearest 30-minute interval
        minute = timestamp.minute
        if minute < 30:
            bucket_minute = 0
        else:
            bucket_minute = 30

        bucket_start = timestamp.replace(minute=bucket_minute, second=0, microsecond=0)
        self.ts_bucket_start = np.uint64(int(bucket_start.timestamp()))

        self.timestamp = timestamp

        self.duration_nano = np.uint64(int(duration.total_seconds() * 1e9))

        # Initialize trace fields
        self.trace_id = trace_id
        self.span_id = span_id
        self.parent_span_id = parent_span_id
        self.trace_state = trace_state
        self.flags = flags
        self.name = name
        self.kind = kind.value
        self.kind_string = kind.name
        self.status_code = status_code.value
        self.status_message = status_message
        self.status_code_string = status_code.name
        self.has_error = status_code == TracesStatusCode.STATUS_CODE_ERROR
        self.is_remote = self._determine_is_remote(flags)

        # Initialize custom fields to empty values
        self.response_status_code = ""
        self.external_http_url = ""
        self.http_url = ""
        self.external_http_method = ""
        self.http_method = ""
        self.http_host = ""
        self.db_name = ""
        self.db_operation = ""

        # Process resources and derive service_name. Spans written before the
        # JSON-resource evolution time only populate resources_string (legacy_only);
        # spans at or after the evolution time dual-write to both columns.
        self.resources_string = {k: str(v) for k, v in resources.items()}
        self.resource_json = {} if resource_write_mode == "legacy_only" else dict(self.resources_string)
        self.service_name = self.resources_string.get("service.name", "default-service")

        for k, v in self.resources_string.items():
            self.tag_attributes.append(
                TracesTagAttributes(
                    timestamp=timestamp,
                    tag_key=k,
                    tag_type="resource",
                    tag_data_type="string",
                    string_value=v,
                    number_value=None,
                )
            )
            self.resource_keys.append(TracesResourceOrAttributeKeys(name=k, datatype="string", tag_type="resource"))
            self.span_attributes.append(
                TracesSpanAttribute(
                    key=k,
                    tag_type="resource",
                    data_type="string",
                    string_value=v,
                )
            )

        # Calculate resource fingerprint
        self.resource_fingerprint = LogsOrTracesFingerprint(self.resources_string).calculate()

        # Process attributes by type and populate custom fields
        self.attribute_string = {}
        self.attributes_number = {}
        self.attributes_bool = {}

        for k, v in attributes.items():
            # Populate custom fields based on attribute keys (following Go exporter logic)
            self._populate_custom_attrs(k, v)

            if isinstance(v, bool):
                self.attributes_bool[k] = v
                self.tag_attributes.append(
                    TracesTagAttributes(
                        timestamp=timestamp,
                        tag_key=k,
                        tag_type="tag",
                        tag_data_type="bool",
                        string_value=None,
                        number_value=None,
                    )
                )
                self.attribute_keys.append(TracesResourceOrAttributeKeys(name=k, datatype="bool", tag_type="tag"))
                self.span_attributes.append(
                    TracesSpanAttribute(
                        key=k,
                        tag_type="tag",
                        data_type="bool",
                        number_value=None,
                    )
                )
            elif isinstance(v, (int, float)):
                self.attributes_number[k] = np.float64(v)
                self.tag_attributes.append(
                    TracesTagAttributes(
                        timestamp=timestamp,
                        tag_key=k,
                        tag_type="tag",
                        tag_data_type="float64",
                        string_value=None,
                        number_value=np.float64(v),
                    )
                )
                self.attribute_keys.append(TracesResourceOrAttributeKeys(name=k, datatype="float64", tag_type="tag"))
                self.span_attributes.append(
                    TracesSpanAttribute(
                        key=k,
                        tag_type="tag",
                        data_type="float64",
                        number_value=np.float64(v),
                    )
                )
            else:
                self.attribute_string[k] = str(v)
                self.tag_attributes.append(
                    TracesTagAttributes(
                        timestamp=timestamp,
                        tag_key=k,
                        tag_type="tag",
                        tag_data_type="string",
                        string_value=str(v),
                        number_value=None,
                    )
                )
                self.attribute_keys.append(TracesResourceOrAttributeKeys(name=k, datatype="string", tag_type="tag"))
                self.span_attributes.append(
                    TracesSpanAttribute(
                        key=k,
                        tag_type="tag",
                        data_type="string",
                        string_value=str(v),
                    )
                )

        # Process events and derive error events
        self.events = []
        for event in events:
            self.events.append(json.dumps([event.name, event.time_unix_nano, event.attribute_map]))

            # Create error events for exception events (following Go exporter logic)
            if event.name == "exception":
                error_event = self._create_error_event(event)
                self.error_events.append(error_event)

        # In Python, when you define a function with a mutable default argument (like a list []), that default object is created once when the function is defined,
        # not each time the function is called.
        # This means all calls to the function share the same default object.
        # https://stackoverflow.com/questions/1132941/least-astonishment-in-python-the-mutable-default-argument
        links_copy = links.copy() if links else []
        if self.parent_span_id != "":
            links_copy.insert(
                0,
                TracesLink(
                    trace_id=self.trace_id,
                    span_id=self.parent_span_id,
                    ref_type=TracesRefType.REF_TYPE_CHILD_OF,
                ),
            )

        self.links = json.dumps([link.__dict__() for link in links_copy], separators=(",", ":"))

        # Initialize resource
        self.resource = []
        self.resource.append(
            TracesResource(
                labels=self.resources_string,
                fingerprint=self.resource_fingerprint,
                seen_at_ts_bucket_start=self.ts_bucket_start,
            )
        )

    def _create_error_event(self, event: TracesEvent) -> TracesErrorEvent:
        """Create error event from exception event (following Go exporter logic)"""
        error_id = str(uuid.uuid4()).replace("-", "")

        # Create error group ID based on exception type and message
        exception_type = event.attribute_map.get("exception.type", "")
        exception_message = event.attribute_map.get("exception.message", "")

        error_group_content = self.service_name + exception_type + exception_message

        error_group_id = hashlib.md5(error_group_content.encode()).hexdigest()

        return TracesErrorEvent(
            event=event,
            error_id=error_id,
            error_group_id=error_group_id,
        )

    def _determine_is_remote(self, flags: np.uint32) -> str:
        """Determine if span is remote based on flags (following Go exporter logic)"""
        has_is_remote_mask = 0x00000100
        is_remote_mask = 0x00000200

        if flags & has_is_remote_mask != 0:
            if flags & is_remote_mask != 0:
                return "yes"

            return "no"
        return "unknown"

    def _populate_custom_attrs(  # pylint: disable=too-many-branches
        self, key: str, value: Any
    ) -> None:
        """Populate custom attributes based on attribute keys (following Go exporter logic)"""
        str_value = str(value)

        if key in ["http.status_code", "http.response.status_code"]:
            # Handle both string/int http status codes
            try:
                status_int = int(str_value)
                self.response_status_code = str(status_int)
            except ValueError:
                self.response_status_code = str_value
        elif key in ["http.url", "url.full"] and self.kind == 3:  # SPAN_KIND_CLIENT
            # For client spans, extract hostname for external URL
            try:
                parsed = urlparse(str_value)
                self.external_http_url = parsed.hostname or str_value
                self.http_url = str_value
            except Exception:  # pylint: disable=broad-exception-caught
                self.external_http_url = str_value
                self.http_url = str_value
        elif key in ["http.method", "http.request.method"] and self.kind == 3:  # SPAN_KIND_CLIENT
            self.external_http_method = str_value
            self.http_method = str_value
        elif key in ["http.url", "url.full"] and self.kind != 3:
            self.http_url = str_value
        elif key in ["http.method", "http.request.method"] and self.kind != 3:
            self.http_method = str_value
        elif key in [
            "http.host",
            "server.address",
            "client.address",
            "http.request.header.host",
        ]:
            self.http_host = str_value
        elif key in ["db.name", "db.namespace"]:
            self.db_name = str_value
        elif key in ["db.operation", "db.operation.name"]:
            self.db_operation = str_value
        elif key == "rpc.grpc.status_code":
            # Handle both string/int status code in GRPC spans
            try:
                status_int = int(str_value)
                self.response_status_code = str(status_int)
            except ValueError:
                self.response_status_code = str_value
        elif key == "rpc.jsonrpc.error_code":
            self.response_status_code = str_value

    def np_arr(self) -> np.array:
        """Return span data as numpy array for database insertion"""
        return np.array(
            [
                self.ts_bucket_start,
                self.resource_fingerprint,
                self.timestamp,
                self.trace_id,
                self.span_id,
                self.trace_state,
                self.parent_span_id,
                self.flags,
                self.name,
                self.kind,
                self.kind_string,
                self.duration_nano,
                self.status_code,
                self.status_message,
                self.status_code_string,
                self.attribute_string,
                self.attributes_number,
                self.attributes_bool,
                self.resources_string,
                self.events,
                self.links,
                self.response_status_code,
                self.external_http_url,
                self.http_url,
                self.external_http_method,
                self.http_method,
                self.http_host,
                self.db_name,
                self.db_operation,
                self.has_error,
                self.is_remote,
                self.resource_json,
            ],
            dtype=object,
        )

    @classmethod
    def from_dict(
        cls,
        data: dict,
    ) -> "Traces":
        """Create a Traces instance from a dict."""
        # parse timestamp from iso format
        timestamp = parse_timestamp(data["timestamp"])
        duration = parse_duration(data.get("duration", "PT1S"))

        kind = TracesKind.from_value(data.get("kind", TracesKind.SPAN_KIND_INTERNAL.value))
        status_code = TracesStatusCode.from_value(data.get("status_code", TracesStatusCode.STATUS_CODE_UNSET.value))

        return cls(
            timestamp=timestamp,
            duration=duration,
            trace_id=data["trace_id"],
            span_id=data["span_id"],
            parent_span_id=data.get("parent_span_id", ""),
            name=data.get("name", "default span"),
            kind=kind,
            status_code=status_code,
            status_message=data.get("status_message", ""),
            resources=data.get("resources", {}),
            attributes=data.get("attributes", {}),
            trace_state=data.get("trace_state", ""),
            flags=data.get("flags", 0),
        )

    @classmethod
    def load_from_file(
        cls,
        file_path: str,
        base_time: datetime.datetime | None = None,
    ) -> list["Traces"]:
        """Load traces from a JSONL file."""

        data_list = []
        with open(file_path, encoding="utf-8") as f:
            for line in f:
                line = line.strip()
                if not line:
                    continue
                data_list.append(json.loads(line))

        if not data_list:
            return []

        # If base_time provided, calculate time offset
        time_offset = datetime.timedelta(0)
        if base_time is not None:
            # Find earliest timestamp
            earliest = None
            for data in data_list:
                ts = parse_timestamp(data["timestamp"])
                if earliest is None or ts < earliest:
                    earliest = ts
            if earliest is not None:
                time_offset = base_time - earliest

        traces = []
        for data in data_list:
            # add time offset to timestamp
            original_ts = parse_timestamp(data["timestamp"])
            duration = parse_duration(data.get("duration", "PT1S"))
            adjusted_ts = original_ts + time_offset
            data["timestamp"] = adjusted_ts.isoformat()
            # parse duration of the span
            data["duration"] = duration
            traces.append(cls.from_dict(data))

        return traces


def insert_traces_to_datastore(conn, traces: list[Traces]) -> None:
    """
    Insert traces into ClickHouse tables following the same logic as the Go exporter.
    Handles insertion into:
    - distributed_o11y_index_v3 (main traces table)
    - distributed_traces_v3_resource (resource fingerprints)
    - distributed_tag_attributes_v2 (tag attributes)
    - distributed_span_attributes_keys (attribute keys)
    - distributed_o11y_error_index_v2 (error events)

    Pure function so the seeder container (tests/seeder/) can reuse the
    exact insert path used by the pytest fixtures. `conn` is a
    clickhouse-connect Client.
    """
    resources: list[TracesResource] = []
    for trace in traces:
        resources.extend(trace.resource)

    if len(resources) > 0:
        conn.insert(
            database="o11y_traces",
            table="distributed_traces_v3_resource",
            data=[resource.np_arr() for resource in resources],
        )

    tag_attributes: list[TracesTagAttributes] = []
    for trace in traces:
        tag_attributes.extend(trace.tag_attributes)

    if len(tag_attributes) > 0:
        conn.insert(
            database="o11y_traces",
            table="distributed_tag_attributes_v2",
            data=[tag_attribute.np_arr() for tag_attribute in tag_attributes],
        )

    attribute_keys: list[TracesResourceOrAttributeKeys] = []
    resource_keys: list[TracesResourceOrAttributeKeys] = []
    for trace in traces:
        attribute_keys.extend(trace.attribute_keys)
        resource_keys.extend(trace.resource_keys)

    if len(attribute_keys) > 0:
        conn.insert(
            database="o11y_traces",
            table="distributed_span_attributes_keys",
            data=[attribute_key.np_arr() for attribute_key in attribute_keys],
        )

    if len(resource_keys) > 0:
        conn.insert(
            database="o11y_traces",
            table="distributed_span_attributes_keys",
            data=[resource_key.np_arr() for resource_key in resource_keys],
        )

    conn.insert(
        database="o11y_traces",
        table="distributed_o11y_index_v3",
        column_names=[
            "ts_bucket_start",
            "resource_fingerprint",
            "timestamp",
            "trace_id",
            "span_id",
            "trace_state",
            "parent_span_id",
            "flags",
            "name",
            "kind",
            "kind_string",
            "duration_nano",
            "status_code",
            "status_message",
            "status_code_string",
            "attributes_string",
            "attributes_number",
            "attributes_bool",
            "resources_string",
            "events",
            "links",
            "response_status_code",
            "external_http_url",
            "http_url",
            "external_http_method",
            "http_method",
            "http_host",
            "db_name",
            "db_operation",
            "has_error",
            "is_remote",
            "resource",
        ],
        data=[trace.np_arr() for trace in traces],
    )

    error_events: list[TracesErrorEvent] = []
    for trace in traces:
        error_events.extend(trace.error_events)

    if len(error_events) > 0:
        conn.insert(
            database="o11y_traces",
            table="distributed_o11y_error_index_v2",
            data=[error_event.np_arr() for error_event in error_events],
        )


_TRACES_TABLES_TO_TRUNCATE = [
    "o11y_index_v3",
    "traces_v3_resource",
    "tag_attributes_v2",
    "span_attributes_keys",
    "o11y_error_index_v2",
]


def truncate_traces_tables(conn, cluster: str) -> None:
    """Truncate all traces tables. Used by the pytest fixture teardown and by
    the seeder's DELETE /telemetry/traces endpoint."""
    for table in _TRACES_TABLES_TO_TRUNCATE:
        conn.query(f"TRUNCATE TABLE o11y_traces.{table} ON CLUSTER '{cluster}' SYNC")


@pytest.fixture(name="insert_traces", scope="function")
def insert_traces(
    datastore: types.TestContainerDatastore,
) -> Generator[Callable[[list[Traces]], None], Any]:
    def _insert_traces(traces: list[Traces]) -> None:
        insert_traces_to_datastore(datastore.conn, traces)

    yield _insert_traces

    truncate_traces_tables(
        datastore.conn,
        datastore.env["O11Y_TELEMETRYSTORE_DATASTORE_CLUSTER"],
    )


@pytest.fixture(name="remove_traces_ttl_and_storage_settings", scope="function")
def remove_traces_ttl_and_storage_settings(o11y: types.O11y):
    """
    Remove any custom TTL settings on traces tables to revert to default retention.
    Also resets storage policy to default by recreating tables if needed.
    """
    tables = [
        "o11y_index_v3",
        "traces_v3_resource",
        "o11y_error_index_v2",
        "usage_explorer",
        "dependency_graph_minutes_v2",
        "trace_summary",
        "span_attributes_keys",
    ]

    for table in tables:
        try:
            o11y.telemetrystore.conn.query(f"ALTER TABLE o11y_traces.{table} ON CLUSTER '{o11y.telemetrystore.env['O11Y_TELEMETRYSTORE_DATASTORE_CLUSTER']}' REMOVE TTL")
            o11y.telemetrystore.conn.query(f"ALTER TABLE o11y_traces.{table} ON CLUSTER '{o11y.telemetrystore.env['O11Y_TELEMETRYSTORE_DATASTORE_CLUSTER']}' RESET SETTING storage_policy;")
        except Exception as e:  # pylint: disable=broad-exception-caught
            print(f"ttl and storage policy reset failed for {table}: {e}")