diff --git a/drift/core/rust_core_binding.py b/drift/core/rust_core_binding.py index b9b298a..8f550fc 100644 --- a/drift/core/rust_core_binding.py +++ b/drift/core/rust_core_binding.py @@ -9,7 +9,10 @@ import json import logging import os -from typing import Any +from typing import TYPE_CHECKING, Any + +if TYPE_CHECKING: + from .types import PackageType, SpanKind, StatusCode logger = logging.getLogger(__name__) @@ -161,16 +164,16 @@ def build_span_proto_bytes( package_name: str, instrumentation_name: str, submodule_name: str, - package_type: int, + package_type: PackageType, environment: str | None, - kind: int, + kind: SpanKind, input_schema: dict[str, Any], output_schema: dict[str, Any], input_schema_hash: str, output_schema_hash: str, input_value_hash: str, output_value_hash: str, - status_code: int, + status_code: StatusCode, status_message: str, is_pre_app_start: bool, is_root_span: bool, @@ -190,6 +193,14 @@ def build_span_proto_bytes( if binding is None: return None try: + from .types import PackageType, SpanKind, StatusCode + + if not isinstance(package_type, PackageType): + raise TypeError(f"package_type must be PackageType, got {type(package_type).__name__}") + if not isinstance(kind, SpanKind): + raise TypeError(f"kind must be SpanKind, got {type(kind).__name__}") + if not isinstance(status_code, StatusCode): + raise TypeError(f"status_code must be StatusCode, got {type(status_code).__name__}") return binding.build_span_proto_bytes_pyobject( trace_id, span_id, @@ -198,16 +209,16 @@ def build_span_proto_bytes( package_name, instrumentation_name, submodule_name, - package_type, + package_type.value, environment, - kind, + kind.value, input_schema, output_schema, input_schema_hash, output_schema_hash, input_value_hash, output_value_hash, - status_code, + status_code.value, status_message, is_pre_app_start, is_root_span, diff --git a/drift/core/tracing/otel_converter.py b/drift/core/tracing/otel_converter.py index fa39561..3eec5a9 100644 --- a/drift/core/tracing/otel_converter.py +++ b/drift/core/tracing/otel_converter.py @@ -420,9 +420,6 @@ def otel_span_to_clean_span_data( # Convert kind kind = otel_span_kind_to_drift(otel_span.kind) - kind_value = kind.value if hasattr(kind, "value") else kind - status_code_value = status.code.value if hasattr(status.code, "value") else status.code - package_type_value = package_type.value if hasattr(package_type, "value") else package_type span_proto_bytes = build_span_proto_bytes( trace_id=trace_id, span_id=span_id, @@ -431,16 +428,16 @@ def otel_span_to_clean_span_data( package_name=package_name, instrumentation_name=instrumentation_name, submodule_name=submodule_name, - package_type=package_type_value, + package_type=package_type, environment=environment, - kind=kind_value, + kind=kind, input_schema=input_schema, output_schema=output_schema, input_schema_hash=input_schema_hash, output_schema_hash=output_schema_hash, input_value_hash=input_value_hash, output_value_hash=output_value_hash, - status_code=status_code_value, + status_code=status.code, status_message=status.message, is_pre_app_start=is_pre_app_start, is_root_span=is_root_span,