jyapayne/dify
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

refactor: Use typed SQLAlchemy base model and fix type errors (#19980)

Browse source 
Signed-off-by: -LAN- <laipz8200@outlook.com>
This commit is contained in:
-LAN- 2025-05-21 15:38:03 +08:00 • committed by GitHub
commit 3196dc2d61
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
20 changed files with 272 additions and 174 deletions
Download patch file Download diff file Expand all files Collapse all files

157
api/models/account.py
View file

@ -1,10 +1,10 @@
import enum
import json
from typing import cast
from typing import Optional, cast
from flask_login import UserMixin # type: ignore
from sqlalchemy import func
from sqlalchemy.orm import Mapped, mapped_column
from sqlalchemy.orm import Mapped, mapped_column, reconstructor
from models.base import Base
@ -12,6 +12,66 @@ from .engine import db
from .types import StringUUID
class TenantAccountRole(enum.StrEnum):
OWNER = "owner"
ADMIN = "admin"
EDITOR = "editor"
NORMAL = "normal"
DATASET_OPERATOR = "dataset_operator"
@staticmethod
def is_valid_role(role: str) -> bool:
if not role:
return False
return role in {
TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod
def is_privileged_role(role: Optional["TenantAccountRole"]) -> bool:
if not role:
return False
return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}
@staticmethod
def is_admin_role(role: Optional["TenantAccountRole"]) -> bool:
if not role:
return False
return role == TenantAccountRole.ADMIN
@staticmethod
def is_non_owner_role(role: Optional["TenantAccountRole"]) -> bool:
if not role:
return False
return role in {
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod
def is_editing_role(role: Optional["TenantAccountRole"]) -> bool:
if not role:
return False
return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}
@staticmethod
def is_dataset_edit_role(role: Optional["TenantAccountRole"]) -> bool:
if not role:
return False
return role in {
TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.DATASET_OPERATOR,
}
class AccountStatus(enum.StrEnum):
PENDING = "pending"
UNINITIALIZED = "uninitialized"
@ -41,24 +101,27 @@ class Account(UserMixin, Base):
created_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
updated_at = db.Column(db.DateTime, nullable=False, server_default=func.current_timestamp())
@reconstructor
def init_on_load(self):
self.role: Optional[TenantAccountRole] = None
self._current_tenant: Optional[Tenant] = None
@property
def is_password_set(self):
return self.password is not None
@property
def current_tenant(self):
return self._current_tenant # type: ignore
return self._current_tenant
@current_tenant.setter
def current_tenant(self, value: "Tenant"):
tenant = value
def current_tenant(self, tenant: "Tenant"):
ta = db.session.query(TenantAccountJoin).filter_by(tenant_id=tenant.id, account_id=self.id).first()
if ta:
tenant.current_role = ta.role
else:
tenant = None # type: ignore
self._current_tenant = tenant
self.role = TenantAccountRole(ta.role)
self._current_tenant = tenant
return
self._current_tenant = None
@property
def current_tenant_id(self) -> str | None:
@ -80,12 +143,12 @@ class Account(UserMixin, Base):
return
tenant, join = tenant_account_join
tenant.current_role = join.role
self.role = join.role
self._current_tenant = tenant
@property
def current_role(self):
return self._current_tenant.current_role
return self.role
def get_status(self) -> AccountStatus:
status_str = self.status
@ -105,23 +168,23 @@ class Account(UserMixin, Base):
# check current_user.current_tenant.current_role in ['admin', 'owner']
@property
def is_admin_or_owner(self):
return TenantAccountRole.is_privileged_role(self._current_tenant.current_role)
return TenantAccountRole.is_privileged_role(self.role)
@property
def is_admin(self):
return TenantAccountRole.is_admin_role(self._current_tenant.current_role)
return TenantAccountRole.is_admin_role(self.role)
@property
def is_editor(self):
return TenantAccountRole.is_editing_role(self._current_tenant.current_role)
return TenantAccountRole.is_editing_role(self.role)
@property
def is_dataset_editor(self):
return TenantAccountRole.is_dataset_edit_role(self._current_tenant.current_role)
return TenantAccountRole.is_dataset_edit_role(self.role)
@property
def is_dataset_operator(self):
return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR
return self.role == TenantAccountRole.DATASET_OPERATOR
class TenantStatus(enum.StrEnum):
@ -129,66 +192,6 @@ class TenantStatus(enum.StrEnum):
ARCHIVE = "archive"
class TenantAccountRole(enum.StrEnum):
OWNER = "owner"
ADMIN = "admin"
EDITOR = "editor"
NORMAL = "normal"
DATASET_OPERATOR = "dataset_operator"
@staticmethod
def is_valid_role(role: str) -> bool:
if not role:
return False
return role in {
TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod
def is_privileged_role(role: str) -> bool:
if not role:
return False
return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}
@staticmethod
def is_admin_role(role: str) -> bool:
if not role:
return False
return role == TenantAccountRole.ADMIN
@staticmethod
def is_non_owner_role(role: str) -> bool:
if not role:
return False
return role in {
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.NORMAL,
TenantAccountRole.DATASET_OPERATOR,
}
@staticmethod
def is_editing_role(role: str) -> bool:
if not role:
return False
return role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}
@staticmethod
def is_dataset_edit_role(role: str) -> bool:
if not role:
return False
return role in {
TenantAccountRole.OWNER,
TenantAccountRole.ADMIN,
TenantAccountRole.EDITOR,
TenantAccountRole.DATASET_OPERATOR,
}
class Tenant(Base):
__tablename__ = "tenants"
__table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),)

6
api/models/base.py
View file

@ -1,5 +1,7 @@
from sqlalchemy.orm import declarative_base
from sqlalchemy.orm import DeclarativeBase
from models.engine import metadata
Base = declarative_base(metadata=metadata)
class Base(DeclarativeBase):
metadata = metadata

4
api/models/tools.py
View file

@ -172,10 +172,6 @@ class WorkflowToolProvider(Base):
db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)")
)
@property
def schema_type(self) -> ApiProviderSchemaType:
return ApiProviderSchemaType.value_of(self.schema_type_str)
@property
def user(self) -> Account | None:
return db.session.query(Account).filter(Account.id == self.user_id).first()

4
api/models/workflow.py
View file

@ -3,7 +3,7 @@ import logging
from collections.abc import Mapping, Sequence
from datetime import UTC, datetime
from enum import Enum, StrEnum
from typing import TYPE_CHECKING, Any, Optional, Self, Union
from typing import TYPE_CHECKING, Any, Optional, Union
from uuid import uuid4
from core.variables import utils as variable_utils
@ -150,7 +150,7 @@ class Workflow(Base):
conversation_variables: Sequence[Variable],
marked_name: str = "",
marked_comment: str = "",
) -> Self:
) -> "Workflow":
workflow = Workflow()
workflow.id = str(uuid4())
workflow.tenant_id = tenant_id