87 lines
2.3 KiB
Python
87 lines
2.3 KiB
Python
from opentelemetry.instrumentation.lancedb.utils import dont_throw
|
|
from opentelemetry.semconv.trace import SpanAttributes
|
|
|
|
from opentelemetry import context as context_api
|
|
from opentelemetry.instrumentation.utils import (
|
|
_SUPPRESS_INSTRUMENTATION_KEY,
|
|
)
|
|
from opentelemetry.semconv_ai import SpanAttributes as AISpanAttributes
|
|
|
|
|
|
def _with_tracer_wrapper(func):
|
|
"""Helper for providing tracer for wrapper functions."""
|
|
|
|
def _with_tracer(tracer, to_wrap):
|
|
def wrapper(wrapped, instance, args, kwargs):
|
|
return func(tracer, to_wrap, wrapped, instance, args, kwargs)
|
|
|
|
return wrapper
|
|
|
|
return _with_tracer
|
|
|
|
|
|
def _set_span_attribute(span, name, value):
|
|
if value is not None:
|
|
if value != "":
|
|
span.set_attribute(name, value)
|
|
return
|
|
|
|
|
|
@_with_tracer_wrapper
|
|
def _wrap(tracer, to_wrap, wrapped, instance, args, kwargs):
|
|
"""Instruments and calls every function defined in TO_WRAP."""
|
|
if context_api.get_value(_SUPPRESS_INSTRUMENTATION_KEY):
|
|
return wrapped(*args, **kwargs)
|
|
|
|
name = to_wrap.get("span_name")
|
|
with tracer.start_as_current_span(name) as span:
|
|
span.set_attribute(SpanAttributes.DB_SYSTEM, "lancedb")
|
|
span.set_attribute(SpanAttributes.DB_OPERATION, to_wrap.get("method"))
|
|
|
|
if to_wrap.get("method") == "add":
|
|
_set_add_attributes(span, kwargs)
|
|
elif to_wrap.get("method") == "search":
|
|
_set_search_attributes(span, kwargs)
|
|
elif to_wrap.get("method") == "delete":
|
|
_set_delete_attributes(span, kwargs)
|
|
|
|
return_value = wrapped(*args, **kwargs)
|
|
|
|
return return_value
|
|
|
|
|
|
def _encode_query(_query):
|
|
_query_str = None
|
|
if _query:
|
|
_query_str = str(_query)
|
|
|
|
return _query_str
|
|
|
|
|
|
def _count_or_none(obj):
|
|
if obj:
|
|
return len(obj)
|
|
|
|
return None
|
|
|
|
|
|
@dont_throw
|
|
def _set_add_attributes(span, kwargs):
|
|
_set_span_attribute(
|
|
span, AISpanAttributes.MILVUS_INSERT_DATA_COUNT, _count_or_none(kwargs.get("data"))
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def _set_search_attributes(span, kwargs):
|
|
_set_span_attribute(
|
|
span, AISpanAttributes.MILVUS_SEARCH_FILTER, _encode_query(kwargs.get("query"))
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def _set_delete_attributes(span, kwargs):
|
|
_set_span_attribute(
|
|
span, AISpanAttributes.CHROMADB_DELETE_WHERE, kwargs.get("where")
|
|
)
|