From c46037cae3edeecdc196c1c6ac2e5c9daa169f94 Mon Sep 17 00:00:00 2001 From: GareArc Date: Mon, 23 Dec 2024 02:19:22 -0500 Subject: [PATCH] reformat --- api/configs/feature/__init__.py | 16 ++++-- api/controllers/console/workspace/account.py | 26 ++++----- api/libs/helper.py | 5 +- api/models/account.py | 8 +-- api/services/account_deletion_log_service.py | 9 +-- api/services/account_service.py | 60 +++++++++++--------- api/services/billing_service.py | 3 +- api/tasks/delete_account_task.py | 40 ++++++------- api/tasks/mail_account_deletion_task.py | 28 ++++----- 9 files changed, 100 insertions(+), 95 deletions(-) diff --git a/api/configs/feature/__init__.py b/api/configs/feature/__init__.py index 0730c8a0b569a9..460e5ce3383f85 100644 --- a/api/configs/feature/__init__.py +++ b/api/configs/feature/__init__.py @@ -1,11 +1,19 @@ from typing import Annotated, Literal, Optional -from configs.feature.hosted_service import HostedServiceConfig -from pydantic import (AliasChoices, Field, HttpUrl, NegativeInt, - NonNegativeInt, PositiveFloat, PositiveInt, - computed_field) +from pydantic import ( + AliasChoices, + Field, + HttpUrl, + NegativeInt, + NonNegativeInt, + PositiveFloat, + PositiveInt, + computed_field, +) from pydantic_settings import BaseSettings +from configs.feature.hosted_service import HostedServiceConfig + class SecurityConfig(BaseSettings): """ diff --git a/api/controllers/console/workspace/account.py b/api/controllers/console/workspace/account.py index 412b8a8b212f54..6e7f7e494b71da 100644 --- a/api/controllers/console/workspace/account.py +++ b/api/controllers/console/workspace/account.py @@ -1,27 +1,27 @@ import datetime import pytz +from flask import request +from flask_login import current_user +from flask_restful import Resource, fields, marshal_with, reqparse + from configs import dify_config from constants.languages import supported_language from controllers.console import api -from controllers.console.workspace.error import (AccountAlreadyInitedError, - CurrentPasswordIncorrectError, - InvalidInvitationCodeError, - RepeatPasswordNotMatchError) -from controllers.console.wraps import (account_initialization_required, - enterprise_license_required, - setup_required) +from controllers.console.workspace.error import ( + AccountAlreadyInitedError, + CurrentPasswordIncorrectError, + InvalidInvitationCodeError, + RepeatPasswordNotMatchError, +) +from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required from extensions.ext_database import db from fields.member_fields import account_fields -from flask import request -from flask_login import current_user -from flask_restful import Resource, fields, marshal_with, reqparse from libs.helper import TimestampField, timezone from libs.login import login_required from models import AccountIntegrate, InvitationCode from services.account_service import AccountService -from services.errors.account import \ - CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError +from services.errors.account import CurrentPasswordIncorrectError as ServiceCurrentPasswordIncorrectError class AccountInitApi(Resource): @@ -243,7 +243,6 @@ def get(self): class AccountDeleteVerifyApi(Resource): - @setup_required @login_required @account_initialization_required @@ -260,7 +259,6 @@ def get(self): class AccountDeleteApi(Resource): - @setup_required @login_required @account_initialization_required diff --git a/api/libs/helper.py b/api/libs/helper.py index 38244b4cc18247..f43e0727e0e30c 100644 --- a/api/libs/helper.py +++ b/api/libs/helper.py @@ -13,12 +13,13 @@ from typing import Any, Optional, Union, cast from zoneinfo import available_timezones +from flask import Response, stream_with_context +from flask_restful import fields + from configs import dify_config from core.app.features.rate_limiting.rate_limit import RateLimitGenerator from core.file import helpers as file_helpers from extensions.ext_redis import redis_client -from flask import Response, stream_with_context -from flask_restful import fields from models.account import Account diff --git a/api/models/account.py b/api/models/account.py index b6843e79067ba9..a1ef04bca586e7 100644 --- a/api/models/account.py +++ b/api/models/account.py @@ -271,9 +271,7 @@ class InvitationCode(db.Model): class AccountDeletionLog(db.Model): __tablename__ = "account_deletion_logs" - __table_args__ = ( - db.PrimaryKeyConstraint("id", name="account_deletion_log_pkey"), - ) + __table_args__ = (db.PrimaryKeyConstraint("id", name="account_deletion_log_pkey"),) id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) email = db.Column(db.String(255), nullable=False) @@ -287,9 +285,7 @@ class AccountDeletionLog(db.Model): class AccountDeletionLogDetail(db.Model): __tablename__ = "account_deletion_log_details" - __table_args__ = ( - db.PrimaryKeyConstraint("id", name="account_deletion_log_detail_pkey"), - ) + __table_args__ = (db.PrimaryKeyConstraint("id", name="account_deletion_log_detail_pkey"),) id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()")) account_id = db.Column(StringUUID, nullable=False) diff --git a/api/services/account_deletion_log_service.py b/api/services/account_deletion_log_service.py index 4a44e3d90fd7f6..1a053bc07739f4 100644 --- a/api/services/account_deletion_log_service.py +++ b/api/services/account_deletion_log_service.py @@ -1,4 +1,3 @@ - import json from datetime import timedelta @@ -22,10 +21,12 @@ def create_account_deletion_log(account: Account, reason): @staticmethod def email_in_freeze(email): - log = db.session.query(AccountDeletionLog) \ - .filter(AccountDeletionLog.email == email) \ - .order_by(AccountDeletionLog.created_at.desc()) \ + log = ( + db.session.query(AccountDeletionLog) + .filter(AccountDeletionLog.email == email) + .order_by(AccountDeletionLog.created_at.desc()) .first() + ) if not log: return False diff --git a/api/services/account_service.py b/api/services/account_service.py index 38409ed785960a..d044b9f4e0b389 100644 --- a/api/services/account_service.py +++ b/api/services/account_service.py @@ -8,6 +8,10 @@ from hashlib import sha256 from typing import Any, Optional +from pydantic import BaseModel +from sqlalchemy import func +from werkzeug.exceptions import Unauthorized + from configs import dify_config from constants.languages import language_timezone_mapping, languages from events.tenant_event import tenant_was_created @@ -17,34 +21,41 @@ from libs.passport import PassportService from libs.password import compare_password, hash_password, valid_password from libs.rsa import generate_key_pair -from models.account import (Account, AccountIntegrate, AccountStatus, Tenant, - TenantAccountJoin, TenantAccountJoinRole, - TenantAccountRole, TenantStatus) +from models.account import ( + Account, + AccountIntegrate, + AccountStatus, + Tenant, + TenantAccountJoin, + TenantAccountJoinRole, + TenantAccountRole, + TenantStatus, +) from models.model import DifySetup -from pydantic import BaseModel from services.account_deletion_log_service import AccountDeletionLogService -from services.errors.account import (AccountAlreadyInTenantError, - AccountLoginError, AccountNotFoundError, - AccountNotLinkTenantError, - AccountPasswordError, - AccountRegisterError, - CannotOperateSelfError, - CurrentPasswordIncorrectError, - InvalidActionError, - LinkAccountIntegrateError, - MemberNotInTenantError, NoPermissionError, - RoleAlreadyAssignedError, - TenantNotFoundError) +from services.errors.account import ( + AccountAlreadyInTenantError, + AccountLoginError, + AccountNotFoundError, + AccountNotLinkTenantError, + AccountPasswordError, + AccountRegisterError, + CannotOperateSelfError, + CurrentPasswordIncorrectError, + InvalidActionError, + LinkAccountIntegrateError, + MemberNotInTenantError, + NoPermissionError, + RoleAlreadyAssignedError, + TenantNotFoundError, +) from services.errors.workspace import WorkSpaceNotAllowedCreateError from services.feature_service import FeatureService -from sqlalchemy import func from tasks.delete_account_task import delete_account_task -from tasks.mail_account_deletion_task import \ - send_account_deletion_verification_code +from tasks.mail_account_deletion_task import send_account_deletion_verification_code from tasks.mail_email_code_login import send_email_code_login_mail_task from tasks.mail_invite_member_task import send_invite_member_mail_task from tasks.mail_reset_password_task import send_reset_password_mail_task -from werkzeug.exceptions import Unauthorized class TokenPair(BaseModel): @@ -250,8 +261,7 @@ def generate_account_deletion_verification_code(account: Account) -> tuple[str, @classmethod def send_account_deletion_verification_email(cls, account: Account, code: str): if cls.email_code_account_deletion_rate_limiter.is_rate_limited(email): - from controllers.console.auth.error import \ - EmailCodeAccountDeletionRateLimitExceededError + from controllers.console.auth.error import EmailCodeAccountDeletionRateLimitExceededError raise EmailCodeAccountDeletionRateLimitExceededError() @@ -385,8 +395,7 @@ def send_reset_password_email( account_email = account.email if account else email if cls.reset_password_rate_limiter.is_rate_limited(account_email): - from controllers.console.auth.error import \ - PasswordResetRateLimitExceededError + from controllers.console.auth.error import PasswordResetRateLimitExceededError raise PasswordResetRateLimitExceededError() @@ -420,8 +429,7 @@ def send_email_code_login_email( cls, account: Optional[Account] = None, email: Optional[str] = None, language: Optional[str] = "en-US" ): if cls.email_code_login_rate_limiter.is_rate_limited(email): - from controllers.console.auth.error import \ - EmailCodeLoginRateLimitExceededError + from controllers.console.auth.error import EmailCodeLoginRateLimitExceededError raise EmailCodeLoginRateLimitExceededError() diff --git a/api/services/billing_service.py b/api/services/billing_service.py index f8d9bac7cad984..2167c76c135262 100644 --- a/api/services/billing_service.py +++ b/api/services/billing_service.py @@ -1,6 +1,7 @@ import os import requests + from extensions.ext_database import db from models.account import TenantAccountJoin, TenantAccountRole @@ -61,6 +62,6 @@ def is_tenant_owner_or_admin(current_user): @classmethod def unsubscripbe_tenant_customer(cls, tenant_id: str): - """ Delete related customer in billing service. Used when tenant is deleted.""" + """Delete related customer in billing service. Used when tenant is deleted.""" params = {"tenant_id": tenant_id} return cls._send_request("DELETE", "/subscription", params=params) diff --git a/api/tasks/delete_account_task.py b/api/tasks/delete_account_task.py index 2ba20f5a3dae12..0db11c7caf0a1a 100644 --- a/api/tasks/delete_account_task.py +++ b/api/tasks/delete_account_task.py @@ -4,10 +4,10 @@ import click from celery import shared_task + from extensions.ext_database import db from libs.helper import serialize_sqlalchemy -from models.account import (Account, AccountDeletionLogDetail, - TenantAccountJoin, TenantAccountJoinRole) +from models.account import Account, AccountDeletionLogDetail, TenantAccountJoin, TenantAccountJoinRole from services.account_deletion_log_service import AccountDeletionLogService from services.billing_service import BillingService from tasks.mail_account_deletion_task import send_deletion_success_task @@ -37,15 +37,13 @@ def delete_account_task(account_id, reason: str): ) except Exception as e: db.session.rollback() - logging.error(f"Failed to delete account {account.email}: {e}.") + logging.exception(f"Failed to delete account {account.email}.") raise def _process_account_deletion(account, reason): # Fetch all tenant-account associations - tenant_account_joins = db.session.query(TenantAccountJoin).filter( - TenantAccountJoin.account_id == account.id - ).all() + tenant_account_joins = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.account_id == account.id).all() for ta in tenant_account_joins: if ta.role == TenantAccountJoinRole.OWNER.value: @@ -53,9 +51,7 @@ def _process_account_deletion(account, reason): else: _remove_account_from_tenant(ta, account.email) - account_deletion_log = AccountDeletionLogService.create_account_deletion_log( - account, reason - ) + account_deletion_log = AccountDeletionLogService.create_account_deletion_log(account, reason) db.session.add(account_deletion_log) db.session.delete(account) logger.info(f"Account {account.email} successfully deleted.") @@ -66,9 +62,7 @@ def _handle_owner_tenant_deletion(ta: TenantAccountJoin): tenant_id = ta.tenant_id # Dismiss all tenant members - members_to_dismiss = db.session.query(TenantAccountJoin).filter( - TenantAccountJoin.tenant_id == tenant_id - ).all() + members_to_dismiss = db.session.query(TenantAccountJoin).filter(TenantAccountJoin.tenant_id == tenant_id).all() for member in members_to_dismiss: db.session.delete(member) logger.info(f"Dismissed {len(members_to_dismiss)} members from tenant {tenant_id}.") @@ -78,17 +72,20 @@ def _handle_owner_tenant_deletion(ta: TenantAccountJoin): BillingService.unsubscripbe_tenant_customer(tenant_id) logger.info(f"Subscription for tenant {tenant_id} deleted successfully.") except Exception as e: - logger.error(f"Failed to delete subscription for tenant {tenant_id}: {e}.") + logger.exception(f"Failed to delete subscription for tenant {tenant_id}.") raise # create account deletion log detail account_deletion_log_detail = AccountDeletionLogDetail() account_deletion_log_detail.account_id = ta.account_id account_deletion_log_detail.tenant_id = tenant_id - account_deletion_log_detail.snapshot = json.dumps({ - "tenant_account_join_info": serialize_sqlalchemy(ta), - "dismissed_members": [serialize_sqlalchemy(member) for member in members_to_dismiss] - }, separators=(",", ":")) + account_deletion_log_detail.snapshot = json.dumps( + { + "tenant_account_join_info": serialize_sqlalchemy(ta), + "dismissed_members": [serialize_sqlalchemy(member) for member in members_to_dismiss], + }, + separators=(",", ":"), + ) account_deletion_log_detail.role = ta.role db.session.add(account_deletion_log_detail) @@ -103,8 +100,11 @@ def _remove_account_from_tenant(ta, email): account_deletion_log_detail = AccountDeletionLogDetail() account_deletion_log_detail.account_id = ta.account_id account_deletion_log_detail.tenant_id = tenant_id - account_deletion_log_detail.snapshot = json.dumps({ - "tenant_account_join_info": serialize_sqlalchemy(ta), - }, separators=(",", ":")) + account_deletion_log_detail.snapshot = json.dumps( + { + "tenant_account_join_info": serialize_sqlalchemy(ta), + }, + separators=(",", ":"), + ) account_deletion_log_detail.role = ta.role db.session.add(account_deletion_log_detail) diff --git a/api/tasks/mail_account_deletion_task.py b/api/tasks/mail_account_deletion_task.py index 6e65733cf45459..d38df3613bd450 100644 --- a/api/tasks/mail_account_deletion_task.py +++ b/api/tasks/mail_account_deletion_task.py @@ -3,9 +3,10 @@ import click from celery import shared_task -from extensions.ext_mail import mail from flask import render_template +from extensions.ext_mail import mail + @shared_task(queue="mail") def send_deletion_success_task(language, to): @@ -17,9 +18,7 @@ def send_deletion_success_task(language, to): if not mail.is_inited(): return - logging.info( - click.style(f"Start send account deletion success email to {to}", fg="green") - ) + logging.info(click.style(f"Start send account deletion success email to {to}", fg="green")) start_at = time.perf_counter() try: @@ -59,31 +58,24 @@ def send_account_deletion_verification_code(language, to, code): if not mail.is_inited(): return - logging.info( - click.style(f"Start send account deletion verification code email to {to}", fg="green") - ) + logging.info(click.style(f"Start send account deletion verification code email to {to}", fg="green")) start_at = time.perf_counter() try: if language == "zh-Hans": - html_content = render_template( - "delete_account_code_email_template_zh-CN.html", - to=to, - code=code - ) + html_content = render_template("delete_account_code_email_template_zh-CN.html", to=to, code=code) mail.send(to=to, subject="Dify 的删除账户验证码", html=html_content) else: - html_content = render_template( - "delete_account_code_email_en-US.html", - to=to, - code=code - ) + html_content = render_template("delete_account_code_email_en-US.html", to=to, code=code) mail.send(to=to, subject="Delete Your Dify Account", html=html_content) end_at = time.perf_counter() logging.info( click.style( - "Send account deletion verification code email to {} succeeded: latency: {}".format(to, end_at - start_at), fg="green" + "Send account deletion verification code email to {} succeeded: latency: {}".format( + to, end_at - start_at + ), + fg="green", ) ) except Exception: