diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 27576bf51a..64be94b09a 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -2273,6 +2273,50 @@ async def user_update(request: Request): pass +@app.get( + "/user/export", tags=["user management"], dependencies=[Depends(user_api_key_auth)] +) +async def export_user_table(): + """ + Exports all user rows as a csv sheet + """ + import csv + from io import StringIO + + users = await prisma_client.get_data(table_name="user", query_type="find_all") + + # Create an in-memory string buffer + buffer = StringIO() + fieldnames = ["user_id", "max_budget", "spend", "user_email", "models"] + writer = csv.DictWriter(buffer, fieldnames=fieldnames) + + # Write the header + writer.writeheader() + + # Write user data + for user in users: + writer.writerow( + { + "user_id": user.user_id, + "max_budget": user.max_budget, + "spend": user.spend, + "user_email": user.user_email, + "models": ",".join( + user.models + ), # convert list to comma-separated string + } + ) + + # Set the buffer's position to the start + buffer.seek(0) + + # Create a CSV response + response = Response(content=buffer.getvalue(), media_type="text/csv") + response.headers["Content-Disposition"] = "attachment; filename=export_users.csv" + + return response + + #### MODEL MANAGEMENT #### diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 82e9e93355..ef8290ce88 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -386,12 +386,21 @@ class PrismaClient: status_code=status.HTTP_401_UNAUTHORIZED, detail="Authentication Error: invalid user key - token does not exist", ) - elif user_id is not None: - response = await self.db.litellm_usertable.find_unique( # type: ignore - where={ - "user_id": user_id, - } - ) + elif user_id is not None or ( + table_name is not None and table_name == "user" + ): + if user_id is not None and query_type == "find_unique": + response = await self.db.litellm_usertable.find_unique( # type: ignore + where={ + "user_id": user_id, + } + ) + elif query_type == "find_all" and user_id is not None: + response = await self.db.litellm_usertable.find_many( + where={"user_id": user_id} + ) + elif query_type == "find_all": + response = await self.db.litellm_usertable.find_many() return response except Exception as e: print_verbose(f"LiteLLM Prisma Client Exception: {e}")