mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(proxy): resolve aliases in audit log responses
Adds object_alias, changed_by_user_email and changed_by_key_alias to
GET /audit and GET /audit/{id}, batch-resolved one query per entity
type per page with a fallback to the audit blobs for deleted objects,
and adds an object_team filter that matches by team id or team alias
This commit is contained in:
parent
973329e986
commit
e7264cacdd
4 changed files with 518 additions and 42 deletions
|
|
@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
|
|||
GET - /audit - Get all audit logs
|
||||
"""
|
||||
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, NamedTuple, Optional, Sequence, Tuple
|
||||
|
||||
#### AUDIT LOGGING ####
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
|
@ -16,11 +16,148 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
|
|||
PaginatedAuditLogResponse,
|
||||
)
|
||||
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmTableNames, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
_KEY_TABLE: Final[str] = LitellmTableNames.KEY_TABLE_NAME.value
|
||||
_TEAM_TABLE: Final[str] = LitellmTableNames.TEAM_TABLE_NAME.value
|
||||
_USER_TABLE: Final[str] = LitellmTableNames.USER_TABLE_NAME.value
|
||||
_ORG_TABLE: Final[str] = "LiteLLM_OrganizationTable"
|
||||
_MODEL_TABLE: Final[str] = LitellmTableNames.PROXY_MODEL_TABLE_NAME.value
|
||||
|
||||
_BLOB_ALIAS_KEYS: Final[Dict[str, Tuple[str, ...]]] = {
|
||||
_KEY_TABLE: ("key_alias",),
|
||||
_TEAM_TABLE: ("team_alias",),
|
||||
_USER_TABLE: ("user_alias", "user_email"),
|
||||
_ORG_TABLE: ("organization_alias",),
|
||||
_MODEL_TABLE: ("model_name",),
|
||||
}
|
||||
|
||||
|
||||
class _AliasMaps(NamedTuple):
|
||||
key_alias_by_token: Dict[str, str]
|
||||
team_alias_by_id: Dict[str, str]
|
||||
user_alias_by_id: Dict[str, str]
|
||||
user_email_by_id: Dict[str, str]
|
||||
org_alias_by_id: Dict[str, str]
|
||||
model_name_by_id: Dict[str, str]
|
||||
|
||||
|
||||
def _object_ids_for_table(audit_logs: Sequence[AuditLogResponse], table_name: str) -> frozenset:
|
||||
return frozenset(log.object_id for log in audit_logs if log.table_name == table_name and log.object_id)
|
||||
|
||||
|
||||
async def _fetch_alias_maps(prisma_client: "PrismaClient", audit_logs: Sequence[AuditLogResponse]) -> _AliasMaps:
|
||||
tokens: Final = _object_ids_for_table(audit_logs, _KEY_TABLE) | frozenset(
|
||||
log.changed_by_api_key for log in audit_logs if log.changed_by_api_key
|
||||
)
|
||||
user_ids: Final = _object_ids_for_table(audit_logs, _USER_TABLE) | frozenset(
|
||||
log.changed_by for log in audit_logs if log.changed_by
|
||||
)
|
||||
team_ids: Final = _object_ids_for_table(audit_logs, _TEAM_TABLE)
|
||||
org_ids: Final = _object_ids_for_table(audit_logs, _ORG_TABLE)
|
||||
model_ids: Final = _object_ids_for_table(audit_logs, _MODEL_TABLE)
|
||||
|
||||
key_rows: Final = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(where={"token": {"in": list(tokens)}})
|
||||
if tokens
|
||||
else []
|
||||
)
|
||||
user_rows: Final = (
|
||||
await prisma_client.db.litellm_usertable.find_many(where={"user_id": {"in": list(user_ids)}})
|
||||
if user_ids
|
||||
else []
|
||||
)
|
||||
team_rows: Final = (
|
||||
await prisma_client.db.litellm_teamtable.find_many(where={"team_id": {"in": list(team_ids)}})
|
||||
if team_ids
|
||||
else []
|
||||
)
|
||||
org_rows: Final = (
|
||||
await prisma_client.db.litellm_organizationtable.find_many(where={"organization_id": {"in": list(org_ids)}})
|
||||
if org_ids
|
||||
else []
|
||||
)
|
||||
model_rows: Final = (
|
||||
await prisma_client.db.litellm_proxymodeltable.find_many(where={"model_id": {"in": list(model_ids)}})
|
||||
if model_ids
|
||||
else []
|
||||
)
|
||||
|
||||
return _AliasMaps(
|
||||
key_alias_by_token={row.token: row.key_alias for row in key_rows if row.key_alias},
|
||||
team_alias_by_id={row.team_id: row.team_alias for row in team_rows if row.team_alias},
|
||||
user_alias_by_id={row.user_id: row.user_alias for row in user_rows if row.user_alias},
|
||||
user_email_by_id={row.user_id: row.user_email for row in user_rows if row.user_email},
|
||||
org_alias_by_id={row.organization_id: row.organization_alias for row in org_rows if row.organization_alias},
|
||||
model_name_by_id={row.model_id: row.model_name for row in model_rows if row.model_name},
|
||||
)
|
||||
|
||||
|
||||
def _db_object_alias(log: AuditLogResponse, aliases: _AliasMaps) -> str | None:
|
||||
if log.table_name == _KEY_TABLE:
|
||||
return aliases.key_alias_by_token.get(log.object_id)
|
||||
if log.table_name == _TEAM_TABLE:
|
||||
return aliases.team_alias_by_id.get(log.object_id)
|
||||
if log.table_name == _USER_TABLE:
|
||||
return aliases.user_alias_by_id.get(log.object_id) or aliases.user_email_by_id.get(log.object_id)
|
||||
if log.table_name == _ORG_TABLE:
|
||||
return aliases.org_alias_by_id.get(log.object_id)
|
||||
if log.table_name == _MODEL_TABLE:
|
||||
return aliases.model_name_by_id.get(log.object_id)
|
||||
return None
|
||||
|
||||
|
||||
def _alias_from_blobs(log: AuditLogResponse, blob_keys: Tuple[str, ...]) -> str | None:
|
||||
for blob in (log.updated_values, log.before_value):
|
||||
if not isinstance(blob, dict):
|
||||
continue
|
||||
for blob_key in blob_keys:
|
||||
value = blob.get(blob_key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _enrich_audit_log(log: AuditLogResponse, aliases: _AliasMaps) -> AuditLogResponse:
|
||||
object_alias: Final = _db_object_alias(log, aliases) or _alias_from_blobs(
|
||||
log, _BLOB_ALIAS_KEYS.get(log.table_name, ())
|
||||
)
|
||||
return log.model_copy(
|
||||
update={
|
||||
"object_alias": object_alias,
|
||||
"changed_by_user_email": aliases.user_email_by_id.get(log.changed_by),
|
||||
"changed_by_key_alias": aliases.key_alias_by_token.get(log.changed_by_api_key),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _enrich_audit_logs(
|
||||
prisma_client: "PrismaClient", audit_logs: Sequence[AuditLogResponse]
|
||||
) -> List[AuditLogResponse]:
|
||||
if not audit_logs:
|
||||
return []
|
||||
aliases: Final = await _fetch_alias_maps(prisma_client, audit_logs)
|
||||
return [_enrich_audit_log(log, aliases) for log in audit_logs]
|
||||
|
||||
|
||||
async def _build_object_team_condition(prisma_client: "PrismaClient", object_team: str) -> Dict[str, Any]:
|
||||
team_rows: Final = await prisma_client.db.litellm_teamtable.find_many(
|
||||
where={"team_alias": {"contains": object_team}}
|
||||
)
|
||||
match_values: Final = dict.fromkeys([object_team, *(row.team_id for row in team_rows)])
|
||||
return {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", object_team),
|
||||
*(_build_json_field_or_condition("team_id", value) for value in match_values),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
|
||||
"""
|
||||
|
|
@ -53,27 +190,24 @@ async def get_audit_logs(
|
|||
page: int = Query(1, ge=1),
|
||||
page_size: int = Query(10, ge=1, le=100),
|
||||
# Filter parameters
|
||||
changed_by: Optional[str] = Query(
|
||||
None, description="Filter by user or system that performed the action"
|
||||
),
|
||||
changed_by_api_key: Optional[str] = Query(
|
||||
None, description="Filter by API key hash that performed the action"
|
||||
),
|
||||
action: Optional[str] = Query(
|
||||
None, description="Filter by action type (create, update, delete)"
|
||||
),
|
||||
table_name: Optional[str] = Query(
|
||||
None, description="Filter by table name that was modified"
|
||||
),
|
||||
object_id: Optional[str] = Query(
|
||||
None, description="Filter by ID of the object that was modified"
|
||||
),
|
||||
changed_by: Optional[str] = Query(None, description="Filter by user or system that performed the action"),
|
||||
changed_by_api_key: Optional[str] = Query(None, description="Filter by API key hash that performed the action"),
|
||||
action: Optional[str] = Query(None, description="Filter by action type (create, update, delete)"),
|
||||
table_name: Optional[str] = Query(None, description="Filter by table name that was modified"),
|
||||
object_id: Optional[str] = Query(None, description="Filter by ID of the object that was modified"),
|
||||
start_date: Optional[str] = Query(None, description="Filter logs after this date"),
|
||||
end_date: Optional[str] = Query(None, description="Filter logs before this date"),
|
||||
object_team_id: Optional[str] = Query(
|
||||
None,
|
||||
description="Filter by team_id present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
),
|
||||
object_team: str | None = Query(
|
||||
None,
|
||||
description=(
|
||||
"Filter by team id or alias: matches team_id or team_alias present in before_value or "
|
||||
"updated_values JSON, or teams whose team_alias contains this value (PostgreSQL only)"
|
||||
),
|
||||
),
|
||||
object_key_hash: Optional[str] = Query(
|
||||
None,
|
||||
description="Filter by token (key hash) present in before_value or updated_values JSON (PostgreSQL only)",
|
||||
|
|
@ -90,8 +224,10 @@ async def get_audit_logs(
|
|||
|
||||
Returns a paginated response of audit logs matching the specified filters.
|
||||
|
||||
Note: object_team_id and object_key_hash use Prisma JSON path filtering,
|
||||
which requires PostgreSQL.
|
||||
Note: object_team_id, object_team and object_key_hash use Prisma JSON path
|
||||
filtering, which requires PostgreSQL. object_team matches a team_id or
|
||||
team_alias in the audit blobs, or any team whose team_alias contains the
|
||||
value.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -131,6 +267,10 @@ async def get_audit_logs(
|
|||
where_conditions["AND"] = where_conditions.get("AND", []) + [
|
||||
_build_json_field_or_condition("token", object_key_hash)
|
||||
]
|
||||
if object_team:
|
||||
where_conditions["AND"] = where_conditions.get("AND", []) + [
|
||||
await _build_object_team_condition(prisma_client, object_team)
|
||||
]
|
||||
|
||||
# Build sort conditions
|
||||
order_by: Dict[str, Any] = {}
|
||||
|
|
@ -151,13 +291,14 @@ async def get_audit_logs(
|
|||
total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions)
|
||||
total_pages = -(-total_count // page_size) # Ceiling division
|
||||
|
||||
enriched_logs: Final = await _enrich_audit_logs(
|
||||
prisma_client,
|
||||
[AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs] if audit_logs else [],
|
||||
)
|
||||
|
||||
# Return paginated response
|
||||
return PaginatedAuditLogResponse(
|
||||
audit_logs=[
|
||||
AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs
|
||||
]
|
||||
if audit_logs
|
||||
else [],
|
||||
audit_logs=enriched_logs,
|
||||
total=total_count,
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
|
|
@ -175,9 +316,7 @@ async def get_audit_logs(
|
|||
500: {"description": "Database connection error"},
|
||||
},
|
||||
)
|
||||
async def get_audit_log_by_id(
|
||||
id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)
|
||||
):
|
||||
async def get_audit_log_by_id(id: str, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth)):
|
||||
"""
|
||||
Get detailed information about a specific audit log entry by its ID.
|
||||
|
||||
|
|
@ -202,9 +341,7 @@ async def get_audit_log_by_id(
|
|||
audit_log = await prisma_client.db.litellm_auditlog.find_unique(where={"id": id})
|
||||
|
||||
if audit_log is None:
|
||||
raise HTTPException(
|
||||
status_code=404, detail={"message": f"Audit log with ID {id} not found"}
|
||||
)
|
||||
raise HTTPException(status_code=404, detail={"message": f"Audit log with ID {id} not found"})
|
||||
|
||||
# Convert to response model
|
||||
return AuditLogResponse(**audit_log.model_dump())
|
||||
enriched_logs: Final = await _enrich_audit_logs(prisma_client, [AuditLogResponse(**audit_log.model_dump())])
|
||||
return enriched_logs[0]
|
||||
|
|
|
|||
|
|
@ -16,15 +16,16 @@ class AuditLogResponse(BaseModel):
|
|||
object_id: str
|
||||
before_value: Optional[Dict[str, Any]] = None
|
||||
updated_values: Optional[Dict[str, Any]] = None
|
||||
object_alias: str | None = None
|
||||
changed_by_user_email: str | None = None
|
||||
changed_by_key_alias: str | None = None
|
||||
|
||||
|
||||
class PaginatedAuditLogResponse(BaseModel):
|
||||
"""Response model for paginated audit logs"""
|
||||
|
||||
audit_logs: List[AuditLogResponse]
|
||||
total: int = Field(
|
||||
..., description="Total number of audit logs matching the filters"
|
||||
)
|
||||
total: int = Field(..., description="Total number of audit logs matching the filters")
|
||||
page: int = Field(..., description="Current page number")
|
||||
page_size: int = Field(..., description="Number of items per page")
|
||||
total_pages: int = Field(..., description="Total number of pages")
|
||||
|
|
|
|||
|
|
@ -1,4 +1,4 @@
|
|||
from datetime import datetime, timedelta
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, patch
|
||||
|
||||
import pytest
|
||||
|
|
@ -34,6 +34,11 @@ def mock_prisma_client():
|
|||
mock.db.litellm_auditlog.find_many = AsyncMock()
|
||||
mock.db.litellm_auditlog.find_unique = AsyncMock()
|
||||
mock.db.litellm_auditlog.count = AsyncMock()
|
||||
mock.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_organizationtable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
yield mock
|
||||
|
||||
|
||||
|
|
@ -41,9 +46,7 @@ def mock_prisma_client():
|
|||
async def test_get_audit_logs(mock_prisma_client):
|
||||
"""Test successful retrieval of audit logs with pagination"""
|
||||
# Mock the database responses
|
||||
mock_prisma_client.db.litellm_auditlog.find_many.return_value = [
|
||||
AuditLogResponse(**MOCK_AUDIT_LOG)
|
||||
]
|
||||
mock_prisma_client.db.litellm_auditlog.find_many.return_value = [AuditLogResponse(**MOCK_AUDIT_LOG)]
|
||||
mock_prisma_client.db.litellm_auditlog.count.return_value = 1
|
||||
|
||||
# Mock the auth dependency
|
||||
|
|
@ -80,9 +83,7 @@ async def test_get_audit_logs(mock_prisma_client):
|
|||
async def test_get_audit_log_by_id(mock_prisma_client):
|
||||
"""Test successful retrieval of a specific audit log by ID"""
|
||||
# Mock the database response
|
||||
mock_prisma_client.db.litellm_auditlog.find_unique.return_value = AuditLogResponse(
|
||||
**MOCK_AUDIT_LOG
|
||||
)
|
||||
mock_prisma_client.db.litellm_auditlog.find_unique.return_value = AuditLogResponse(**MOCK_AUDIT_LOG)
|
||||
|
||||
# Mock the auth dependency
|
||||
with patch("litellm.proxy.auth.user_api_key_auth.user_api_key_auth") as mock_auth:
|
||||
|
|
|
|||
|
|
@ -0,0 +1,337 @@
|
|||
"""
|
||||
Tests for audit log alias enrichment and the combined object_team filter (LIT-4997).
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi import FastAPI
|
||||
from fastapi.testclient import TestClient
|
||||
from litellm_enterprise.proxy.audit_logging_endpoints import (
|
||||
_build_json_field_or_condition,
|
||||
_build_object_team_condition,
|
||||
_enrich_audit_logs,
|
||||
)
|
||||
from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_router
|
||||
from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse
|
||||
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
|
||||
class FakeTable:
|
||||
def __init__(self, rows=()):
|
||||
self.rows = list(rows)
|
||||
self.find_many_calls = []
|
||||
|
||||
async def find_many(self, where=None, **kwargs):
|
||||
self.find_many_calls.append(where)
|
||||
return self.rows
|
||||
|
||||
|
||||
class FakeAuditLogTable(FakeTable):
|
||||
def __init__(self, rows=()):
|
||||
super().__init__(rows)
|
||||
self.count_calls = []
|
||||
|
||||
async def count(self, where=None):
|
||||
self.count_calls.append(where)
|
||||
return len(self.rows)
|
||||
|
||||
async def find_unique(self, where):
|
||||
return next((row for row in self.rows if row.id == where["id"]), None)
|
||||
|
||||
|
||||
class FakeDb:
|
||||
def __init__(
|
||||
self,
|
||||
audit_logs=(),
|
||||
keys=(),
|
||||
users=(),
|
||||
teams=(),
|
||||
orgs=(),
|
||||
models=(),
|
||||
):
|
||||
self.litellm_auditlog = FakeAuditLogTable(audit_logs)
|
||||
self.litellm_verificationtoken = FakeTable(keys)
|
||||
self.litellm_usertable = FakeTable(users)
|
||||
self.litellm_teamtable = FakeTable(teams)
|
||||
self.litellm_organizationtable = FakeTable(orgs)
|
||||
self.litellm_proxymodeltable = FakeTable(models)
|
||||
|
||||
|
||||
class FakePrismaClient:
|
||||
def __init__(self, db: FakeDb):
|
||||
self.db = db
|
||||
|
||||
|
||||
def make_log(**overrides) -> AuditLogResponse:
|
||||
defaults = {
|
||||
"id": "log-1",
|
||||
"updated_at": datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
"changed_by": "",
|
||||
"changed_by_api_key": "",
|
||||
"action": "updated",
|
||||
"table_name": "LiteLLM_TeamTable",
|
||||
"object_id": "obj-1",
|
||||
"before_value": None,
|
||||
"updated_values": None,
|
||||
}
|
||||
return AuditLogResponse(**{**defaults, **overrides})
|
||||
|
||||
|
||||
async def test_enrichment_resolves_object_alias_per_table_name():
|
||||
"""object_alias comes from the right table per table_name; changed_by fields resolve too."""
|
||||
logs = [
|
||||
make_log(id="l1", table_name="LiteLLM_VerificationToken", object_id="hash-1"),
|
||||
make_log(id="l2", table_name="LiteLLM_TeamTable", object_id="team-1"),
|
||||
make_log(id="l3", table_name="LiteLLM_UserTable", object_id="user-1"),
|
||||
make_log(id="l4", table_name="LiteLLM_UserTable", object_id="user-2"),
|
||||
make_log(id="l5", table_name="LiteLLM_OrganizationTable", object_id="org-1"),
|
||||
make_log(id="l6", table_name="LiteLLM_ProxyModelTable", object_id="model-1"),
|
||||
make_log(
|
||||
id="l7",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
),
|
||||
]
|
||||
db = FakeDb(
|
||||
keys=[
|
||||
SimpleNamespace(token="hash-1", key_alias="prod-key"),
|
||||
SimpleNamespace(token="hash-admin", key_alias="admin-key"),
|
||||
],
|
||||
users=[
|
||||
SimpleNamespace(user_id="user-1", user_alias="Alice", user_email="alice@example.com"),
|
||||
SimpleNamespace(user_id="user-2", user_alias=None, user_email="bob@example.com"),
|
||||
SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com"),
|
||||
],
|
||||
teams=[SimpleNamespace(team_id="team-1", team_alias="ml-team")],
|
||||
orgs=[SimpleNamespace(organization_id="org-1", organization_alias="acme-org")],
|
||||
models=[SimpleNamespace(model_id="model-1", model_name="gpt-5.2")],
|
||||
)
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
by_id = {log.id: log for log in enriched}
|
||||
|
||||
assert by_id["l1"].object_alias == "prod-key"
|
||||
assert by_id["l2"].object_alias == "ml-team"
|
||||
assert by_id["l3"].object_alias == "Alice"
|
||||
assert by_id["l4"].object_alias == "bob@example.com"
|
||||
assert by_id["l5"].object_alias == "acme-org"
|
||||
assert by_id["l6"].object_alias == "gpt-5.2"
|
||||
assert by_id["l7"].changed_by_user_email == "admin@example.com"
|
||||
assert by_id["l7"].changed_by_key_alias == "admin-key"
|
||||
assert by_id["l1"].changed_by_user_email is None
|
||||
assert by_id["l1"].changed_by_key_alias is None
|
||||
|
||||
|
||||
async def test_enrichment_runs_one_query_per_entity_type():
|
||||
"""A page with many rows triggers at most one find_many per entity table, ids batched via `in`."""
|
||||
logs = [
|
||||
make_log(id="l1", table_name="LiteLLM_VerificationToken", object_id="hash-1", changed_by="u1"),
|
||||
make_log(id="l2", table_name="LiteLLM_VerificationToken", object_id="hash-2", changed_by="u2"),
|
||||
make_log(
|
||||
id="l3",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
changed_by="u1",
|
||||
changed_by_api_key="hash-caller",
|
||||
),
|
||||
]
|
||||
db = FakeDb()
|
||||
|
||||
await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
|
||||
assert len(db.litellm_verificationtoken.find_many_calls) == 1
|
||||
assert len(db.litellm_usertable.find_many_calls) == 1
|
||||
assert len(db.litellm_teamtable.find_many_calls) == 1
|
||||
assert len(db.litellm_organizationtable.find_many_calls) == 0
|
||||
assert len(db.litellm_proxymodeltable.find_many_calls) == 0
|
||||
assert set(db.litellm_verificationtoken.find_many_calls[0]["token"]["in"]) == {
|
||||
"hash-1",
|
||||
"hash-2",
|
||||
"hash-caller",
|
||||
}
|
||||
assert set(db.litellm_usertable.find_many_calls[0]["user_id"]["in"]) == {"u1", "u2"}
|
||||
assert db.litellm_teamtable.find_many_calls[0] == {"team_id": {"in": ["team-1"]}}
|
||||
|
||||
|
||||
async def test_enrichment_falls_back_to_blobs_for_deleted_objects():
|
||||
"""When DB lookups miss (deleted objects), aliases come from updated_values then before_value."""
|
||||
logs = [
|
||||
make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="gone-team",
|
||||
action="deleted",
|
||||
before_value={"team_id": "gone-team", "team_alias": "old-team"},
|
||||
),
|
||||
make_log(
|
||||
id="l2",
|
||||
table_name="LiteLLM_VerificationToken",
|
||||
object_id="gone-hash",
|
||||
before_value={"key_alias": "old-alias"},
|
||||
updated_values={"key_alias": "new-alias"},
|
||||
),
|
||||
make_log(
|
||||
id="l3",
|
||||
table_name="LiteLLM_UserTable",
|
||||
object_id="gone-user",
|
||||
updated_values={"user_email": "gone@example.com"},
|
||||
),
|
||||
make_log(
|
||||
id="l4",
|
||||
table_name="LiteLLM_OrganizationTable",
|
||||
object_id="gone-org",
|
||||
before_value={"organization_alias": "old-org"},
|
||||
),
|
||||
make_log(id="l5", table_name="SomeUnknownTable", object_id="x", updated_values={"team_alias": "nope"}),
|
||||
]
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(FakeDb()), logs)
|
||||
by_id = {log.id: log for log in enriched}
|
||||
|
||||
assert by_id["l1"].object_alias == "old-team"
|
||||
assert by_id["l2"].object_alias == "new-alias"
|
||||
assert by_id["l3"].object_alias == "gone@example.com"
|
||||
assert by_id["l4"].object_alias == "old-org"
|
||||
assert by_id["l5"].object_alias is None
|
||||
|
||||
|
||||
async def test_enrichment_db_lookup_wins_over_blob():
|
||||
"""A live DB row beats a stale alias captured in the audit blobs."""
|
||||
logs = [
|
||||
make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
before_value={"team_alias": "stale-alias"},
|
||||
)
|
||||
]
|
||||
db = FakeDb(teams=[SimpleNamespace(team_id="team-1", team_alias="current-alias")])
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
|
||||
assert enriched[0].object_alias == "current-alias"
|
||||
|
||||
|
||||
async def test_build_object_team_condition_matches_id_and_alias():
|
||||
"""object_team ORs the raw value with every team_id whose team_alias contains it."""
|
||||
db = FakeDb(
|
||||
teams=[
|
||||
SimpleNamespace(team_id="team-1", team_alias="prod-team"),
|
||||
SimpleNamespace(team_id="team-2", team_alias="prod-eu"),
|
||||
]
|
||||
)
|
||||
|
||||
condition = await _build_object_team_condition(FakePrismaClient(db), "prod")
|
||||
|
||||
assert db.litellm_teamtable.find_many_calls == [{"team_alias": {"contains": "prod"}}]
|
||||
assert condition == {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "prod"),
|
||||
_build_json_field_or_condition("team_id", "prod"),
|
||||
_build_json_field_or_condition("team_id", "team-1"),
|
||||
_build_json_field_or_condition("team_id", "team-2"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
async def test_build_object_team_condition_deleted_team_matches_blob_alias():
|
||||
"""With no live team rows the condition still matches blob team_alias and the raw value as team_id."""
|
||||
condition = await _build_object_team_condition(FakePrismaClient(FakeDb()), "gone-team")
|
||||
|
||||
assert condition == {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "gone-team"),
|
||||
_build_json_field_or_condition("team_id", "gone-team"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _client_for(db: FakeDb) -> TestClient:
|
||||
app = FastAPI()
|
||||
app.include_router(audit_router)
|
||||
app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(user_role="proxy_admin")
|
||||
return TestClient(app)
|
||||
|
||||
|
||||
def test_get_audit_logs_object_team_filter_and_enrichment():
|
||||
"""GET /audit?object_team=... ANDs in the combined id/alias condition and returns enriched rows."""
|
||||
audit_row = make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
changed_by="admin-user",
|
||||
updated_values={"team_id": "team-1"},
|
||||
)
|
||||
db = FakeDb(
|
||||
audit_logs=[audit_row],
|
||||
users=[SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com")],
|
||||
teams=[SimpleNamespace(team_id="team-1", team_alias="prod-team")],
|
||||
)
|
||||
client = _client_for(db)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit?object_team=prod")
|
||||
|
||||
assert response.status_code == 200
|
||||
where = db.litellm_auditlog.find_many_calls[0]
|
||||
assert where["AND"] == [
|
||||
{
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "prod"),
|
||||
_build_json_field_or_condition("team_id", "prod"),
|
||||
_build_json_field_or_condition("team_id", "team-1"),
|
||||
]
|
||||
}
|
||||
]
|
||||
log = response.json()["audit_logs"][0]
|
||||
assert log["object_alias"] == "prod-team"
|
||||
assert log["changed_by_user_email"] == "admin@example.com"
|
||||
assert log["changed_by_key_alias"] is None
|
||||
|
||||
|
||||
def test_get_audit_logs_object_team_id_filter_unchanged():
|
||||
"""The pre-existing object_team_id param still builds its exact condition, no alias lookup."""
|
||||
db = FakeDb()
|
||||
client = _client_for(db)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit?object_team_id=team-1")
|
||||
|
||||
assert response.status_code == 200
|
||||
where = db.litellm_auditlog.find_many_calls[0]
|
||||
assert where["AND"] == [_build_json_field_or_condition("team_id", "team-1")]
|
||||
assert db.litellm_teamtable.find_many_calls == []
|
||||
|
||||
|
||||
def test_get_audit_log_by_id_is_enriched():
|
||||
"""GET /audit/{id} carries the same alias enrichment as the list endpoint."""
|
||||
audit_row = make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_VerificationToken",
|
||||
object_id="gone-hash",
|
||||
action="deleted",
|
||||
before_value={"key_alias": "deleted-key"},
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
)
|
||||
db = FakeDb(
|
||||
audit_logs=[audit_row],
|
||||
keys=[SimpleNamespace(token="hash-admin", key_alias="admin-key")],
|
||||
users=[SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com")],
|
||||
)
|
||||
client = _client_for(db)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit/l1")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["object_alias"] == "deleted-key"
|
||||
assert body["changed_by_user_email"] == "admin@example.com"
|
||||
assert body["changed_by_key_alias"] == "admin-key"
|
||||
Loading…
Add table
Reference in a new issue