Skip to content

Commit

Permalink
reformat
Browse files Browse the repository at this point in the history
  • Loading branch information
GareArc committed Dec 23, 2024
1 parent d9687ee commit d9bba39
Show file tree
Hide file tree
Showing 9 changed files with 100 additions and 95 deletions.
16 changes: 12 additions & 4 deletions api/configs/feature/__init__.py
Original file line number Diff line number Diff line change
@@ -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):
"""
Expand Down
26 changes: 12 additions & 14 deletions api/controllers/console/workspace/account.py
Original file line number Diff line number Diff line change
@@ -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):
Expand Down Expand Up @@ -243,7 +243,6 @@ def get(self):


class AccountDeleteVerifyApi(Resource):

@setup_required
@login_required
@account_initialization_required
Expand All @@ -260,7 +259,6 @@ def get(self):


class AccountDeleteApi(Resource):

@setup_required
@login_required
@account_initialization_required
Expand Down
5 changes: 3 additions & 2 deletions api/libs/helper.py
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand Down
8 changes: 2 additions & 6 deletions api/models/account.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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)
Expand Down
9 changes: 5 additions & 4 deletions api/services/account_deletion_log_service.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,3 @@

import json
from datetime import timedelta

Expand All @@ -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
Expand Down
60 changes: 34 additions & 26 deletions api/services/account_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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):
Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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()

Expand Down Expand Up @@ -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()

Expand Down
3 changes: 2 additions & 1 deletion api/services/billing_service.py
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
import os

import requests

from extensions.ext_database import db
from models.account import TenantAccountJoin, TenantAccountRole

Expand Down Expand Up @@ -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)
40 changes: 20 additions & 20 deletions api/tasks/delete_account_task.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -37,25 +37,21 @@ 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:
_handle_owner_tenant_deletion(ta)
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.")
Expand All @@ -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}.")
Expand All @@ -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)

Expand All @@ -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)
Loading

0 comments on commit d9bba39

Please sign in to comment.