Skip to content

Commit

Permalink
[ENH]: add tenant, database, and collection IDs to trace when availab…
Browse files Browse the repository at this point in the history
…le (#2889)
  • Loading branch information
codetheweb authored Oct 1, 2024
1 parent 9a46096 commit 91ba2c9
Show file tree
Hide file tree
Showing 2 changed files with 17 additions and 1 deletion.
15 changes: 15 additions & 0 deletions chromadb/server/fastapi/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,7 @@
from chromadb.telemetry.opentelemetry import (
OpenTelemetryClient,
OpenTelemetryGranularity,
add_attributes_to_current_span,
trace_method,
)
from chromadb.types import Collection as CollectionModel
Expand Down Expand Up @@ -383,6 +384,13 @@ def auth_and_get_tenant_and_database_for_request(
- The user has access to a single tenant and/or single database.
"""
if not self.authn_provider:
add_attributes_to_current_span(
{
"tenant": tenant,
"database": database,
"collection": collection,
}
)
return (tenant, database)

user_identity = self.authn_provider.authenticate_or_raise(dict(headers))
Expand All @@ -407,6 +415,13 @@ def auth_and_get_tenant_and_database_for_request(
)

self.authz_provider.authorize_or_raise(user_identity, action, authz_resource)
add_attributes_to_current_span(
{
"tenant": tenant,
"database": database,
"collection": collection,
}
)
return (tenant, database)

@trace_method("FastAPI.create_database", OpenTelemetryGranularity.OPERATION)
Expand Down
3 changes: 2 additions & 1 deletion chromadb/telemetry/opentelemetry/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -166,6 +166,7 @@ def add_attributes_to_current_span(
Sequence[bool],
Sequence[float],
Sequence[int],
None,
],
]
) -> None:
Expand All @@ -176,4 +177,4 @@ def add_attributes_to_current_span(
if not tracer:
return
span = trace.get_current_span()
span.set_attributes(attributes)
span.set_attributes({k: v for k, v in attributes.items() if v is not None})

0 comments on commit 91ba2c9

Please sign in to comment.