mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-09 03:18:44 +00:00
feat(otel): attribute Prisma database spans to PostgreSQL instead of localhost (#36595)
* feat(otel): attribute Prisma database spans to PostgreSQL instead of localhost Prisma reaches PostgreSQL through a query engine on loopback, so transport instrumentation attributes database waits to localhost and operators cannot tell the work is PostgreSQL or correlate it with the database's own metrics. Datastore service spans now carry db.system.name, db.system, db.operation.name and, for PostgreSQL, server.address, server.port and db.namespace derived from DATABASE_URL, and are emitted as CLIENT spans. Only host, port, database and schema are read, so no credential reaches an exporter. Endpoint attributes are omitted when a read replica is configured, because routing is decided per Prisma call underneath the span. * fix(otel): reject a mis-split DSN authority and name socket-only databases An unencoded '/' in the password truncates the URL authority, so urlparse reports the username as the host and the password tail as the database, which put credential material in db.namespace. Postgres drivers reject that DSN outright, so the only safe reading is no endpoint at all. A hostless 'postgresql:///litellm' is a valid local-socket DSN that Prisma accepts, and it now yields db.namespace with no server address rather than nothing. The default schema is matched case-insensitively, since an unquoted PostgreSQL identifier folds and one deployment must yield one namespace. * fix(otel): keep a non-default schema in db.namespace Prisma quotes the schema name, so a DSN with ?schema=PUBLIC provisions a second schema alongside public rather than reusing it. Observed on a live proxy: the PUBLIC schema came up with its own 70 tables next to public's 78, and a key created under one was not visible under the other. Case-folding the two into a single namespace therefore reported two different schemas as the same database, which is the misattribution this feature exists to remove. Match the default literally. * fix(otel): reject any DSN whose userinfo fell outside the authority An unencoded '#' or '?' in the password sends the tail to the fragment or query, leaving the path empty, so the marker check on the database segment never fired and urlparse's hostname (the database username) was exported as server.address. The stranded userinfo '@' is the general tell for every mis-split, so guard on that instead of enumerating the characters that cause it. * fix(otel): allow an at-sign inside a well-formed DSN query The previous guard rejected any DSN whose userinfo at-sign fell outside the netloc, which also caught libpq parameters that legitimately carry one, so ?application_name=svc@prod and ?user=admin@company.com lost their endpoint attributes. Discriminate instead: a PostgreSQL DSN never has a fragment, its database name cannot hold an unencoded at-sign or slash, and an at-sign in the query is only suspicious when the query did not parse as parameters. * fix(otel): resolve the database endpoint per span instead of once per process The endpoint was cached for the process lifetime on the premise that DATABASE_URL is deployment-static. It is not. The RDS IAM refresh rebuilds the URL from DATABASE_HOST/PORT/NAME/SCHEMA on every rotation, the reconnect path re-reads DATABASE_URL, and the DB-backed environment_variables config overlay sets arbitrary keys post-startup with no blocklist covering DATABASE_*. A process that had genuinely failed over kept exporting the old server. Read the environment per span, which is also what Prisma connects with, so the span can no longer name a different server than the one serving the query; get_secret_str consulted a secret manager first and could diverge from it. Only the parse is memoized, keyed on the URL. * fix(otel): reject a question-mark mis-split whose tail parses as parameters A '?' in a password strands the rest of the authority in the query, and that tail can still parse as key=value, so testing only for an unparseable query let the login through as server.address. One spelling hijacked the host= parameter and put the password suffix there directly. A legitimate at-sign in a query always follows a database path, and a '?'-mis-split never leaves one, so require both. * refactor(otel): drop the DSN parse cache that retained rotated credentials The cache was keyed on the full DATABASE_URL, so up to eight complete DSNs, each carrying a password or a retired IAM token, stayed referenced for the process lifetime and outlived the rotation that replaced them. Nothing reached a span, but a heap dump or crash report would have surfaced them. Parsing costs about four microseconds against a span emission that costs orders of magnitude more, so the cache bought nothing worth that. * fix(otel): avoid a set construction the tightened LIT002 budget rejects * fix(otel): refuse an ambiguous DSN authority instead of guessing at it A password holding both an unencoded slash and a query-like tail defeated all three shape checks: the slash left a clean path carrying the password remainder, the query still parsed as parameters, and no fragment survived. The login went out as server.address, the password's leading digits as server.port and the rest as db.namespace. A DSN whose at-sign sits in a query parameter is indistinguishable from that mis-split by any property of the parse; both leave no userinfo, a host, a port and a path. Guessing wrong publishes a credential fragment, so the ambiguity now resolves to refusing the endpoint. Such a DSN loses server.address and db.namespace and keeps the rest of the span; percent-encoding the at-sign restores them. Also honour port= alongside host=, which libpq pairs and this read ignored. * docs(otel): fix a spliced sentence and a stale cache claim in db_endpoint Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yucheng-berri
GitHub
parent
1de398d9aa
commit
71044bf5ea
8 changed files with 653 additions and 25 deletions
|
|
@ -20,6 +20,7 @@ from litellm.integrations.opentelemetry_utils.gen_ai_semconv import (
|
|||
OTELSemconvCategory,
|
||||
parse_semconv_opt_in,
|
||||
)
|
||||
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
|
||||
from litellm.integrations.otel.model.semconv import Metric
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.litellm_core_utils.secret_redaction import redact_string
|
||||
|
|
@ -718,6 +719,28 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time):
|
||||
self._handle_failure(kwargs, response_obj, start_time, end_time)
|
||||
|
||||
def _start_service_span(self, payload: ServiceLoggerPayload, parent_otel_span: Span, start_time_ns: int) -> Span:
|
||||
"""Open a service span, named and classified by what the service is.
|
||||
|
||||
A datastore call is an outbound CLIENT span carrying ``db.*`` semconv.
|
||||
Without those a Postgres span says only ``service=postgres``, so the
|
||||
backend falls back to the transport peer, which for Prisma is the local
|
||||
query engine on loopback. Everything else stays an INTERNAL span.
|
||||
"""
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.trace import SpanKind
|
||||
|
||||
attributes: Final = db_span_attributes(payload.service.value, payload.call_type)
|
||||
span: Final = self.tracer.start_span(
|
||||
name=payload.service,
|
||||
context=trace.set_span_in_context(parent_otel_span),
|
||||
start_time=start_time_ns,
|
||||
kind=SpanKind.CLIENT if attributes else SpanKind.INTERNAL,
|
||||
)
|
||||
for key, value in attributes.items():
|
||||
self.safe_set_attribute(span=span, key=key, value=value)
|
||||
return span
|
||||
|
||||
async def async_service_success_hook(
|
||||
self,
|
||||
payload: ServiceLoggerPayload,
|
||||
|
|
@ -726,7 +749,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
end_time: datetime | float | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
):
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
_start_time_ns = 0
|
||||
|
|
@ -743,12 +765,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
_end_time_ns = self._to_ns(end_time)
|
||||
|
||||
if parent_otel_span is not None:
|
||||
_span_name: Final = payload.service
|
||||
service_logging_span: Final = self.tracer.start_span(
|
||||
name=_span_name,
|
||||
context=trace.set_span_in_context(parent_otel_span),
|
||||
start_time=_start_time_ns,
|
||||
)
|
||||
service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns)
|
||||
self.safe_set_attribute(
|
||||
span=service_logging_span,
|
||||
key="call_type",
|
||||
|
|
@ -786,7 +803,6 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
end_time: float | datetime | None = None,
|
||||
event_metadata: dict | None = None,
|
||||
):
|
||||
from opentelemetry import trace
|
||||
from opentelemetry.trace import Status, StatusCode
|
||||
|
||||
_start_time_ns = 0
|
||||
|
|
@ -803,12 +819,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger):
|
|||
_end_time_ns = self._to_ns(end_time)
|
||||
|
||||
if parent_otel_span is not None:
|
||||
_span_name: Final = payload.service
|
||||
service_logging_span: Final = self.tracer.start_span(
|
||||
name=_span_name,
|
||||
context=trace.set_span_in_context(parent_otel_span),
|
||||
start_time=_start_time_ns,
|
||||
)
|
||||
service_logging_span: Final = self._start_service_span(payload, parent_otel_span, _start_time_ns)
|
||||
self.safe_set_attribute(
|
||||
span=service_logging_span,
|
||||
key="call_type",
|
||||
|
|
|
|||
|
|
@ -18,6 +18,7 @@ from litellm.integrations.otel.mappers.utils import (
|
|||
serialize_messages,
|
||||
tool_definition_attrs,
|
||||
)
|
||||
from litellm.integrations.otel.model.db_endpoint import db_span_attributes
|
||||
from litellm.integrations.otel.model.payloads import (
|
||||
GuardrailSpanData,
|
||||
LLMCallSpanData,
|
||||
|
|
@ -27,7 +28,6 @@ from litellm.integrations.otel.model.payloads import (
|
|||
ToolDefinition,
|
||||
)
|
||||
from litellm.integrations.otel.model.semconv import (
|
||||
DB,
|
||||
MCP,
|
||||
Error,
|
||||
GenAI,
|
||||
|
|
@ -36,7 +36,6 @@ from litellm.integrations.otel.model.semconv import (
|
|||
RpcSystem,
|
||||
Server,
|
||||
)
|
||||
from litellm.integrations.otel.model.spans import db_system
|
||||
|
||||
|
||||
class GenAIMapper:
|
||||
|
|
@ -182,12 +181,8 @@ class GenAIMapper:
|
|||
def _service(cls, data: ServiceSpanData) -> AttributeMap:
|
||||
attrs: Final = collect(cls._SERVICE_ATTRS, data)
|
||||
# An outbound datastore call (DB_CALL / CLIENT span) also carries db.*
|
||||
# semconv. Internal services (router, budget jobs, …) have no db.system,
|
||||
# so they get only the litellm.service.* keys above.
|
||||
system: Final = db_system(data.service_name)
|
||||
if system is not None:
|
||||
attrs[DB.SYSTEM_NAME] = system
|
||||
if data.call_type:
|
||||
attrs[DB.OPERATION_NAME] = data.call_type
|
||||
# semconv naming the server it reached. Internal services (router, budget
|
||||
# jobs, …) have no db.system, so they get only the litellm.service.* keys.
|
||||
attrs.update(db_span_attributes(data.service_name, data.call_type))
|
||||
attrs.update({f"{LiteLLM.METADATA_PREFIX}{key}": value for key, value in data.event_metadata.items()})
|
||||
return attrs
|
||||
|
|
|
|||
164
litellm/integrations/otel/model/db_endpoint.py
Normal file
164
litellm/integrations/otel/model/db_endpoint.py
Normal file
|
|
@ -0,0 +1,164 @@
|
|||
"""OTel ``db.*`` / ``server.*`` attributes naming the database litellm talks to.
|
||||
|
||||
Prisma reaches PostgreSQL through a query engine listening on loopback, so
|
||||
transport-level instrumentation attributes the work to ``localhost`` and an
|
||||
operator cannot tell it is a PostgreSQL call or correlate it with the database's
|
||||
own metrics. These attributes name the real server on litellm's DB spans.
|
||||
|
||||
Only the host, port, database and schema of the DSN are read, so no credential
|
||||
can reach an exporter.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping, Sequence
|
||||
from dataclasses import dataclass
|
||||
from types import MappingProxyType
|
||||
from typing import Final
|
||||
from urllib.parse import ParseResult, parse_qs, unquote, urlparse
|
||||
|
||||
from litellm.integrations.otel.model.semconv import DB, Server
|
||||
from litellm.integrations.otel.model.spans import POSTGRESQL, db_system
|
||||
|
||||
_DATABASE_URL_ENV: Final = "DATABASE_URL"
|
||||
_READ_REPLICA_ENV: Final = "DATABASE_URL_READ_REPLICA"
|
||||
_DEFAULT_POSTGRES_PORT: Final = 5432
|
||||
_DEFAULT_POSTGRES_SCHEMA: Final = "public"
|
||||
_POSTGRES_SCHEMES: Final = frozenset({"postgres", "postgresql"})
|
||||
_EMPTY_ATTRIBUTES: Final[Mapping[str, str | int]] = MappingProxyType({})
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class DatabaseEndpoint:
|
||||
"""The non-sensitive identity of a PostgreSQL server, parsed from a DSN."""
|
||||
|
||||
address: str | None
|
||||
port: int | None
|
||||
namespace: str | None
|
||||
|
||||
|
||||
def parse_database_endpoint(url: str | None) -> DatabaseEndpoint | None:
|
||||
"""Parse a PostgreSQL DSN into its exportable endpoint identity.
|
||||
|
||||
Returns ``None`` for an absent, malformed or non-PostgreSQL URL rather than
|
||||
raising: an unparseable DSN must degrade to a span without endpoint
|
||||
attributes, never break the request that emitted it.
|
||||
"""
|
||||
if not url:
|
||||
return None
|
||||
try:
|
||||
parsed: Final = urlparse(url)
|
||||
if parsed.scheme not in _POSTGRES_SCHEMES:
|
||||
return None
|
||||
query: Final = parse_qs(parsed.query)
|
||||
raw_database: Final = (parsed.path or "").lstrip("/")
|
||||
if _is_misparsed_authority(parsed, url, raw_database):
|
||||
return None
|
||||
# ``host=`` beats the netloc: it is how libpq names a Unix socket
|
||||
# directory and how the Cloud SQL connector sits behind a localhost
|
||||
# netloc, where the netloc is the very answer this module replaces.
|
||||
address: Final = _first(query.get("host")) or parsed.hostname
|
||||
# ``port=`` accompanies ``host=`` in a libpq URI, so honour it the same way.
|
||||
port: Final = _port(_first(query.get("port")), parsed.port) if address else None
|
||||
namespace: Final = _namespace(unquote(raw_database), _first(query.get("schema")))
|
||||
except ValueError:
|
||||
return None
|
||||
if address is None and namespace is None:
|
||||
return None
|
||||
return DatabaseEndpoint(address=address, port=port, namespace=namespace)
|
||||
|
||||
|
||||
def _is_misparsed_authority(parsed: ParseResult, url: str, raw_database: str) -> bool:
|
||||
"""Whether the URL authority may have been truncated by an unencoded character.
|
||||
|
||||
``/``, ``#`` or ``?`` in a password ends the netloc early, so urlparse hands
|
||||
back the username as the host, the leading digits of the password as the
|
||||
port, and the rest of the credential as the path, query or fragment. The
|
||||
stranded userinfo ``@`` is the only surviving evidence.
|
||||
|
||||
A database name cannot hold an unencoded slash either, so a second path
|
||||
segment is the same evidence.
|
||||
|
||||
A DSN that carries the at-sign in a query parameter instead, such as
|
||||
``?application_name=svc@prod``, is indistinguishable from a mis-split by any
|
||||
property of the parse: both leave no userinfo, a host, a port and a path.
|
||||
Since guessing wrong publishes a credential fragment to a tracing backend,
|
||||
that ambiguity resolves to refusing the endpoint. Such a DSN loses
|
||||
``server.address`` and ``db.namespace`` and keeps the rest of the span,
|
||||
which is the cheaper error of the two. Percent-encode the at-sign to keep
|
||||
them.
|
||||
"""
|
||||
if "/" in raw_database:
|
||||
return True
|
||||
return "@" in url and "@" not in parsed.netloc
|
||||
|
||||
|
||||
def _first(values: Sequence[str] | None) -> str:
|
||||
return values[0] if values else ""
|
||||
|
||||
|
||||
def _port(from_query: str, from_netloc: int | None) -> int:
|
||||
return int(from_query) if from_query.isdigit() else (from_netloc or _DEFAULT_POSTGRES_PORT)
|
||||
|
||||
|
||||
def _namespace(database: str, schema: str) -> str | None:
|
||||
"""``{database}|{schema}`` per the PostgreSQL semconv, dropping absent halves.
|
||||
|
||||
Only Prisma's literal default schema stays implicit. The match is
|
||||
case-sensitive because Prisma quotes the name, so ``?schema=PUBLIC`` builds
|
||||
a second schema alongside ``public`` and the two must not collapse to one
|
||||
namespace.
|
||||
"""
|
||||
qualifier: Final = "" if schema == _DEFAULT_POSTGRES_SCHEMA else schema
|
||||
return "|".join(part for part in (database, qualifier) if part) or None
|
||||
|
||||
|
||||
def postgres_endpoint() -> DatabaseEndpoint | None:
|
||||
"""The PostgreSQL endpoint the process is currently connected to.
|
||||
|
||||
Read from ``os.environ`` on every span, deliberately, on both counts.
|
||||
|
||||
The environment is what Prisma itself connects with, so the span cannot
|
||||
disagree with the connection; ``get_secret_str`` would consult a configured
|
||||
secret manager first and could name a different server than the one serving
|
||||
the query. And the value is not static: the RDS IAM refresh rebuilds the URL
|
||||
from ``DATABASE_HOST``/``PORT``/``NAME``/``SCHEMA`` every rotation, the
|
||||
reconnect path re-reads ``DATABASE_URL``, and the DB-backed
|
||||
``environment_variables`` config overlay can rewrite any of them after
|
||||
startup, so a value cached for the process lifetime goes stale against a
|
||||
connection that has genuinely moved. Nothing is memoized either: a cache
|
||||
keyed on the URL would hold a rotated credential past its rotation, and the
|
||||
parse is a single ``urlparse`` on a short string.
|
||||
|
||||
A configured read replica yields ``None``: ``RoutingPrismaWrapper`` picks
|
||||
reader or writer per Prisma call, underneath the span, so naming the writer
|
||||
would attribute replica reads to the primary.
|
||||
"""
|
||||
if os.environ.get(_READ_REPLICA_ENV):
|
||||
return None
|
||||
return parse_database_endpoint(os.environ.get(_DATABASE_URL_ENV, ""))
|
||||
|
||||
|
||||
def db_span_attributes(service_name: str, call_type: str | None = None) -> Mapping[str, str | int]:
|
||||
"""The ``db.*``/``server.*`` attributes for a datastore service call.
|
||||
|
||||
Empty for services that are not outbound datastore calls. Endpoint
|
||||
attributes are PostgreSQL-only: ``DATABASE_URL`` says nothing about where
|
||||
the redis-backed services point. ``db.system`` rides alongside the current
|
||||
``db.system.name`` because Datadog's OTLP intake still types a database span
|
||||
from the older key.
|
||||
"""
|
||||
system: Final = db_system(service_name)
|
||||
if system is None:
|
||||
return _EMPTY_ATTRIBUTES
|
||||
endpoint: Final = postgres_endpoint() if system == POSTGRESQL else None
|
||||
pairs: Final[tuple[tuple[str, str | int | None], ...]] = (
|
||||
(DB.SYSTEM_NAME, system),
|
||||
(DB.SYSTEM_LEGACY, system),
|
||||
(DB.OPERATION_NAME, call_type),
|
||||
(Server.ADDRESS, endpoint.address if endpoint is not None else None),
|
||||
(Server.PORT, endpoint.port if endpoint is not None else None),
|
||||
(DB.NAMESPACE, endpoint.namespace if endpoint is not None else None),
|
||||
)
|
||||
return MappingProxyType({key: value for key, value in pairs if value})
|
||||
|
|
@ -238,7 +238,11 @@ class DB:
|
|||
"""
|
||||
|
||||
SYSTEM_NAME: Final = "db.system.name"
|
||||
# Superseded by SYSTEM_NAME, dual-emitted because Datadog's OTLP intake
|
||||
# still infers a span's database type from this key.
|
||||
SYSTEM_LEGACY: Final = "db.system"
|
||||
OPERATION_NAME: Final = "db.operation.name"
|
||||
NAMESPACE: Final = "db.namespace"
|
||||
|
||||
|
||||
class HTTP:
|
||||
|
|
|
|||
|
|
@ -115,10 +115,12 @@ SPAN_REGISTRY: Final[dict[SpanRole, SpanSpec]] = {
|
|||
# redis-backed spend queues. Any service not mapped here is litellm-internal work
|
||||
# and stays an INTERNAL ``SERVICE`` span. This table is the single source of
|
||||
# datastore knowledge — both the role classifier and the mapper read it.
|
||||
POSTGRESQL: Final = "postgresql"
|
||||
|
||||
_DB_SYSTEM_BY_SERVICE: Final[dict[str, str]] = {
|
||||
"redis": "redis",
|
||||
"postgres": "postgresql",
|
||||
"batch_write_to_db": "postgresql",
|
||||
"postgres": POSTGRESQL,
|
||||
"batch_write_to_db": POSTGRESQL,
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
312
tests/test_litellm/integrations/otel/test_db_endpoint.py
Normal file
312
tests/test_litellm/integrations/otel/test_db_endpoint.py
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
"""Tests for litellm/integrations/otel/model/db_endpoint.py
|
||||
|
||||
Prisma talks to PostgreSQL through a loopback query engine, so a DB span with no
|
||||
``server.address`` gets attributed to ``localhost`` by the backend. These cover
|
||||
the endpoint derivation that names the real server, for the local engine and for
|
||||
remote and read-replica deployments, and pin the rule that no credential is ever
|
||||
exported.
|
||||
"""
|
||||
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
|
||||
from litellm.integrations.otel.model.db_endpoint import (
|
||||
DatabaseEndpoint,
|
||||
db_span_attributes,
|
||||
parse_database_endpoint,
|
||||
postgres_endpoint,
|
||||
)
|
||||
|
||||
LOCAL_DSN = "postgresql://llmproxy:dbpassword9090@localhost:5432/litellm"
|
||||
REMOTE_DSN = "postgresql://llmproxy:s3cr3t@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting&sslmode=require"
|
||||
REPLICA_DSN = "postgresql://reader:r3ad0nly@litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com/litellm_replica"
|
||||
|
||||
|
||||
def _resolve(service, call_type=None, database_url=None, read_replica_url=None):
|
||||
"""Resolve attributes with the two DB env vars set, as the proxy sets them."""
|
||||
env = {k: v for k, v in (("DATABASE_URL", database_url), ("DATABASE_URL_READ_REPLICA", read_replica_url)) if v}
|
||||
with patch.dict(os.environ, env, clear=False):
|
||||
for absent in {"DATABASE_URL", "DATABASE_URL_READ_REPLICA"} - set(env):
|
||||
os.environ.pop(absent, None)
|
||||
return dict(db_span_attributes(service, call_type))
|
||||
|
||||
|
||||
def test_local_prisma_engine_endpoint_is_the_postgres_server_not_the_engine():
|
||||
assert parse_database_endpoint(LOCAL_DSN) == DatabaseEndpoint(
|
||||
address="localhost", port=5432, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_remote_endpoint_keeps_host_port_and_schema_qualified_namespace():
|
||||
assert parse_database_endpoint(REMOTE_DSN) == DatabaseEndpoint(
|
||||
address="litellm-prod.abc123.us-east-1.rds.amazonaws.com",
|
||||
port=6432,
|
||||
namespace="litellm|reporting",
|
||||
)
|
||||
|
||||
|
||||
def test_read_replica_dsn_parses_to_the_replica_host_and_database():
|
||||
assert parse_database_endpoint(REPLICA_DSN) == DatabaseEndpoint(
|
||||
address="litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com",
|
||||
port=5432,
|
||||
namespace="litellm_replica",
|
||||
)
|
||||
|
||||
|
||||
def test_default_schema_is_not_spelled_out_in_the_namespace():
|
||||
"""``?schema=public`` and no schema at all are the same deployment, so they
|
||||
must not split a group-by on db.namespace."""
|
||||
assert parse_database_endpoint("postgresql://u:p@db.internal/litellm?schema=public") == parse_database_endpoint(
|
||||
"postgresql://u:p@db.internal/litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_unix_socket_host_parameter_wins_over_the_netloc():
|
||||
"""libpq and the Cloud SQL connector both put the real target in ``host=``
|
||||
behind a localhost netloc, which is the attribution this module removes."""
|
||||
assert parse_database_endpoint(
|
||||
"postgresql://u:p@localhost:5432/litellm?host=/cloudsql/proj:us-east1:inst"
|
||||
) == DatabaseEndpoint(address="/cloudsql/proj:us-east1:inst", port=5432, namespace="litellm")
|
||||
|
||||
|
||||
def test_socket_only_dsn_without_a_netloc_host_still_resolves():
|
||||
assert parse_database_endpoint("postgresql:///litellm?host=/var/run/postgresql") == DatabaseEndpoint(
|
||||
address="/var/run/postgresql", port=5432, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_percent_encoded_database_name_is_decoded():
|
||||
endpoint = parse_database_endpoint("postgresql://u:p@db.internal/litellm%20prod")
|
||||
assert endpoint is not None and endpoint.namespace == "litellm prod"
|
||||
|
||||
|
||||
MISPARSED_AUTHORITY_DSNS = (
|
||||
("postgresql://litellm:/kJ8xQz+9wT@db.internal:5432/litellm", "kJ8xQz+9wT"),
|
||||
("postgresql://litellm:12345/aBcD@db.internal:5432/litellm", "aBcD"),
|
||||
# '#' sends the tail to the fragment and '?' to the query, so the path is
|
||||
# empty and only the stranded userinfo '@' reveals the mis-split.
|
||||
("postgresql://litellm:12345#aBcD@db.internal/litellm", "aBcD"),
|
||||
("postgresql://litellm:12345?aBcD@db.internal/litellm", "aBcD"),
|
||||
# A '?'-stranded tail that happens to parse as parameters, including one
|
||||
# that hijacks the host= parameter into server.address.
|
||||
("postgresql://litellm:12345?a=aBcD@db.internal/litellm", "aBcD"),
|
||||
("postgresql://litellm:12345?host=aBcD@db.internal/litellm", "aBcD"),
|
||||
# Both '/' and '?key=value' together: the slash leaves a clean path holding
|
||||
# the password remainder and the query still parses, so only the stranded
|
||||
# at-sign gives it away.
|
||||
("postgresql://litellm:12345/aBcD?x=1@db.internal/litellm", "aBcD"),
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("dsn", "secret"), MISPARSED_AUTHORITY_DSNS)
|
||||
def test_unencoded_slash_in_password_never_yields_an_endpoint(dsn, secret):
|
||||
"""An unencoded '/' truncates the authority, so urlparse reports the username
|
||||
as the host and the password tail as the database. Postgres drivers reject
|
||||
such a DSN outright, so the only safe reading is no endpoint at all."""
|
||||
assert parse_database_endpoint(dsn) is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("dsn", "secret"), MISPARSED_AUTHORITY_DSNS)
|
||||
def test_unencoded_slash_in_password_never_reaches_a_span(dsn, secret):
|
||||
attrs = _resolve("postgres", "get_data", database_url=dsn)
|
||||
exported = " ".join(str(value) for value in attrs.values())
|
||||
assert secret not in exported
|
||||
assert "db.namespace" not in attrs
|
||||
assert "server.address" not in attrs
|
||||
|
||||
|
||||
def test_extra_path_segment_yields_no_endpoint():
|
||||
"""A database name cannot hold an unencoded '/', so a second path segment
|
||||
means the authority was mis-split even when no '@' survived into the path."""
|
||||
assert parse_database_endpoint("postgresql://db.internal:5432/litellm/extra") is None
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dsn", [d for d, _ in MISPARSED_AUTHORITY_DSNS])
|
||||
def test_a_mis_split_authority_never_exports_the_database_username(dsn):
|
||||
"""The username lands in ``parsed.hostname`` when the authority truncates, so
|
||||
a span would name the DB user as the server."""
|
||||
attrs = _resolve("postgres", "get_data", database_url=dsn)
|
||||
assert "server.address" not in attrs
|
||||
assert "litellm" not in " ".join(str(v) for v in attrs.values())
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dsn",
|
||||
[
|
||||
"postgresql://db.internal:5432/litellm?application_name=svc@prod",
|
||||
"postgresql://db.internal:5432/litellm?user=admin@company.com",
|
||||
],
|
||||
)
|
||||
def test_an_unencoded_at_sign_in_a_query_forfeits_the_endpoint(dsn):
|
||||
"""This shape is byte-for-byte indistinguishable from a mis-split password,
|
||||
so it resolves to no endpoint rather than risking a credential fragment.
|
||||
Percent-encoding the at-sign restores the attributes."""
|
||||
assert parse_database_endpoint(dsn) is None
|
||||
assert parse_database_endpoint(dsn.replace("@", "%40")) is not None
|
||||
|
||||
|
||||
def test_host_and_port_query_parameters_are_honoured_together():
|
||||
assert parse_database_endpoint("postgresql://ignored/litellm?host=real.internal&port=6543") == DatabaseEndpoint(
|
||||
address="real.internal", port=6543, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_percent_encoded_password_still_resolves_the_endpoint():
|
||||
"""The encoded spelling is the one a driver accepts, so it must keep working."""
|
||||
assert parse_database_endpoint("postgresql://litellm:pa%2Fssw0rd@db.internal:5432/litellm") == DatabaseEndpoint(
|
||||
address="db.internal", port=5432, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_hostless_socket_dsn_still_names_the_database():
|
||||
"""``postgresql:///litellm`` is a valid local-socket DSN that Prisma accepts,
|
||||
so the database is knowable even though no server address is."""
|
||||
assert parse_database_endpoint("postgresql:///litellm") == DatabaseEndpoint(
|
||||
address=None, port=None, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
def test_hostless_socket_dsn_emits_namespace_without_a_server():
|
||||
attrs = _resolve("postgres", "get_data", database_url="postgresql:///litellm")
|
||||
assert attrs["db.namespace"] == "litellm"
|
||||
assert "server.address" not in attrs
|
||||
assert "server.port" not in attrs
|
||||
|
||||
|
||||
def test_dsn_with_neither_host_nor_database_yields_no_endpoint():
|
||||
assert parse_database_endpoint("postgresql://") is None
|
||||
|
||||
|
||||
def test_prisma_default_schema_is_left_implicit():
|
||||
endpoint = parse_database_endpoint("postgresql://u:p@db.internal/litellm?schema=public")
|
||||
assert endpoint is not None and endpoint.namespace == "litellm"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ["PUBLIC", "Public", "reporting"])
|
||||
def test_a_non_default_schema_stays_in_the_namespace(spelling):
|
||||
"""Prisma quotes the schema name, so ``?schema=PUBLIC`` provisions a second
|
||||
schema alongside ``public`` with its own tables. Case-folding them into one
|
||||
namespace would report two different schemas as the same database."""
|
||||
endpoint = parse_database_endpoint(f"postgresql://u:p@db.internal/litellm?schema={spelling}")
|
||||
assert endpoint is not None and endpoint.namespace == f"litellm|{spelling}"
|
||||
|
||||
|
||||
def test_postgres_scheme_alias_is_accepted():
|
||||
assert parse_database_endpoint("postgres://u:p@db.internal/litellm") == DatabaseEndpoint(
|
||||
address="db.internal", port=5432, namespace="litellm"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"dsn",
|
||||
[
|
||||
None,
|
||||
"",
|
||||
"mysql://u:p@db.internal:3306/litellm",
|
||||
"postgresql://u:p@db.internal:not-a-port/litellm",
|
||||
"not a url at all",
|
||||
],
|
||||
)
|
||||
def test_unusable_dsn_degrades_to_no_endpoint(dsn):
|
||||
assert parse_database_endpoint(dsn) is None
|
||||
|
||||
|
||||
def test_database_without_name_or_schema_has_no_namespace():
|
||||
assert parse_database_endpoint("postgresql://u:p@db.internal:5432/") == DatabaseEndpoint(
|
||||
address="db.internal", port=5432, namespace=None
|
||||
)
|
||||
|
||||
|
||||
def test_postgres_service_span_carries_system_operation_and_endpoint():
|
||||
assert _resolve("postgres", "get_data", database_url=REMOTE_DSN) == {
|
||||
"db.system.name": "postgresql",
|
||||
"db.system": "postgresql",
|
||||
"db.operation.name": "get_data",
|
||||
"server.address": "litellm-prod.abc123.us-east-1.rds.amazonaws.com",
|
||||
"server.port": 6432,
|
||||
"db.namespace": "litellm|reporting",
|
||||
}
|
||||
|
||||
|
||||
def test_legacy_db_system_is_dual_emitted_for_datadog():
|
||||
"""Datadog's OTLP intake infers the database span type from ``db.system``,
|
||||
not from the semconv-current ``db.system.name``."""
|
||||
assert _resolve("postgres", "get_data", database_url=LOCAL_DSN)["db.system"] == "postgresql"
|
||||
assert _resolve("redis", "set")["db.system"] == "redis"
|
||||
|
||||
|
||||
def test_batch_write_service_is_also_attributed_to_postgres():
|
||||
attrs = _resolve("batch_write_to_db", "_PROXY_track_cost_callback", database_url=REMOTE_DSN)
|
||||
assert attrs["db.system.name"] == "postgresql"
|
||||
assert attrs["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com"
|
||||
|
||||
|
||||
def test_redis_service_never_borrows_the_postgres_endpoint():
|
||||
assert _resolve("redis", "set", database_url=REMOTE_DSN) == {
|
||||
"db.system.name": "redis",
|
||||
"db.system": "redis",
|
||||
"db.operation.name": "set",
|
||||
}
|
||||
|
||||
|
||||
def test_non_datastore_service_gets_no_db_attributes():
|
||||
assert _resolve("reset_budget_job", "reset_budget", database_url=REMOTE_DSN) == {}
|
||||
|
||||
|
||||
def test_configured_read_replica_suppresses_the_endpoint_rather_than_naming_the_primary():
|
||||
"""Reads are routed to the replica per Prisma call, underneath the span, so
|
||||
naming the writer would pin replica latency onto the primary."""
|
||||
attrs = _resolve("postgres", "get_data", database_url=REMOTE_DSN, read_replica_url=REPLICA_DSN)
|
||||
assert attrs == {
|
||||
"db.system.name": "postgresql",
|
||||
"db.system": "postgresql",
|
||||
"db.operation.name": "get_data",
|
||||
}
|
||||
|
||||
|
||||
def test_endpoint_attributes_are_omitted_when_database_url_is_unset():
|
||||
assert _resolve("postgres", "get_data") == {
|
||||
"db.system.name": "postgresql",
|
||||
"db.system": "postgresql",
|
||||
"db.operation.name": "get_data",
|
||||
}
|
||||
|
||||
|
||||
def test_blank_call_type_does_not_emit_an_empty_operation_attribute():
|
||||
assert "db.operation.name" not in _resolve("postgres", "")
|
||||
assert "db.operation.name" not in _resolve("postgres", None)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("dsn", "secrets"),
|
||||
[
|
||||
(LOCAL_DSN, ("dbpassword9090", "llmproxy")),
|
||||
(REMOTE_DSN, ("s3cr3t", "llmproxy", "sslmode")),
|
||||
(REPLICA_DSN, ("r3ad0nly", "reader")),
|
||||
],
|
||||
)
|
||||
def test_no_credential_reaches_any_exported_attribute(dsn, secrets):
|
||||
attrs = _resolve("postgres", "get_data", database_url=dsn)
|
||||
assert attrs["server.address"]
|
||||
exported = " ".join(str(value) for value in attrs.values())
|
||||
for secret in secrets:
|
||||
assert secret not in exported
|
||||
|
||||
|
||||
def test_a_runtime_endpoint_change_is_reflected_on_the_next_span():
|
||||
"""The RDS IAM refresh, the reconnect path and the DB-backed
|
||||
environment_variables overlay can all rewrite DATABASE_URL after startup, so
|
||||
a value cached for the process lifetime would report a server the process no
|
||||
longer talks to."""
|
||||
first = _resolve("postgres", "get_data", database_url=LOCAL_DSN)
|
||||
assert first["server.address"] == "localhost"
|
||||
moved = _resolve("postgres", "get_data", database_url=REMOTE_DSN)
|
||||
assert moved["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com"
|
||||
|
||||
|
||||
def test_a_replica_configured_after_the_first_span_suppresses_the_endpoint():
|
||||
assert _resolve("postgres", "get_data", database_url=REMOTE_DSN)["server.address"]
|
||||
later = _resolve("postgres", "get_data", database_url=REMOTE_DSN, read_replica_url=REPLICA_DSN)
|
||||
assert "server.address" not in later
|
||||
|
|
@ -8,6 +8,8 @@ hooks, proxy SERVER span lifecycle (start + setters), parent-context resolution
|
|||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import os
|
||||
from unittest.mock import patch
|
||||
from datetime import datetime, timedelta, timezone
|
||||
|
||||
import pytest
|
||||
|
|
@ -1521,6 +1523,36 @@ def test_async_service_success_hook_emits_service_span():
|
|||
assert span.status.status_code is StatusCode.UNSET
|
||||
|
||||
|
||||
def test_postgres_db_span_names_the_database_server_not_the_prisma_engine():
|
||||
"""Prisma reaches Postgres over loopback, so without server.address the
|
||||
backend attributes the wait to localhost."""
|
||||
dsn = "postgresql://llmproxy:dbpassword9090@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting"
|
||||
logger, exporter = _logger()
|
||||
parent = _service_parent(logger)
|
||||
try:
|
||||
with patch.dict(os.environ, {"DATABASE_URL": dsn}, clear=False):
|
||||
os.environ.pop("DATABASE_URL_READ_REPLICA", None)
|
||||
asyncio.run(
|
||||
logger.async_service_success_hook(
|
||||
payload=_ServicePayload("postgres", "get_data"),
|
||||
parent_otel_span=parent,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
parent.end()
|
||||
span = {s.name: s for s in exporter.get_finished_spans()}["postgres get_data"]
|
||||
assert span.kind is SpanKind.CLIENT
|
||||
assert span.attributes["db.system.name"] == "postgresql"
|
||||
assert span.attributes["db.operation.name"] == "get_data"
|
||||
assert span.attributes["server.address"] == "litellm-prod.abc123.us-east-1.rds.amazonaws.com"
|
||||
assert span.attributes["server.port"] == 6432
|
||||
assert span.attributes["db.namespace"] == "litellm|reporting"
|
||||
assert span.attributes["db.system"] == "postgresql"
|
||||
exported = " ".join(str(value) for value in span.attributes.values())
|
||||
assert "dbpassword9090" not in exported
|
||||
assert "llmproxy" not in exported
|
||||
|
||||
|
||||
def test_async_service_failure_hook_marks_error_status():
|
||||
logger, exporter = _logger()
|
||||
parent = _service_parent(logger)
|
||||
|
|
|
|||
|
|
@ -34,6 +34,7 @@ from litellm.integrations.opentelemetry import (
|
|||
_normalize_team_metadata_keys,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.types.services import ServiceLoggerPayload, ServiceTypes
|
||||
|
||||
|
||||
class TestOpenTelemetryGuardrails(unittest.TestCase):
|
||||
|
|
@ -6212,3 +6213,110 @@ class TestDynamicTracerProviderCache(unittest.TestCase):
|
|||
|
||||
self.assertTrue(entry.owns_exporter)
|
||||
self.assertIsNotNone(entry.provider._atexit_handler)
|
||||
class TestOpenTelemetryDatabaseSemconvAttributes(unittest.TestCase):
|
||||
"""A Postgres service span must name the PostgreSQL server it reached.
|
||||
|
||||
Without ``db.system`` and ``server.address``, the only host in the trace is
|
||||
the loopback address of Prisma's local query engine, so the backend
|
||||
attributes the wait to ``localhost`` and it cannot be correlated with the
|
||||
database's own metrics.
|
||||
"""
|
||||
|
||||
DSN = "postgresql://llmproxy:dbpassword9090@litellm-prod.abc123.us-east-1.rds.amazonaws.com:6432/litellm?schema=reporting"
|
||||
REPLICA_DSN = "postgresql://reader:r3ad0nly@litellm-prod-ro.abc123.us-east-1.rds.amazonaws.com/litellm"
|
||||
|
||||
def _service_span(self, service, call_type, dsn, error=None, replica_dsn=None):
|
||||
exporter = InMemorySpanExporter()
|
||||
provider = TracerProvider()
|
||||
provider.add_span_processor(SimpleSpanProcessor(exporter))
|
||||
otel = OpenTelemetry()
|
||||
otel.tracer = provider.get_tracer(__name__)
|
||||
parent = otel.tracer.start_span("Received Proxy Server Request")
|
||||
payload = ServiceLoggerPayload(
|
||||
is_error=error is not None,
|
||||
error=error,
|
||||
service=service,
|
||||
duration=0.25,
|
||||
call_type=call_type,
|
||||
event_metadata=None,
|
||||
)
|
||||
hook = otel.async_service_failure_hook if error else otel.async_service_success_hook
|
||||
kwargs = {"error": error} if error else {}
|
||||
env = {k: v for k, v in (("DATABASE_URL", dsn), ("DATABASE_URL_READ_REPLICA", replica_dsn)) if v}
|
||||
with patch.dict(os.environ, env, clear=False):
|
||||
for absent in {"DATABASE_URL", "DATABASE_URL_READ_REPLICA"} - set(env):
|
||||
os.environ.pop(absent, None)
|
||||
asyncio.run(
|
||||
hook(
|
||||
payload=payload,
|
||||
parent_otel_span=parent,
|
||||
start_time=datetime.now(),
|
||||
end_time=datetime.now(),
|
||||
**kwargs,
|
||||
)
|
||||
)
|
||||
parent.end()
|
||||
return next(s for s in exporter.get_finished_spans() if s.name == service.value)
|
||||
|
||||
def test_postgres_span_names_the_database_server(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN)
|
||||
self.assertEqual(span.attributes["db.system.name"], "postgresql")
|
||||
self.assertEqual(span.attributes["db.operation.name"], "get_data")
|
||||
self.assertEqual(
|
||||
span.attributes["server.address"],
|
||||
"litellm-prod.abc123.us-east-1.rds.amazonaws.com",
|
||||
)
|
||||
self.assertEqual(span.attributes["server.port"], 6432)
|
||||
self.assertEqual(span.attributes["db.namespace"], "litellm|reporting")
|
||||
|
||||
def test_datastore_span_is_a_client_span_carrying_the_legacy_db_system(self):
|
||||
"""Datadog types a span as a database call from CLIENT kind plus
|
||||
``db.system``; an INTERNAL span is classified as custom work."""
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN)
|
||||
self.assertEqual(span.kind, trace.SpanKind.CLIENT)
|
||||
self.assertEqual(span.attributes["db.system"], "postgresql")
|
||||
|
||||
def test_internal_service_span_stays_internal(self):
|
||||
span = self._service_span(ServiceTypes.RESET_BUDGET_JOB, "reset_budget", self.DSN)
|
||||
self.assertEqual(span.kind, trace.SpanKind.INTERNAL)
|
||||
self.assertNotIn("db.system.name", span.attributes)
|
||||
self.assertNotIn("server.address", span.attributes)
|
||||
|
||||
def test_existing_service_and_call_type_attributes_are_unchanged(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN)
|
||||
self.assertEqual(span.attributes["service"], "postgres")
|
||||
self.assertEqual(span.attributes["call_type"], "get_data")
|
||||
|
||||
def test_failed_postgres_span_also_names_the_database_server(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN, error="connection refused")
|
||||
self.assertEqual(span.attributes["db.system.name"], "postgresql")
|
||||
self.assertEqual(span.kind, trace.SpanKind.CLIENT)
|
||||
self.assertEqual(
|
||||
span.attributes["server.address"],
|
||||
"litellm-prod.abc123.us-east-1.rds.amazonaws.com",
|
||||
)
|
||||
self.assertEqual(span.attributes["error"], "connection refused")
|
||||
|
||||
def test_no_credential_from_the_dsn_lands_on_the_span(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN)
|
||||
exported = " ".join(str(value) for value in span.attributes.values())
|
||||
self.assertIn("litellm-prod.abc123.us-east-1.rds.amazonaws.com", exported)
|
||||
self.assertNotIn("dbpassword9090", exported)
|
||||
self.assertNotIn("llmproxy", exported)
|
||||
|
||||
def test_redis_span_does_not_borrow_the_postgres_endpoint(self):
|
||||
span = self._service_span(ServiceTypes.REDIS, "async_set_cache", self.DSN)
|
||||
self.assertEqual(span.attributes["db.system.name"], "redis")
|
||||
self.assertEqual(span.kind, trace.SpanKind.CLIENT)
|
||||
self.assertNotIn("server.address", span.attributes)
|
||||
|
||||
def test_configured_read_replica_suppresses_the_endpoint(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", self.DSN, replica_dsn=self.REPLICA_DSN)
|
||||
self.assertEqual(span.attributes["db.system.name"], "postgresql")
|
||||
self.assertNotIn("server.address", span.attributes)
|
||||
self.assertNotIn("db.namespace", span.attributes)
|
||||
|
||||
def test_unset_database_url_leaves_the_span_without_endpoint_attributes(self):
|
||||
span = self._service_span(ServiceTypes.DB, "get_data", None)
|
||||
self.assertEqual(span.attributes["db.system.name"], "postgresql")
|
||||
self.assertNotIn("server.address", span.attributes)
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue