fix(proxy): degrade to flat groups when membership table isn't ready

8 existing tests broke because they pass a plain MagicMock as the
prisma_client and the new awaited find_many trips on await.

same can hit prod - the proxy can boot before litellm-proxy-extras
finishes prisma migrate deploy, so the model is briefly missing
from the prisma client.

wrap find_many in a narrow try/except (AttributeError, TypeError),
return {} on miss. auth path falls back to today's flat-group
behavior instead of 500'ing. two regression tests so we don't
forget.

refs #28032
This commit is contained in:
Ashwin Upadhyay 2026-05-17 00:48:06 +05:30
parent a6139fab89
commit 917b28588c
2 changed files with 43 additions and 1 deletions

View file

@ -86,8 +86,23 @@ async def get_group_memberships_from_db(
"""
Build parent_group -> [child_groups] map from the membership table.
Single query, in-memory bucketing - no N+1.
Resilient by design: if the table isn't available (Prisma client predates
this migration, the proxy started before `prisma migrate deploy` finished,
or the membership Prisma model was stripped from a downstream build) we
return an empty map. The auth path then falls back to today's flat-group
semantics instead of 500-ing the whole request.
"""
rows = await prisma_client.db.litellm_accessgroupmembership.find_many()
try:
rows = await prisma_client.db.litellm_accessgroupmembership.find_many()
except (AttributeError, TypeError) as e:
verbose_proxy_logger.debug(
"litellm_accessgroupmembership unavailable - "
"skipping nested group resolution: %s",
e,
)
return {}
memberships: Dict[str, List[str]] = {}
for row in rows:
memberships.setdefault(row.parent_group, []).append(row.child_group)

View file

@ -86,6 +86,33 @@ async def test_get_group_memberships_empty_table_returns_empty_dict():
assert await get_group_memberships_from_db(prisma_client=prisma) == {}
@pytest.mark.asyncio
async def test_get_group_memberships_returns_empty_when_table_unavailable():
"""
Defensive fallback: if the prisma client predates the migration (or any
other reason find_many raises AttributeError/TypeError), we return {} so
auth-path callers fall back to today's flat-group behavior.
"""
# Plain MagicMock - prisma_client.db.litellm_accessgroupmembership.find_many()
# returns a non-awaitable MagicMock, raising TypeError on `await`.
plain = MagicMock()
assert await get_group_memberships_from_db(prisma_client=plain) == {}
@pytest.mark.asyncio
async def test_get_group_memberships_returns_empty_when_table_attribute_missing():
"""
AttributeError on the table accessor (model stripped from a downstream
Prisma client build) also falls back to empty.
"""
class NoMembershipTable:
class db:
pass
assert await get_group_memberships_from_db(prisma_client=NoMembershipTable()) == {}
# ---------------------------------------------------------------------------
# upsert_group_memberships
# ---------------------------------------------------------------------------