CostEvalPlatform/api/manage_tanant.py

# -*- coding: utf-8 -*-
# @Time    : 2024/11/19 下午8:04
# @FileName: manage_tanant.py
# @Software: PyCharm
from fastapi import APIRouter
from fastapi import HTTPException, Response, Depends, APIRouter
from typing import Optional, Annotated
from datetime import datetime, timedelta
from jose import JWTError, jwt
from sqlalchemy import delete

from sqlmodel import select

from models import Tenant, User, Project
from dependencies import *

router = APIRouter()

from fastapi import HTTPException, Response
from sqlalchemy.orm import Session
from models import Tenant, User  # 假设你已导入 Tenant 和 User 模型
from dependencies import SessionDep  # 假设 SessionDep 是数据库会话的依赖

#列举所有租户
@router.get("/api/s1/tenant")
async def get_tenant(response: Response, session: SessionDep):
    tenants = session.query(Tenant).all()  # 获取所有租户
    if not tenants:
        raise HTTPException(status_code=404, detail="No tenants found")

    tenant_data = []
    for tenant in tenants:
        # 获取该租户中 role=1 的第一个用户（如果存在）
        tenant_user = next((user for user in tenant.users if user.role == 1), None)

        # 获取该租户中除了 role=1 以外的用户数量
        user_num = len([user for user in tenant.users if user.role != 1])

        # 构建租户信息
        tenant_info = {
            "name": tenant.name,
            "username": tenant_user.username if tenant_user else None,  # 如果找到 role=1 的用户，返回其 username
            "user_num": user_num  # 除去 role=1 的用户数量
        }

        tenant_data.append(tenant_info)

    return {"tenants": tenant_data}

# 新增和修改租户
@router.post("/api/s1/tenant")
async def create_or_update_tenant(data: dict, session: SessionDep):
    name = data["name"]
    username = data["username"]
    password = data.get("password", "")  # 默认为空字符串

    # 验证是否缺少必要参数
    if not name or not username:
        raise HTTPException(status_code=400, detail="Need more name/username")

    # 查找用户
    user_query = select(User).where(User.username == username)
    existing_user = session.exec(user_query).first()

    # 如果密码为空，更新租户信息
    if password == "":
        print("密码为空")   #测试用
        # 如果用户不存在，返回错误
        if not existing_user:
            raise HTTPException(status_code=404, detail="User not found")
        else:
            # 如果找到了对应的 User
            # 使用 user.tenant_id 查找对应的 Tenant
            tenant = session.get(Tenant, existing_user.tenant_id)
            # 如果 Tenant 存在，更新 Tenant 的 name 字段
            if tenant:
                tenant.name = name
                session.commit()  # 提交更新
            else:
                raise HTTPException(status_code=404, detail="Tenant not found")
            return {"message": "Tenant and User update successfully"}
    else:
        print("密码不为空")  #测试用
        # 如果密码不为空，执行创建新租户和用户的操作
        if existing_user:
            # 如果用户已存在，返回错误
            raise HTTPException(status_code=409, detail="User already exists")

        # 检查租户是否已存在
        tenant_query = select(Tenant).where(Tenant.name == name)
        existing_tenant = session.exec(tenant_query).first()

        if existing_tenant:
            raise HTTPException(status_code=409, detail="Tenant name already exists")

        # 创建新租户
        tenant = Tenant(
            name=name,
            username=username,
            password=password,  # 实际使用时应加密密码
        )
        session.add(tenant)
        session.commit()
        session.refresh(tenant)

        # 创建新用户
        user = User(
            username=username,
            password=password,  # 同样需要加密密码
            role=1,  # 默认role为1
            tenant_id = tenant.id,
        )
        session.add(user)

        # 提交事务
        session.commit()
        session.refresh(tenant)

        return {"message": "Tenant and User added successfully"}

#删除租户
@router.delete("/api/s1/tenant")
async def delete_tenant(data: dict, session: SessionDep):
    tenant_name = data.get("name")

    if not tenant_name:
        raise HTTPException(status_code=400, detail="Tenant name is required")

    # 查找租户
    tenant = session.exec(
        select(Tenant).where(Tenant.name == tenant_name)).first()
    if not tenant:
        raise HTTPException(status_code=404, detail="Tenant not found")

    # 删除与租户相关的用户
    stmt = delete(User).where(User.tenant_id == tenant.id)
    session.execute(stmt)

    # 删除租户
    session.delete(tenant)
    session.commit()

    return {"detail": "Tenant deleted successfully"}