Dify智能体平台源码二次开发笔记(5) - 多租户的SAAS版实现(2)
目录
前言
用户的查询
controller层
添加路由
service层
用户的添加
controller层
添加路由
service层-添加用户
service层-添加用户和租户关系
验证结果
结果
前言
完成租户添加功能后,下一步需要实现租户下的用户管理。基础功能包括:查询租户用户列表接口,添加用户接口
用户的查询
controller层
class AccountListApi(Resource):
"""Resource for getting account list."""
@validate_token
def 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层
@staticmethod
def 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_token
def 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_user
tenant = TenantService.get_tenant_by_id(args["tenant_id"])
if not tenant:
return {"result": "fail", "message": "Tenant not found"}, 404
try:
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层-添加用户
@staticmethod
def 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 = email
account.name = name
if password:
# generate password salt
salt = secrets.token_bytes(16)
base64_salt = base64.b64encode(salt).decode()
# encrypt password with salt
password_hashed = hash_password(password, salt)
base64_password_hashed = base64.b64encode(password_hashed).decode()
account.password = base64_password_hashed
account.password_salt = base64_salt
account.interface_language = interface_language
account.interface_theme = interface_theme
# Set timezone based on language
account.timezone = language_timezone_mapping.get(interface_language, "UTC")
db.session.add(account)
db.session.commit()
return account
service层-添加用户和租户关系
@staticmethod
def 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 = role
else:
ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
db.session.add(ta)
db.session.commit()
return ta
这里直接调用已有方法
TenantAccountJoin
验证结果
用户表
用户租户关联表
结果
用创建的用户和密码从前端登录进去后,智能体、知识库、插件、模型等都完全隔离了。