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

Signed-off-by: -LAN- <laipz8200@outlook.com>

api/models/account.py

@@ -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"),)