feat(proxy_server.py): new /user/export endpoint

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 ####
 
 

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}")