Underground-Digital/Workflow-Engine

1import enum2import json3 4from flask_login import UserMixin5 6from extensions.ext_database import db7 8from .types import StringUUID9 10 11class AccountStatus(str, enum.Enum):12    PENDING = "pending"13    UNINITIALIZED = "uninitialized"14    ACTIVE = "active"15    BANNED = "banned"16    CLOSED = "closed"17 18 19class Account(UserMixin, db.Model):20    __tablename__ = "accounts"21    __table_args__ = (db.PrimaryKeyConstraint("id", name="account_pkey"), db.Index("account_email_idx", "email"))22 23    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))24    name = db.Column(db.String(255), nullable=False)25    email = db.Column(db.String(255), nullable=False)26    password = db.Column(db.String(255), nullable=True)27    password_salt = db.Column(db.String(255), nullable=True)28    avatar = db.Column(db.String(255))29    interface_language = db.Column(db.String(255))30    interface_theme = db.Column(db.String(255))31    timezone = db.Column(db.String(255))32    last_login_at = db.Column(db.DateTime)33    last_login_ip = db.Column(db.String(255))34    last_active_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))35    status = db.Column(db.String(16), nullable=False, server_default=db.text("'active'::character varying"))36    initialized_at = db.Column(db.DateTime)37    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))38    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))39 40    @property41    def is_password_set(self):42        return self.password is not None43 44    @property45    def current_tenant(self):46        return self._current_tenant47 48    @current_tenant.setter49    def current_tenant(self, value: "Tenant"):50        tenant = value51        ta = TenantAccountJoin.query.filter_by(tenant_id=tenant.id, account_id=self.id).first()52        if ta:53            tenant.current_role = ta.role54        else:55            tenant = None56        self._current_tenant = tenant57 58    @property59    def current_tenant_id(self):60        return self._current_tenant.id61 62    @current_tenant_id.setter63    def current_tenant_id(self, value: str):64        try:65            tenant_account_join = (66                db.session.query(Tenant, TenantAccountJoin)67                .filter(Tenant.id == value)68                .filter(TenantAccountJoin.tenant_id == Tenant.id)69                .filter(TenantAccountJoin.account_id == self.id)70                .one_or_none()71            )72 73            if tenant_account_join:74                tenant, ta = tenant_account_join75                tenant.current_role = ta.role76            else:77                tenant = None78        except:79            tenant = None80 81        self._current_tenant = tenant82 83    @property84    def current_role(self):85        return self._current_tenant.current_role86 87    def get_status(self) -> AccountStatus:88        status_str = self.status89        return AccountStatus(status_str)90 91    @classmethod92    def get_by_openid(cls, provider: str, open_id: str) -> db.Model:93        account_integrate = (94            db.session.query(AccountIntegrate)95            .filter(AccountIntegrate.provider == provider, AccountIntegrate.open_id == open_id)96            .one_or_none()97        )98        if account_integrate:99            return db.session.query(Account).filter(Account.id == account_integrate.account_id).one_or_none()100        return None101 102    def get_integrates(self) -> list[db.Model]:103        ai = db.Model104        return db.session.query(ai).filter(ai.account_id == self.id).all()105 106    # check current_user.current_tenant.current_role in ['admin', 'owner']107    @property108    def is_admin_or_owner(self):109        return TenantAccountRole.is_privileged_role(self._current_tenant.current_role)110 111    @property112    def is_editor(self):113        return TenantAccountRole.is_editing_role(self._current_tenant.current_role)114 115    @property116    def is_dataset_editor(self):117        return TenantAccountRole.is_dataset_edit_role(self._current_tenant.current_role)118 119    @property120    def is_dataset_operator(self):121        return self._current_tenant.current_role == TenantAccountRole.DATASET_OPERATOR122 123 124class TenantStatus(str, enum.Enum):125    NORMAL = "normal"126    ARCHIVE = "archive"127 128 129class TenantAccountRole(str, enum.Enum):130    OWNER = "owner"131    ADMIN = "admin"132    EDITOR = "editor"133    NORMAL = "normal"134    DATASET_OPERATOR = "dataset_operator"135 136    @staticmethod137    def is_valid_role(role: str) -> bool:138        return role and role in {139            TenantAccountRole.OWNER,140            TenantAccountRole.ADMIN,141            TenantAccountRole.EDITOR,142            TenantAccountRole.NORMAL,143            TenantAccountRole.DATASET_OPERATOR,144        }145 146    @staticmethod147    def is_privileged_role(role: str) -> bool:148        return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN}149 150    @staticmethod151    def is_non_owner_role(role: str) -> bool:152        return role and role in {153            TenantAccountRole.ADMIN,154            TenantAccountRole.EDITOR,155            TenantAccountRole.NORMAL,156            TenantAccountRole.DATASET_OPERATOR,157        }158 159    @staticmethod160    def is_editing_role(role: str) -> bool:161        return role and role in {TenantAccountRole.OWNER, TenantAccountRole.ADMIN, TenantAccountRole.EDITOR}162 163    @staticmethod164    def is_dataset_edit_role(role: str) -> bool:165        return role and role in {166            TenantAccountRole.OWNER,167            TenantAccountRole.ADMIN,168            TenantAccountRole.EDITOR,169            TenantAccountRole.DATASET_OPERATOR,170        }171 172 173class Tenant(db.Model):174    __tablename__ = "tenants"175    __table_args__ = (db.PrimaryKeyConstraint("id", name="tenant_pkey"),)176 177    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))178    name = db.Column(db.String(255), nullable=False)179    encrypt_public_key = db.Column(db.Text)180    plan = db.Column(db.String(255), nullable=False, server_default=db.text("'basic'::character varying"))181    status = db.Column(db.String(255), nullable=False, server_default=db.text("'normal'::character varying"))182    custom_config = db.Column(db.Text)183    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))184    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))185 186    def get_accounts(self) -> list[Account]:187        return (188            db.session.query(Account)189            .filter(Account.id == TenantAccountJoin.account_id, TenantAccountJoin.tenant_id == self.id)190            .all()191        )192 193    @property194    def custom_config_dict(self) -> dict:195        return json.loads(self.custom_config) if self.custom_config else {}196 197    @custom_config_dict.setter198    def custom_config_dict(self, value: dict):199        self.custom_config = json.dumps(value)200 201 202class TenantAccountJoinRole(enum.Enum):203    OWNER = "owner"204    ADMIN = "admin"205    NORMAL = "normal"206    DATASET_OPERATOR = "dataset_operator"207 208 209class TenantAccountJoin(db.Model):210    __tablename__ = "tenant_account_joins"211    __table_args__ = (212        db.PrimaryKeyConstraint("id", name="tenant_account_join_pkey"),213        db.Index("tenant_account_join_account_id_idx", "account_id"),214        db.Index("tenant_account_join_tenant_id_idx", "tenant_id"),215        db.UniqueConstraint("tenant_id", "account_id", name="unique_tenant_account_join"),216    )217 218    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))219    tenant_id = db.Column(StringUUID, nullable=False)220    account_id = db.Column(StringUUID, nullable=False)221    current = db.Column(db.Boolean, nullable=False, server_default=db.text("false"))222    role = db.Column(db.String(16), nullable=False, server_default="normal")223    invited_by = db.Column(StringUUID, nullable=True)224    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))225    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))226 227 228class AccountIntegrate(db.Model):229    __tablename__ = "account_integrates"230    __table_args__ = (231        db.PrimaryKeyConstraint("id", name="account_integrate_pkey"),232        db.UniqueConstraint("account_id", "provider", name="unique_account_provider"),233        db.UniqueConstraint("provider", "open_id", name="unique_provider_open_id"),234    )235 236    id = db.Column(StringUUID, server_default=db.text("uuid_generate_v4()"))237    account_id = db.Column(StringUUID, nullable=False)238    provider = db.Column(db.String(16), nullable=False)239    open_id = db.Column(db.String(255), nullable=False)240    encrypted_token = db.Column(db.String(255), nullable=False)241    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))242    updated_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))243 244 245class InvitationCode(db.Model):246    __tablename__ = "invitation_codes"247    __table_args__ = (248        db.PrimaryKeyConstraint("id", name="invitation_code_pkey"),249        db.Index("invitation_codes_batch_idx", "batch"),250        db.Index("invitation_codes_code_idx", "code", "status"),251    )252 253    id = db.Column(db.Integer, nullable=False)254    batch = db.Column(db.String(255), nullable=False)255    code = db.Column(db.String(32), nullable=False)256    status = db.Column(db.String(16), nullable=False, server_default=db.text("'unused'::character varying"))257    used_at = db.Column(db.DateTime)258    used_by_tenant_id = db.Column(StringUUID)259    used_by_account_id = db.Column(StringUUID)260    deprecated_at = db.Column(db.DateTime)261    created_at = db.Column(db.DateTime, nullable=False, server_default=db.text("CURRENT_TIMESTAMP(0)"))262