hz/ebiz-dify-ai
1
0
Fork 0
mirror of http://112.124.100.131/huang.ze/ebiz-dify-ai.git synced 2025-12-22 01:06:52 +08:00
Code Issues Packages Projects Releases Wiki Activity

chore(api/services): apply ruff reformatting (#7599)

Browse Source 
Co-authored-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
Bowen Liang
2024-08-26 13:43:57 +08:00
committed by GitHub
parent 979422cdc6
commit 17fd773a30
49 changed files with 2630 additions and 2655 deletions
Download Patch File Download Diff File Expand all files Collapse all files

288
api/services/account_service.py
View File

@@ -39,12 +39,7 @@ from tasks.mail_reset_password_task import send_reset_password_mail_task
class AccountService:
reset_password_rate_limiter = RateLimiter(
prefix="reset_password_rate_limit",
max_attempts=5,
time_window=60 * 60
)
reset_password_rate_limiter = RateLimiter(prefix="reset_password_rate_limit", max_attempts=5, time_window=60 * 60)
@staticmethod
def load_user(user_id: str) -> None | Account:
@@ -55,12 +50,15 @@ class AccountService:
if account.status in [AccountStatus.BANNED.value, AccountStatus.CLOSED.value]:
raise Unauthorized("Account is banned or closed.")
current_tenant: TenantAccountJoin = TenantAccountJoin.query.filter_by(account_id=account.id, current=True).first()
current_tenant: TenantAccountJoin = TenantAccountJoin.query.filter_by(
account_id=account.id, current=True
).first()
if current_tenant:
account.current_tenant_id = current_tenant.tenant_id
else:
available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
.order_by(TenantAccountJoin.id.asc()).first()
available_ta = (
TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
)
if not available_ta:
return None
@@ -74,14 +72,13 @@ class AccountService:
return account
@staticmethod
def get_account_jwt_token(account, *, exp: timedelta = timedelta(days=30)):
payload = {
"user_id": account.id,
"exp": datetime.now(timezone.utc).replace(tzinfo=None) + exp,
"iss": dify_config.EDITION,
"sub": 'Console API Passport',
"sub": "Console API Passport",
}
token = PassportService().issue(payload)
@@ -93,10 +90,10 @@ class AccountService:
account = Account.query.filter_by(email=email).first()
if not account:
raise AccountLoginError('Invalid email or password.')
raise AccountLoginError("Invalid email or password.")
if account.status == AccountStatus.BANNED.value or account.status == AccountStatus.CLOSED.value:
raise AccountLoginError('Account is banned or closed.')
raise AccountLoginError("Account is banned or closed.")
if account.status == AccountStatus.PENDING.value:
account.status = AccountStatus.ACTIVE.value
@@ -104,7 +101,7 @@ class AccountService:
db.session.commit()
if account.password is None or not compare_password(password, account.password, account.password_salt):
raise AccountLoginError('Invalid email or password.')
raise AccountLoginError("Invalid email or password.")
return account
@staticmethod
@@ -129,11 +126,9 @@ class AccountService:
return account
@staticmethod
def create_account(email: str,
name: str,
interface_language: str,
password: Optional[str] = None,
interface_theme: str = 'light') -> Account:
def create_account(
email: str, name: str, interface_language: str, password: Optional[str] = None, interface_theme: str = "light"
) -> Account:
"""create account"""
account = Account()
account.email = email
@@ -155,7 +150,7 @@ class AccountService:
account.interface_theme = interface_theme
# Set timezone based on language
account.timezone = language_timezone_mapping.get(interface_language, 'UTC')
account.timezone = language_timezone_mapping.get(interface_language, "UTC")
db.session.add(account)
db.session.commit()
@@ -166,8 +161,9 @@ class AccountService:
"""Link account integrate"""
try:
# Query whether there is an existing binding record for the same provider
account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(account_id=account.id,
provider=provider).first()
account_integrate: Optional[AccountIntegrate] = AccountIntegrate.query.filter_by(
account_id=account.id, provider=provider
).first()
if account_integrate:
# If it exists, update the record
@@ -176,15 +172,16 @@ class AccountService:
account_integrate.updated_at = datetime.now(timezone.utc).replace(tzinfo=None)
else:
# If it does not exist, create a new record
account_integrate = AccountIntegrate(account_id=account.id, provider=provider, open_id=open_id,
encrypted_token="")
account_integrate = AccountIntegrate(
account_id=account.id, provider=provider, open_id=open_id, encrypted_token=""
)
db.session.add(account_integrate)
db.session.commit()
logging.info(f'Account {account.id} linked {provider} account {open_id}.')
logging.info(f"Account {account.id} linked {provider} account {open_id}.")
except Exception as e:
logging.exception(f'Failed to link {provider} account {open_id} to Account {account.id}')
raise LinkAccountIntegrateError('Failed to link account.') from e
logging.exception(f"Failed to link {provider} account {open_id} to Account {account.id}")
raise LinkAccountIntegrateError("Failed to link account.") from e
@staticmethod
def close_account(account: Account) -> None:
@@ -218,7 +215,7 @@ class AccountService:
AccountService.update_last_login(account, ip_address=ip_address)
exp = timedelta(days=30)
token = AccountService.get_account_jwt_token(account, exp=exp)
redis_client.set(_get_login_cache_key(account_id=account.id, token=token), '1', ex=int(exp.total_seconds()))
redis_client.set(_get_login_cache_key(account_id=account.id, token=token), "1", ex=int(exp.total_seconds()))
return token
@staticmethod
@@ -236,22 +233,18 @@ class AccountService:
if cls.reset_password_rate_limiter.is_rate_limited(account.email):
raise RateLimitExceededError(f"Rate limit exceeded for email: {account.email}. Please try again later.")
token = TokenManager.generate_token(account, 'reset_password')
send_reset_password_mail_task.delay(
language=account.interface_language,
to=account.email,
token=token
)
token = TokenManager.generate_token(account, "reset_password")
send_reset_password_mail_task.delay(language=account.interface_language, to=account.email, token=token)
cls.reset_password_rate_limiter.increment_rate_limit(account.email)
return token
@classmethod
def revoke_reset_password_token(cls, token: str):
TokenManager.revoke_token(token, 'reset_password')
TokenManager.revoke_token(token, "reset_password")
@classmethod
def get_reset_password_data(cls, token: str) -> Optional[dict[str, Any]]:
return TokenManager.get_token_data(token, 'reset_password')
return TokenManager.get_token_data(token, "reset_password")
def _get_login_cache_key(*, account_id: str, token: str):
@@ -259,7 +252,6 @@ def _get_login_cache_key(*, account_id: str, token: str):
class TenantService:
@staticmethod
def create_tenant(name: str) -> Tenant:
"""Create tenant"""
@@ -275,31 +267,28 @@ class TenantService:
@staticmethod
def create_owner_tenant_if_not_exist(account: Account):
"""Create owner tenant if not exist"""
available_ta = TenantAccountJoin.query.filter_by(account_id=account.id) \
.order_by(TenantAccountJoin.id.asc()).first()
available_ta = (
TenantAccountJoin.query.filter_by(account_id=account.id).order_by(TenantAccountJoin.id.asc()).first()
)
if available_ta:
return
tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
TenantService.create_tenant_member(tenant, account, role='owner')
TenantService.create_tenant_member(tenant, account, role="owner")
account.current_tenant = tenant
db.session.commit()
tenant_was_created.send(tenant)
@staticmethod
def create_tenant_member(tenant: Tenant, account: Account, role: str = 'normal') -> TenantAccountJoin:
def create_tenant_member(tenant: Tenant, account: Account, role: str = "normal") -> TenantAccountJoin:
"""Create tenant member"""
if role == TenantAccountJoinRole.OWNER.value:
if TenantService.has_roles(tenant, [TenantAccountJoinRole.OWNER]):
logging.error(f'Tenant {tenant.id} has already an owner.')
raise Exception('Tenant already has an owner.')
logging.error(f"Tenant {tenant.id} has already an owner.")
raise Exception("Tenant already has an owner.")
ta = TenantAccountJoin(
tenant_id=tenant.id,
account_id=account.id,
role=role
)
ta = TenantAccountJoin(tenant_id=tenant.id, account_id=account.id, role=role)
db.session.add(ta)
db.session.commit()
return ta
@@ -307,9 +296,12 @@ class TenantService:
@staticmethod
def get_join_tenants(account: Account) -> list[Tenant]:
"""Get account join tenants"""
return db.session.query(Tenant).join(
TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id
).filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL).all()
return (
db.session.query(Tenant)
.join(TenantAccountJoin, Tenant.id == TenantAccountJoin.tenant_id)
.filter(TenantAccountJoin.account_id == account.id, Tenant.status == TenantStatus.NORMAL)
.all()
)
@staticmethod
def get_current_tenant_by_account(account: Account):
@@ -333,16 +325,23 @@ class TenantService:
if tenant_id is None:
raise ValueError("Tenant ID must be provided.")
tenant_account_join = db.session.query(TenantAccountJoin).join(Tenant, TenantAccountJoin.tenant_id == Tenant.id).filter(
TenantAccountJoin.account_id == account.id,
TenantAccountJoin.tenant_id == tenant_id,
Tenant.status == TenantStatus.NORMAL,
).first()
tenant_account_join = (
db.session.query(TenantAccountJoin)
.join(Tenant, TenantAccountJoin.tenant_id == Tenant.id)
.filter(
TenantAccountJoin.account_id == account.id,
TenantAccountJoin.tenant_id == tenant_id,
Tenant.status == TenantStatus.NORMAL,
)
.first()
)
if not tenant_account_join:
raise AccountNotLinkTenantError("Tenant not found or account is not a member of the tenant.")
else:
TenantAccountJoin.query.filter(TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id).update({'current': False})
TenantAccountJoin.query.filter(
TenantAccountJoin.account_id == account.id, TenantAccountJoin.tenant_id != tenant_id
).update({"current": False})
tenant_account_join.current = True
# Set the current tenant for the account
account.current_tenant_id = tenant_account_join.tenant_id
@@ -354,9 +353,7 @@ class TenantService:
query = (
db.session.query(Account, TenantAccountJoin.role)
.select_from(Account)
.join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(TenantAccountJoin.tenant_id == tenant.id)
)
@@ -375,11 +372,9 @@ class TenantService:
query = (
db.session.query(Account, TenantAccountJoin.role)
.select_from(Account)
.join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(TenantAccountJoin.tenant_id == tenant.id)
.filter(TenantAccountJoin.role == 'dataset_operator')
.filter(TenantAccountJoin.role == "dataset_operator")
)
# Initialize an empty list to store the updated accounts
@@ -395,20 +390,25 @@ class TenantService:
def has_roles(tenant: Tenant, roles: list[TenantAccountJoinRole]) -> bool:
"""Check if user has any of the given roles for a tenant"""
if not all(isinstance(role, TenantAccountJoinRole) for role in roles):
raise ValueError('all roles must be TenantAccountJoinRole')
raise ValueError("all roles must be TenantAccountJoinRole")
return db.session.query(TenantAccountJoin).filter(
TenantAccountJoin.tenant_id == tenant.id,
TenantAccountJoin.role.in_([role.value for role in roles])
).first() is not None
return (
db.session.query(TenantAccountJoin)
.filter(
TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.role.in_([role.value for role in roles])
)
.first()
is not None
)
@staticmethod
def get_user_role(account: Account, tenant: Tenant) -> Optional[TenantAccountJoinRole]:
"""Get the role of the current account for a given tenant"""
join = db.session.query(TenantAccountJoin).filter(
TenantAccountJoin.tenant_id == tenant.id,
TenantAccountJoin.account_id == account.id
).first()
join = (
db.session.query(TenantAccountJoin)
.filter(TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == account.id)
.first()
)
return join.role if join else None
@staticmethod
@@ -420,29 +420,26 @@ class TenantService:
def check_member_permission(tenant: Tenant, operator: Account, member: Account, action: str) -> None:
"""Check member permission"""
perms = {
'add': [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
'remove': [TenantAccountRole.OWNER],
'update': [TenantAccountRole.OWNER]
"add": [TenantAccountRole.OWNER, TenantAccountRole.ADMIN],
"remove": [TenantAccountRole.OWNER],
"update": [TenantAccountRole.OWNER],
}
if action not in ['add', 'remove', 'update']:
if action not in ["add", "remove", "update"]:
raise InvalidActionError("Invalid action.")
if member:
if operator.id == member.id:
raise CannotOperateSelfError("Cannot operate self.")
ta_operator = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=operator.id
).first()
ta_operator = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=operator.id).first()
if not ta_operator or ta_operator.role not in perms[action]:
raise NoPermissionError(f'No permission to {action} member.')
raise NoPermissionError(f"No permission to {action} member.")
@staticmethod
def remove_member_from_tenant(tenant: Tenant, account: Account, operator: Account) -> None:
"""Remove member from tenant"""
if operator.id == account.id and TenantService.check_member_permission(tenant, operator, account, 'remove'):
if operator.id == account.id and TenantService.check_member_permission(tenant, operator, account, "remove"):
raise CannotOperateSelfError("Cannot operate self.")
ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
@@ -455,23 +452,17 @@ class TenantService:
@staticmethod
def update_member_role(tenant: Tenant, member: Account, new_role: str, operator: Account) -> None:
"""Update member role"""
TenantService.check_member_permission(tenant, operator, member, 'update')
TenantService.check_member_permission(tenant, operator, member, "update")
target_member_join = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=member.id
).first()
target_member_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=member.id).first()
if target_member_join.role == new_role:
raise RoleAlreadyAssignedError("The provided role is already assigned to the member.")
if new_role == 'owner':
if new_role == "owner":
# Find the current owner and change their role to 'admin'
current_owner_join = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
role='owner'
).first()
current_owner_join.role = 'admin'
current_owner_join = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, role="owner").first()
current_owner_join.role = "admin"
# Update the role of the target member
target_member_join.role = new_role
@@ -480,8 +471,8 @@ class TenantService:
@staticmethod
def dissolve_tenant(tenant: Tenant, operator: Account) -> None:
"""Dissolve tenant"""
if not TenantService.check_member_permission(tenant, operator, operator, 'remove'):
raise NoPermissionError('No permission to dissolve tenant.')
if not TenantService.check_member_permission(tenant, operator, operator, "remove"):
raise NoPermissionError("No permission to dissolve tenant.")
db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id).delete()
db.session.delete(tenant)
db.session.commit()
@@ -494,10 +485,9 @@ class TenantService:
class RegisterService:
@classmethod
def _get_invitation_token_key(cls, token: str) -> str:
return f'member_invite:token:{token}'
return f"member_invite:token:{token}"
@classmethod
def setup(cls, email: str, name: str, password: str, ip_address: str) -> None:
@@ -523,9 +513,7 @@ class RegisterService:
TenantService.create_owner_tenant_if_not_exist(account)
dify_setup = DifySetup(
version=dify_config.CURRENT_VERSION
)
dify_setup = DifySetup(version=dify_config.CURRENT_VERSION)
db.session.add(dify_setup)
db.session.commit()
except Exception as e:
@@ -535,34 +523,35 @@ class RegisterService:
db.session.query(Tenant).delete()
db.session.commit()
logging.exception(f'Setup failed: {e}')
raise ValueError(f'Setup failed: {e}')
logging.exception(f"Setup failed: {e}")
raise ValueError(f"Setup failed: {e}")
@classmethod
def register(cls, email, name,
password: Optional[str] = None,
open_id: Optional[str] = None,
provider: Optional[str] = None,
language: Optional[str] = None,
status: Optional[AccountStatus] = None) -> Account:
def register(
cls,
email,
name,
password: Optional[str] = None,
open_id: Optional[str] = None,
provider: Optional[str] = None,
language: Optional[str] = None,
status: Optional[AccountStatus] = None,
) -> Account:
db.session.begin_nested()
"""Register account"""
try:
account = AccountService.create_account(
email=email,
name=name,
interface_language=language if language else languages[0],
password=password
email=email, name=name, interface_language=language if language else languages[0], password=password
)
account.status = AccountStatus.ACTIVE.value if not status else status.value
account.initialized_at = datetime.now(timezone.utc).replace(tzinfo=None)
if open_id is not None or provider is not None:
AccountService.link_account_integrate(provider, open_id, account)
if dify_config.EDITION != 'SELF_HOSTED':
if dify_config.EDITION != "SELF_HOSTED":
tenant = TenantService.create_tenant(f"{account.name}'s Workspace")
TenantService.create_tenant_member(tenant, account, role='owner')
TenantService.create_tenant_member(tenant, account, role="owner")
account.current_tenant = tenant
tenant_was_created.send(tenant)
@@ -570,30 +559,29 @@ class RegisterService:
db.session.commit()
except Exception as e:
db.session.rollback()
logging.error(f'Register failed: {e}')
raise AccountRegisterError(f'Registration failed: {e}') from e
logging.error(f"Register failed: {e}")
raise AccountRegisterError(f"Registration failed: {e}") from e
return account
@classmethod
def invite_new_member(cls, tenant: Tenant, email: str, language: str, role: str = 'normal', inviter: Account = None) -> str:
def invite_new_member(
cls, tenant: Tenant, email: str, language: str, role: str = "normal", inviter: Account = None
) -> str:
"""Invite new member"""
account = Account.query.filter_by(email=email).first()
if not account:
TenantService.check_member_permission(tenant, inviter, None, 'add')
name = email.split('@')[0]
TenantService.check_member_permission(tenant, inviter, None, "add")
name = email.split("@")[0]
account = cls.register(email=email, name=name, language=language, status=AccountStatus.PENDING)
# Create new tenant member for invited tenant
TenantService.create_tenant_member(tenant, account, role)
TenantService.switch_tenant(account, tenant.id)
else:
TenantService.check_member_permission(tenant, inviter, account, 'add')
ta = TenantAccountJoin.query.filter_by(
tenant_id=tenant.id,
account_id=account.id
).first()
TenantService.check_member_permission(tenant, inviter, account, "add")
ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=account.id).first()
if not ta:
TenantService.create_tenant_member(tenant, account, role)
@@ -609,7 +597,7 @@ class RegisterService:
language=account.interface_language,
to=email,
token=token,
inviter_name=inviter.name if inviter else 'Dify',
inviter_name=inviter.name if inviter else "Dify",
workspace_name=tenant.name,
)
@@ -619,23 +607,19 @@ class RegisterService:
def generate_invite_token(cls, tenant: Tenant, account: Account) -> str:
token = str(uuid.uuid4())
invitation_data = {
'account_id': account.id,
'email': account.email,
'workspace_id': tenant.id,
"account_id": account.id,
"email": account.email,
"workspace_id": tenant.id,
}
expiryHours = dify_config.INVITE_EXPIRY_HOURS
redis_client.setex(
cls._get_invitation_token_key(token),
expiryHours * 60 * 60,
json.dumps(invitation_data)
)
redis_client.setex(cls._get_invitation_token_key(token), expiryHours * 60 * 60, json.dumps(invitation_data))
return token
@classmethod
def revoke_token(cls, workspace_id: str, email: str, token: str):
if workspace_id and email:
email_hash = sha256(email.encode()).hexdigest()
cache_key = 'member_invite_token:{}, {}:{}'.format(workspace_id, email_hash, token)
cache_key = "member_invite_token:{}, {}:{}".format(workspace_id, email_hash, token)
redis_client.delete(cache_key)
else:
redis_client.delete(cls._get_invitation_token_key(token))
@@ -646,17 +630,21 @@ class RegisterService:
if not invitation_data:
return None
tenant = db.session.query(Tenant).filter(
Tenant.id == invitation_data['workspace_id'],
Tenant.status == 'normal'
).first()
tenant = (
db.session.query(Tenant)
.filter(Tenant.id == invitation_data["workspace_id"], Tenant.status == "normal")
.first()
)
if not tenant:
return None
tenant_account = db.session.query(Account, TenantAccountJoin.role).join(
TenantAccountJoin, Account.id == TenantAccountJoin.account_id
).filter(Account.email == invitation_data['email'], TenantAccountJoin.tenant_id == tenant.id).first()
tenant_account = (
db.session.query(Account, TenantAccountJoin.role)
.join(TenantAccountJoin, Account.id == TenantAccountJoin.account_id)
.filter(Account.email == invitation_data["email"], TenantAccountJoin.tenant_id == tenant.id)
.first()
)
if not tenant_account:
return None
@@ -665,29 +653,29 @@ class RegisterService:
if not account:
return None
if invitation_data['account_id'] != str(account.id):
if invitation_data["account_id"] != str(account.id):
return None
return {
'account': account,
'data': invitation_data,
'tenant': tenant,
"account": account,
"data": invitation_data,
"tenant": tenant,
}
@classmethod
def _get_invitation_by_token(cls, token: str, workspace_id: str, email: str) -> Optional[dict[str, str]]:
if workspace_id is not None and email is not None:
email_hash = sha256(email.encode()).hexdigest()
cache_key = f'member_invite_token:{workspace_id}, {email_hash}:{token}'
cache_key = f"member_invite_token:{workspace_id}, {email_hash}:{token}"
account_id = redis_client.get(cache_key)
if not account_id:
return None
return {
'account_id': account_id.decode('utf-8'),
'email': email,
'workspace_id': workspace_id,
"account_id": account_id.decode("utf-8"),
"email": email,
"workspace_id": workspace_id,
}
else:
data = redis_client.get(cls._get_invitation_token_key(token))