Dify智能体平台源码二次开发笔记（5） - 多租户的SAAS版实现（2）

前言

用户的查询

controller层

添加路由

service层

用户的添加

controller层

添加路由

service层-添加用户

service层-添加用户和租户关系

验证结果

结果

前言

完成租户添加功能后，下一步需要实现租户下的用户管理。基础功能包括：查询租户用户列表接口，添加用户接口

用户的查询

controller层

class AccountListApi(Resource):"""Resource for getting account list."""@validate_tokendef get(self):"""Get account list."""parser = reqparse.RequestParser()parser.add_argument("tenant_id", type=str, required=True, location="json")parser.add_argument("page", type=int, required=False, default=1, location="args")parser.add_argument("limit", type=int, required=False, default=20, location="args")parser.add_argument("search", type=str, required=False, location="args")args = parser.parse_args()accounts = AccountService.get_accounts_by_tenant(tenant_id=args["tenant_id"],page=args["page"],limit=args["limit"],search=args["search"])return {"result": "success","data": accounts}

添加路由

api.add_resource(AccountListApi, "/accounts")

service层

@staticmethoddef get_accounts_by_tenant(tenant_id: str, page: int = 1, limit: int = 20, search: str = None, status: str = None) -> dict:query = (db.session.query(Account, TenantAccountJoin.role).select_from(Account).join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id).filter(TenantAccountJoin.tenant_id == tenant_id))if search:search = f"%{search}%"query = query.filter(db.or_(Account.name.ilike(search),Account.email.ilike(search)))if status:query = query.filter(Account.status == status)query = query.order_by(Account.name, Account.email)total = query.count()query = query.offset((page - 1) * limit).limit(limit)results = query.all()account_list = [{"id": str(account[0].id),"name": account[0].name,"email": account[0].email,"created_at": account[0].created_at.isoformat(),"role": account[1]} for account in results]return {'items': account_list,'total': total,'page': page,'limit': limit}

用户的添加

controller层

class AccountCreateApi(Resource):"""Resource for creating a new account."""@validate_tokendef post(self):"""Create a new account."""parser = reqparse.RequestParser()parser.add_argument("name", type=str, required=True, location="json")parser.add_argument("email", type=str, required=True, location="json")parser.add_argument("password", type=str, required=True, location="json")parser.add_argument("tenant_id", type=str, required=True, location="json")args = parser.parse_args()account = current_usertenant = TenantService.get_tenant_by_id(args["tenant_id"])if not tenant:return {"result": "fail", "message": "Tenant not found"}, 404try:new_account = AccountService.create_account(email=args["email"],name=args["name"],interface_language="zh-Hans",password=args["password"])TenantService.create_tenant_member(tenant, new_account, role="owner")return {"result": "success","data": {"id": str(new_account.id),"name": new_account.name,"email": new_account.email,"created_at": new_account.created_at.isoformat()}}except Exception as e:return {"result": "error", "message": str(e)}, 400

传入租户id和用户信息，我这里就直接默认语言是中文。

添加路由

api.add_resource(AccountCreateApi, "/accounts/create")

service层-添加用户

@staticmethoddef create_account(email: str,name: str,interface_language: str,password: Optional[str] = None,interface_theme: str = "light",is_setup: Optional[bool] = False,) -> Account:"""create account"""# if not FeatureService.get_system_features().is_allow_register and not is_setup:#     from controllers.console.error import AccountNotFound##     raise AccountNotFound()if dify_config.BILLING_ENABLED and BillingService.is_email_in_freeze(email):raise AccountRegisterError(description=("This email account has been deleted within the past ""30 days and is temporarily unavailable for new account registration"))account = Account()account.email = emailaccount.name = nameif password:# generate password saltsalt = secrets.token_bytes(16)base64_salt = base64.b64encode(salt).decode()# encrypt password with saltpassword_hashed = hash_password(password, salt)base64_password_hashed = base64.b64encode(password_hashed).decode()account.password = base64_password_hashedaccount.password_salt = base64_saltaccount.interface_language = interface_languageaccount.interface_theme = interface_theme# Set timezone based on languageaccount.timezone = language_timezone_mapping.get(interface_language, "UTC")db.session.add(account)db.session.commit()return account

service层-添加用户和租户关系

@staticmethoddef create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin:"""Create tenant member"""if role == TenantAccountRole.OWNER.value:if TenantService.has_roles(tenant, [TenantAccountRole.OWNER]):logging.error(f"Tenant {tenant.id} has already an owner.")raise Exception("Tenant already has an owner.")ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=account.id).first()if ta:ta.role = roleelse:ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)db.session.add(ta)db.session.commit()return ta

这里直接调用已有方法