api/models/account.py
Rama main · réplica de lectura del repositorio original.
import enum
import json
from dataclasses import field
from datetime import datetime
from typing import TypedDict
from uuid import uuid4
from warnings import deprecated
import sqlalchemy as sa
from flask_login import UserMixin
from sqlalchemy import DateTime, String, func, select
from sqlalchemy.orm import Mapped, Session, mapped_column
from configs import dify_config
from enums.account import TenantAccountRole
from .base import TypeBase
from .types import EnumText, LongText, StringUUID
class AccountStatus(enum.StrEnum):
PENDING = "pending"
UNINITIALIZED = "uninitialized"
ACTIVE = "active"
BANNED = "banned"
CLOSED = "closed"
class Account(UserMixin, TypeBase):
__tablename__ = "accounts"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="account_pkey"),
sa.Index("account_email_idx", "email"),
sa.Index("account_normalized_email_idx", "normalized_email"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
name: Mapped[str] = mapped_column(String(255))
email: Mapped[str] = mapped_column(String(255))
normalized_email: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
password: Mapped[str | None] = mapped_column(String(255), default=None)
password_salt: Mapped[str | None] = mapped_column(String(255), default=None)
avatar: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
interface_language: Mapped[str | None] = mapped_column(String(255), default=None)
interface_theme: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
timezone: Mapped[str | None] = mapped_column(String(255), default=None)
last_login_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
last_login_ip: Mapped[str | None] = mapped_column(String(255), nullable=True, default=None)
last_active_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False
)
status: Mapped[AccountStatus] = mapped_column(EnumText(AccountStatus, length=16), default=AccountStatus.ACTIVE)
initialized_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False, onupdate=func.current_timestamp()
)
role: TenantAccountRole | None = field(default=None, init=False)
_current_tenant: "Tenant | None" = field(default=None, init=False)
@property
def is_password_set(self):
return self.password is not None
@property
def current_tenant(self):
return self._current_tenant
def set_current_tenant_with_session(self, tenant: "Tenant", *, session: Session) -> None:
"""Set the current tenant and role using the caller-owned session."""
tenant_join_query = select(TenantAccountJoin).where(
TenantAccountJoin.tenant_id == tenant.id, TenantAccountJoin.account_id == self.id
)
tenant_join = session.scalar(tenant_join_query)
tenant_query = select(Tenant).where(Tenant.id == tenant.id)
tenant_reloaded = session.scalars(tenant_query).one()
if tenant_join:
self.role = TenantAccountRole(tenant_join.role)
self._current_tenant = tenant_reloaded
return
self._current_tenant = None
@property
def current_tenant_id(self) -> str | None:
return self._current_tenant.id if self._current_tenant else None
def set_tenant_id_with_session(self, tenant_id: str, *, session: Session) -> None:
"""Set the current tenant by id using the caller-owned session."""
query = select(Tenant, TenantAccountJoin).where(
Tenant.id == tenant_id, TenantAccountJoin.tenant_id == Tenant.id, TenantAccountJoin.account_id == self.id
)
tenant_account_join = session.execute(query).first()
if not tenant_account_join:
return
tenant, join = tenant_account_join
self.role = TenantAccountRole(join.role)
self._current_tenant = tenant
@property
def current_role(self):
return self.role
def get_status(self) -> AccountStatus:
return self.status
# check current_user.current_tenant.current_role in ['admin', 'owner']
@property
def is_admin_or_owner(self):
if dify_config.RBAC_ENABLED:
return True
return TenantAccountRole.is_privileged_role(self.role)
@property
def is_admin(self):
if dify_config.RBAC_ENABLED:
return True
return TenantAccountRole.is_admin_role(self.role)
@property
@deprecated("Use has_edit_permission instead.")
def is_editor(self):
"""Determines if the account has edit permissions in their current tenant (workspace).
This property checks if the current role has editing privileges, which includes:
- `OWNER`
- `ADMIN`
- `EDITOR`
Note: This checks for any role with editing permission, not just the 'EDITOR' role specifically.
"""
return self.has_edit_permission
@property
def has_edit_permission(self):
"""Determines if the account has editing permissions in their current tenant (workspace).
This property checks if the current role has editing privileges, which includes:
- `OWNER`
- `ADMIN`
- `EDITOR`
"""
if dify_config.RBAC_ENABLED:
return True
return TenantAccountRole.is_editing_role(self.role)
@property
def is_dataset_editor(self):
if dify_config.RBAC_ENABLED:
return True
return TenantAccountRole.is_dataset_edit_role(self.role)
@property
def is_dataset_operator(self):
if dify_config.RBAC_ENABLED:
return True
return self.role == TenantAccountRole.DATASET_OPERATOR
class TenantStatus(enum.StrEnum):
NORMAL = "normal"
ARCHIVE = "archive"
class TenantCustomConfigDict(TypedDict, total=False):
remove_webapp_brand: bool
replace_webapp_logo: str | None
class Tenant(TypeBase):
__tablename__ = "tenants"
__table_args__ = (sa.PrimaryKeyConstraint("id", name="tenant_pkey"),)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
name: Mapped[str] = mapped_column(String(255))
encrypt_public_key: Mapped[str | None] = mapped_column(LongText, default=None)
plan: Mapped[str] = mapped_column(String(255), default="basic")
status: Mapped[TenantStatus] = mapped_column(EnumText(TenantStatus, length=255), default=TenantStatus.NORMAL)
custom_config: Mapped[str | None] = mapped_column(LongText, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), init=False, onupdate=func.current_timestamp()
)
def get_accounts(self, *, session: Session) -> list[Account]:
return list(
session.scalars(
select(Account).where(
Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id
)
).all()
)
@property
def custom_config_dict(self) -> TenantCustomConfigDict:
return json.loads(self.custom_config) if self.custom_config else {}
@custom_config_dict.setter
def custom_config_dict(self, value: TenantCustomConfigDict) -> None:
self.custom_config = json.dumps(value)
class TenantAccountJoin(TypeBase):
__tablename__ = "tenant_account_joins"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_account_join_pkey"),
sa.Index("tenant_account_join_account_id_idx", "account_id"),
sa.Index("tenant_account_join_tenant_id_idx", "tenant_id"),
sa.UniqueConstraint("tenant_id", "account_id", name="unique_tenant_account_join"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID)
account_id: Mapped[str] = mapped_column(StringUUID)
current: Mapped[bool] = mapped_column(sa.Boolean, default=False)
role: Mapped[TenantAccountRole] = mapped_column(
EnumText(TenantAccountRole, length=16), default=TenantAccountRole.NORMAL
)
invited_by: Mapped[str | None] = mapped_column(StringUUID, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False, onupdate=func.current_timestamp()
)
last_opened_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
class AccountIntegrate(TypeBase):
__tablename__ = "account_integrates"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="account_integrate_pkey"),
sa.UniqueConstraint("account_id", "provider", name="unique_account_provider"),
sa.UniqueConstraint("provider", "open_id", name="unique_provider_open_id"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
account_id: Mapped[str] = mapped_column(StringUUID)
provider: Mapped[str] = mapped_column(String(16))
open_id: Mapped[str] = mapped_column(String(255))
encrypted_token: Mapped[str] = mapped_column(String(255))
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, server_default=func.current_timestamp(), nullable=False, init=False, onupdate=func.current_timestamp()
)
class InvitationCodeStatus(enum.StrEnum):
UNUSED = "unused"
USED = "used"
class InvitationCode(TypeBase):
__tablename__ = "invitation_codes"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="invitation_code_pkey"),
sa.Index("invitation_codes_batch_idx", "batch"),
sa.Index("invitation_codes_code_idx", "code", "status"),
)
id: Mapped[int] = mapped_column(sa.Integer, init=False)
batch: Mapped[str] = mapped_column(String(255))
code: Mapped[str] = mapped_column(String(32))
status: Mapped[InvitationCodeStatus] = mapped_column(
EnumText(InvitationCodeStatus, length=16),
default=InvitationCodeStatus.UNUSED,
)
used_at: Mapped[datetime | None] = mapped_column(DateTime, default=None)
used_by_tenant_id: Mapped[str | None] = mapped_column(StringUUID, default=None)
used_by_account_id: Mapped[str | None] = mapped_column(StringUUID, default=None)
deprecated_at: Mapped[datetime | None] = mapped_column(DateTime, nullable=True, default=None)
created_at: Mapped[datetime] = mapped_column(
DateTime, server_default=sa.func.current_timestamp(), nullable=False, init=False
)
class TenantPluginInstallPermission(enum.StrEnum):
EVERYONE = "everyone"
ADMINS = "admins"
NOBODY = "noone"
class TenantPluginDebugPermission(enum.StrEnum):
EVERYONE = "everyone"
ADMINS = "admins"
NOBODY = "noone"
class TenantPluginPermission(TypeBase):
__tablename__ = "account_plugin_permissions"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="account_plugin_permission_pkey"),
sa.UniqueConstraint("tenant_id", name="unique_tenant_plugin"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
install_permission: Mapped[TenantPluginInstallPermission] = mapped_column(
EnumText(TenantPluginInstallPermission, length=16),
nullable=False,
default=TenantPluginInstallPermission.ADMINS,
)
debug_permission: Mapped[TenantPluginDebugPermission] = mapped_column(
EnumText(TenantPluginDebugPermission, length=16),
nullable=False,
default=TenantPluginDebugPermission.NOBODY,
)
class TenantPluginAutoUpgradeCategory(enum.StrEnum):
TOOL = "tool"
MODEL = "model"
EXTENSION = "extension"
AGENT_STRATEGY = "agent-strategy"
DATASOURCE = "datasource"
TRIGGER = "trigger"
class TenantPluginAutoUpgradeStrategySetting(enum.StrEnum):
DISABLED = "disabled"
FIX_ONLY = "fix_only"
LATEST = "latest"
class TenantPluginAutoUpgradeMode(enum.StrEnum):
ALL = "all"
PARTIAL = "partial"
EXCLUDE = "exclude"
class TenantPluginAutoUpgradeStrategy(TypeBase):
__tablename__ = "tenant_plugin_auto_upgrade_strategies"
__table_args__ = (
sa.PrimaryKeyConstraint("id", name="tenant_plugin_auto_upgrade_strategy_pkey"),
sa.UniqueConstraint("tenant_id", "category", name="unique_tenant_plugin_auto_upgrade_strategy"),
sa.Index("idx_tenant_plugin_auto_upgrade_strategy_time", "upgrade_time_of_day"),
)
id: Mapped[str] = mapped_column(
StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
)
tenant_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
category: Mapped[TenantPluginAutoUpgradeCategory] = mapped_column(
EnumText(TenantPluginAutoUpgradeCategory, length=32),
nullable=False,
default=TenantPluginAutoUpgradeCategory.TOOL,
)
strategy_setting: Mapped[TenantPluginAutoUpgradeStrategySetting] = mapped_column(
EnumText(TenantPluginAutoUpgradeStrategySetting, length=16),
nullable=False,
default=TenantPluginAutoUpgradeStrategySetting.FIX_ONLY,
)
upgrade_mode: Mapped[TenantPluginAutoUpgradeMode] = mapped_column(
EnumText(TenantPluginAutoUpgradeMode, length=16),
nullable=False,
default=TenantPluginAutoUpgradeMode.EXCLUDE,
)
exclude_plugins: Mapped[list[str]] = mapped_column(sa.JSON, nullable=False, default_factory=list)
include_plugins: Mapped[list[str]] = mapped_column(sa.JSON, nullable=False, default_factory=list)
upgrade_time_of_day: Mapped[int] = mapped_column(sa.Integer, nullable=False, default=0)
created_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False
)
updated_at: Mapped[datetime] = mapped_column(
DateTime, nullable=False, server_default=func.current_timestamp(), init=False, onupdate=func.current_timestamp()
)