didi-lot1-ai/ai_platform/modules/dashboard/src/dashboard/cli.py

103 lines
3.3 KiB
Python

"""CLI utilities for admin tasks (create user, list users, etc.).
Usage (inside the container):
python -m dashboard.cli create-user didi
python -m dashboard.cli list-users
python -m dashboard.cli delete-user didi
"""
import argparse
import asyncio
import sys
from sqlalchemy import delete, select
from dashboard.auth import generate_token, upsert_user
from dashboard.config import SettingsCache
from dashboard.db.models import User
from dashboard.db.session import close_engine, get_session_factory, init_engine
async def _create_user(username: str, email: str | None, role: str) -> None:
settings = SettingsCache.get()
init_engine(settings)
factory = get_session_factory()
async with factory() as session:
token = generate_token()
user = await upsert_user(session, username, token, email=email, role=role)
print("User created:")
print(f" id: {user.id}")
print(f" username: {user.username}")
print(f" role: {user.role}")
print(f" email: {user.email or ''}")
print()
print("Bearer token (SAVE THIS — it is not shown again):")
print(f" {token}")
await close_engine()
async def _list_users() -> None:
settings = SettingsCache.get()
init_engine(settings)
factory = get_session_factory()
async with factory() as session:
rows = (await session.execute(select(User).order_by(User.id))).scalars().all()
if not rows:
print("No users yet.")
await close_engine()
return
print(f"{'ID':<4} {'USERNAME':<20} {'ROLE':<10} {'EMAIL':<30} LAST LOGIN")
print("-" * 100)
for u in rows:
last = u.last_login.strftime("%Y-%m-%d %H:%M") if u.last_login else "never"
print(
f"{u.id:<4} {u.username:<20} {u.role:<10} {(u.email or ''):<30} {last}"
)
await close_engine()
async def _delete_user(username: str) -> None:
settings = SettingsCache.get()
init_engine(settings)
factory = get_session_factory()
async with factory() as session:
result = await session.execute(
delete(User).where(User.username == username)
)
await session.commit()
if result.rowcount:
print(f"Deleted user {username}")
else:
print(f"User {username} not found")
await close_engine()
def main() -> None:
parser = argparse.ArgumentParser(prog="dashboard.cli", description="Dashboard admin")
sub = parser.add_subparsers(dest="command", required=True)
p_create = sub.add_parser("create-user", help="Create or update a user")
p_create.add_argument("username")
p_create.add_argument("--email", default=None)
p_create.add_argument("--role", default="admin", choices=["admin", "viewer"])
sub.add_parser("list-users", help="List all users")
p_delete = sub.add_parser("delete-user", help="Delete a user by username")
p_delete.add_argument("username")
args = parser.parse_args()
if args.command == "create-user":
asyncio.run(_create_user(args.username, args.email, args.role))
elif args.command == "list-users":
asyncio.run(_list_users())
elif args.command == "delete-user":
asyncio.run(_delete_user(args.username))
else:
parser.print_help()
sys.exit(1)
if __name__ == "__main__":
main()