diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 9363db12385..f2bd18cd046 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -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", diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index 032441535e0..79487e69ac4 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -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 diff --git a/litellm/integrations/otel/model/db_endpoint.py b/litellm/integrations/otel/model/db_endpoint.py new file mode 100644 index 00000000000..562162a8f31 --- /dev/null +++ b/litellm/integrations/otel/model/db_endpoint.py @@ -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}) diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 3d585c36b67..ada2822ba66 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -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: diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 0f67f0e7a7c..08318f78b7c 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -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, } diff --git a/tests/test_litellm/integrations/otel/test_db_endpoint.py b/tests/test_litellm/integrations/otel/test_db_endpoint.py new file mode 100644 index 00000000000..5ab0a927b52 --- /dev/null +++ b/tests/test_litellm/integrations/otel/test_db_endpoint.py @@ -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 diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py index 82b074220fa..bb2d970e9c7 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_logger.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_logger.py @@ -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) diff --git a/tests/test_litellm/integrations/test_opentelemetry.py b/tests/test_litellm/integrations/test_opentelemetry.py index fa1c9fa8a79..28d2cb22e7c 100644 --- a/tests/test_litellm/integrations/test_opentelemetry.py +++ b/tests/test_litellm/integrations/test_opentelemetry.py @@ -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)