refactor(proxy): read org admin's own team ids via get_user_object(check_db_only=True)

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/management_endpoints/team_endpoints.py

@@ -4858,9 +4858,24 @@ async def _get_org_admin_org_ids(
     return org_ids if org_ids else None
 
 
-async def _get_user_team_ids_from_db(user_id: str, prisma_client: PrismaClient) -> tuple[str, ...]:
-    user_row: Final = await _user_db(prisma_client).find_unique(where={"user_id": user_id})
-    return tuple(user_row.teams or ()) if user_row is not None else ()
+async def _get_user_team_ids_from_db(
+    user_id: str,
+    prisma_client: PrismaClient,
+    user_api_key_cache: UserApiKeyCache,
+    proxy_logging_obj: ProxyLogging,
+) -> tuple[str, ...]:
+    try:
+        user: Final = await get_user_object(
+            user_id=user_id,
+            prisma_client=prisma_client,
+            user_api_key_cache=user_api_key_cache,
+            user_id_upsert=False,
+            proxy_logging_obj=proxy_logging_obj,
+            check_db_only=True,
+        )
+    except ValueError:
+        return ()
+    return tuple(user.teams or ()) if user is not None else ()
 
 
 async def _build_team_list_where_conditions(
@@ -5080,7 +5095,12 @@ async def _enforce_list_team_v2_access(
             )
         is_own_query: Final = user_id is None or user_id == caller_user_id
         own_team_ids: Final = (
-            await _get_user_team_ids_from_db(user_id=caller_user_id, prisma_client=prisma_client)
+            await _get_user_team_ids_from_db(
+                user_id=caller_user_id,
+                prisma_client=prisma_client,
+                user_api_key_cache=user_api_key_cache,
+                proxy_logging_obj=proxy_logging_obj,
+            )
             if is_own_query
             else ()
         )

tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py

@@ -4162,7 +4162,14 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(
     prisma_client.db.litellm_teamtable.count = AsyncMock(side_effect=count)
     prisma_client.db.litellm_verificationtoken.group_by = AsyncMock(return_value=[])
     prisma_client.db.litellm_usertable.find_unique = AsyncMock(
-        return_value=SimpleNamespace(teams=["team_in_org_A", "team_in_org_B"])
+        return_value=LiteLLM_UserTable(
+            user_id="org_admin_user",
+            teams=["team_in_org_A", "team_in_org_B"],
+            organization_memberships=[
+                _org_membership("org_admin_user", "org_A", "org_admin"),
+                _org_membership("org_admin_user", "org_B", "internal_user"),
+            ],
+        )
     )
     monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
     monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache)
@@ -4193,7 +4200,9 @@ async def test_list_team_v2_org_admin_own_query_keeps_memberships_in_other_orgs(
     assert await list_teams(None) == own_view
     assert await list_teams("org_admin_user", search="team_in_org_B") == ["team_in_org_B"]
     assert await list_teams("other_user") == ["other_team_in_org_A"]
-    prisma_client.db.litellm_usertable.find_unique.assert_awaited_with(where={"user_id": "org_admin_user"})
+    prisma_client.db.litellm_usertable.find_unique.assert_awaited_with(
+        where={"user_id": "org_admin_user"}, include={"organization_memberships": True}
+    )
 
 
 @pytest.mark.asyncio