Underground-Digital/Workflow-Engine

1import logging2 3from flask import request4from flask_login import current_user5from flask_restful import Resource, fields, inputs, marshal, marshal_with, reqparse6from werkzeug.exceptions import Unauthorized7 8import services9from controllers.common.errors import FilenameNotExistsError10from controllers.console import api11from controllers.console.admin import admin_required12from controllers.console.datasets.error import (13    FileTooLargeError,14    NoFileUploadedError,15    TooManyFilesError,16    UnsupportedFileTypeError,17)18from controllers.console.error import AccountNotLinkTenantError19from controllers.console.wraps import (20    account_initialization_required,21    cloud_edition_billing_resource_check,22    setup_required,23)24from extensions.ext_database import db25from libs.helper import TimestampField26from libs.login import login_required27from models.account import Tenant, TenantStatus28from services.account_service import TenantService29from services.file_service import FileService30from services.workspace_service import WorkspaceService31 32provider_fields = {33    "provider_name": fields.String,34    "provider_type": fields.String,35    "is_valid": fields.Boolean,36    "token_is_set": fields.Boolean,37}38 39tenant_fields = {40    "id": fields.String,41    "name": fields.String,42    "plan": fields.String,43    "status": fields.String,44    "created_at": TimestampField,45    "role": fields.String,46    "in_trial": fields.Boolean,47    "trial_end_reason": fields.String,48    "custom_config": fields.Raw(attribute="custom_config"),49}50 51tenants_fields = {52    "id": fields.String,53    "name": fields.String,54    "plan": fields.String,55    "status": fields.String,56    "created_at": TimestampField,57    "current": fields.Boolean,58}59 60workspace_fields = {"id": fields.String, "name": fields.String, "status": fields.String, "created_at": TimestampField}61 62 63class TenantListApi(Resource):64    @setup_required65    @login_required66    @account_initialization_required67    def get(self):68        tenants = TenantService.get_join_tenants(current_user)69 70        for tenant in tenants:71            if tenant.id == current_user.current_tenant_id:72                tenant.current = True  # Set current=True for current tenant73        return {"workspaces": marshal(tenants, tenants_fields)}, 20074 75 76class WorkspaceListApi(Resource):77    @setup_required78    @admin_required79    def get(self):80        parser = reqparse.RequestParser()81        parser.add_argument("page", type=inputs.int_range(1, 99999), required=False, default=1, location="args")82        parser.add_argument("limit", type=inputs.int_range(1, 100), required=False, default=20, location="args")83        args = parser.parse_args()84 85        tenants = (86            db.session.query(Tenant)87            .order_by(Tenant.created_at.desc())88            .paginate(page=args["page"], per_page=args["limit"])89        )90 91        has_more = False92        if len(tenants.items) == args["limit"]:93            current_page_first_tenant = tenants[-1]94            rest_count = (95                db.session.query(Tenant)96                .filter(97                    Tenant.created_at < current_page_first_tenant.created_at, Tenant.id != current_page_first_tenant.id98                )99                .count()100            )101 102            if rest_count > 0:103                has_more = True104        total = db.session.query(Tenant).count()105        return {106            "data": marshal(tenants.items, workspace_fields),107            "has_more": has_more,108            "limit": args["limit"],109            "page": args["page"],110            "total": total,111        }, 200112 113 114class TenantApi(Resource):115    @setup_required116    @login_required117    @account_initialization_required118    @marshal_with(tenant_fields)119    def get(self):120        if request.path == "/info":121            logging.warning("Deprecated URL /info was used.")122 123        tenant = current_user.current_tenant124 125        if tenant.status == TenantStatus.ARCHIVE:126            tenants = TenantService.get_join_tenants(current_user)127            # if there is any tenant, switch to the first one128            if len(tenants) > 0:129                TenantService.switch_tenant(current_user, tenants[0].id)130                tenant = tenants[0]131            # else, raise Unauthorized132            else:133                raise Unauthorized("workspace is archived")134 135        return WorkspaceService.get_tenant_info(tenant), 200136 137 138class SwitchWorkspaceApi(Resource):139    @setup_required140    @login_required141    @account_initialization_required142    def post(self):143        parser = reqparse.RequestParser()144        parser.add_argument("tenant_id", type=str, required=True, location="json")145        args = parser.parse_args()146 147        # check if tenant_id is valid, 403 if not148        try:149            TenantService.switch_tenant(current_user, args["tenant_id"])150        except Exception:151            raise AccountNotLinkTenantError("Account not link tenant")152 153        new_tenant = db.session.query(Tenant).get(args["tenant_id"])  # Get new tenant154 155        return {"result": "success", "new_tenant": marshal(WorkspaceService.get_tenant_info(new_tenant), tenant_fields)}156 157 158class CustomConfigWorkspaceApi(Resource):159    @setup_required160    @login_required161    @account_initialization_required162    @cloud_edition_billing_resource_check("workspace_custom")163    def post(self):164        parser = reqparse.RequestParser()165        parser.add_argument("remove_webapp_brand", type=bool, location="json")166        parser.add_argument("replace_webapp_logo", type=str, location="json")167        args = parser.parse_args()168 169        tenant = db.session.query(Tenant).filter(Tenant.id == current_user.current_tenant_id).one_or_404()170 171        custom_config_dict = {172            "remove_webapp_brand": args["remove_webapp_brand"],173            "replace_webapp_logo": args["replace_webapp_logo"]174            if args["replace_webapp_logo"] is not None175            else tenant.custom_config_dict.get("replace_webapp_logo"),176        }177 178        tenant.custom_config_dict = custom_config_dict179        db.session.commit()180 181        return {"result": "success", "tenant": marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)}182 183 184class WebappLogoWorkspaceApi(Resource):185    @setup_required186    @login_required187    @account_initialization_required188    @cloud_edition_billing_resource_check("workspace_custom")189    def post(self):190        # get file from request191        file = request.files["file"]192 193        # check file194        if "file" not in request.files:195            raise NoFileUploadedError()196 197        if len(request.files) > 1:198            raise TooManyFilesError()199 200        if not file.filename:201            raise FilenameNotExistsError202 203        extension = file.filename.split(".")[-1]204        if extension.lower() not in {"svg", "png"}:205            raise UnsupportedFileTypeError()206 207        try:208            upload_file = FileService.upload_file(209                filename=file.filename,210                content=file.read(),211                mimetype=file.mimetype,212                user=current_user,213            )214 215        except services.errors.file.FileTooLargeError as file_too_large_error:216            raise FileTooLargeError(file_too_large_error.description)217        except services.errors.file.UnsupportedFileTypeError:218            raise UnsupportedFileTypeError()219 220        return {"id": upload_file.id}, 201221 222 223api.add_resource(TenantListApi, "/workspaces")  # GET for getting all tenants224api.add_resource(WorkspaceListApi, "/all-workspaces")  # GET for getting all tenants225api.add_resource(TenantApi, "/workspaces/current", endpoint="workspaces_current")  # GET for getting current tenant info226api.add_resource(TenantApi, "/info", endpoint="info")  # Deprecated227api.add_resource(SwitchWorkspaceApi, "/workspaces/switch")  # POST for switching tenant228api.add_resource(CustomConfigWorkspaceApi, "/workspaces/custom-config")229api.add_resource(WebappLogoWorkspaceApi, "/workspaces/custom-config/webapp-logo/upload")230