dify/api/controllers/console/workspace/workspace.py

import logging

from flask import request
from flask_login import current_user
from flask_restful import Resource, fields, inputs, marshal, marshal_with, reqparse
from werkzeug.exceptions import Unauthorized

import services
from controllers.common.errors import FilenameNotExistsError
from controllers.console import api
from controllers.console.admin import admin_required
from controllers.console.datasets.error import (
    FileTooLargeError,
    NoFileUploadedError,
    TooManyFilesError,
    UnsupportedFileTypeError,
)
from controllers.console.error import AccountNotLinkTenantError
from controllers.console.wraps import (
    account_initialization_required,
    cloud_edition_billing_resource_check,
    setup_required,
)
from extensions.ext_database import db
from libs.helper import TimestampField
from libs.login import login_required
from models.account import Tenant, TenantStatus
from services.account_service import TenantService
from services.file_service import FileService
from services.workspace_service import WorkspaceService

provider_fields = {
    "provider_name": fields.String,
    "provider_type": fields.String,
    "is_valid": fields.Boolean,
    "token_is_set": fields.Boolean,
}

tenant_fields = {
    "id": fields.String,
    "name": fields.String,
    "plan": fields.String,
    "status": fields.String,
    "created_at": TimestampField,
    "role": fields.String,
    "in_trial": fields.Boolean,
    "trial_end_reason": fields.String,
    "custom_config": fields.Raw(attribute="custom_config"),
}

tenants_fields = {
    "id": fields.String,
    "name": fields.String,
    "plan": fields.String,
    "status": fields.String,
    "created_at": TimestampField,
    "current": fields.Boolean,
}

workspace_fields = {"id": fields.String, "name": fields.String, "status": fields.String, "created_at": TimestampField}


class TenantListApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def get(self):
        tenants = TenantService.get_join_tenants(current_user)

        for tenant in tenants:
            if tenant.id == current_user.current_tenant_id:
                tenant.current = True  # Set current=True for current tenant
        return {"workspaces": marshal(tenants, tenants_fields)}, 200


class WorkspaceListApi(Resource):
    @setup_required
    @admin_required
    def get(self):
        parser = reqparse.RequestParser()
        parser.add_argument("page", type=inputs.int_range(1, 99999), required=False, default=1, location="args")
        parser.add_argument("limit", type=inputs.int_range(1, 100), required=False, default=20, location="args")
        args = parser.parse_args()

        tenants = (
            db.session.query(Tenant)
            .order_by(Tenant.created_at.desc())
            .paginate(page=args["page"], per_page=args["limit"])
        )

        has_more = False
        if len(tenants.items) == args["limit"]:
            current_page_first_tenant = tenants[-1]
            rest_count = (
                db.session.query(Tenant)
                .filter(
                    Tenant.created_at < current_page_first_tenant.created_at, Tenant.id != current_page_first_tenant.id
                )
                .count()
            )

            if rest_count > 0:
                has_more = True
        total = db.session.query(Tenant).count()
        return {
            "data": marshal(tenants.items, workspace_fields),
            "has_more": has_more,
            "limit": args["limit"],
            "page": args["page"],
            "total": total,
        }, 200


class TenantApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @marshal_with(tenant_fields)
    def get(self):
        if request.path == "/info":
            logging.warning("Deprecated URL /info was used.")

        tenant = current_user.current_tenant

        if tenant.status == TenantStatus.ARCHIVE:
            tenants = TenantService.get_join_tenants(current_user)
            # if there is any tenant, switch to the first one
            if len(tenants) > 0:
                TenantService.switch_tenant(current_user, tenants[0].id)
                tenant = tenants[0]
            # else, raise Unauthorized
            else:
                raise Unauthorized("workspace is archived")

        return WorkspaceService.get_tenant_info(tenant), 200


class SwitchWorkspaceApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument("tenant_id", type=str, required=True, location="json")
        args = parser.parse_args()

        # check if tenant_id is valid, 403 if not
        try:
            TenantService.switch_tenant(current_user, args["tenant_id"])
        except Exception:
            raise AccountNotLinkTenantError("Account not link tenant")

        new_tenant = db.session.query(Tenant).get(args["tenant_id"])  # Get new tenant

        return {"result": "success", "new_tenant": marshal(WorkspaceService.get_tenant_info(new_tenant), tenant_fields)}


class CustomConfigWorkspaceApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @cloud_edition_billing_resource_check("workspace_custom")
    def post(self):
        parser = reqparse.RequestParser()
        parser.add_argument("remove_webapp_brand", type=bool, location="json")
        parser.add_argument("replace_webapp_logo", type=str, location="json")
        args = parser.parse_args()

        tenant = db.session.query(Tenant).filter(Tenant.id == current_user.current_tenant_id).one_or_404()

        custom_config_dict = {
            "remove_webapp_brand": args["remove_webapp_brand"],
            "replace_webapp_logo": args["replace_webapp_logo"]
            if args["replace_webapp_logo"] is not None
            else tenant.custom_config_dict.get("replace_webapp_logo"),
        }

        tenant.custom_config_dict = custom_config_dict
        db.session.commit()

        return {"result": "success", "tenant": marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)}


class WebappLogoWorkspaceApi(Resource):
    @setup_required
    @login_required
    @account_initialization_required
    @cloud_edition_billing_resource_check("workspace_custom")
    def post(self):
        # get file from request
        file = request.files["file"]

        # check file
        if "file" not in request.files:
            raise NoFileUploadedError()

        if len(request.files) > 1:
            raise TooManyFilesError()

        if not file.filename:
            raise FilenameNotExistsError

        extension = file.filename.split(".")[-1]
        if extension.lower() not in {"svg", "png"}:
            raise UnsupportedFileTypeError()

        try:
            upload_file = FileService.upload_file(
                filename=file.filename,
                content=file.read(),
                mimetype=file.mimetype,
                user=current_user,
            )

        except services.errors.file.FileTooLargeError as file_too_large_error:
            raise FileTooLargeError(file_too_large_error.description)
        except services.errors.file.UnsupportedFileTypeError:
            raise UnsupportedFileTypeError()

        return {"id": upload_file.id}, 201


api.add_resource(TenantListApi, "/workspaces")  # GET for getting all tenants
api.add_resource(WorkspaceListApi, "/all-workspaces")  # GET for getting all tenants
api.add_resource(TenantApi, "/workspaces/current", endpoint="workspaces_current")  # GET for getting current tenant info
api.add_resource(TenantApi, "/info", endpoint="info")  # Deprecated
api.add_resource(SwitchWorkspaceApi, "/workspaces/switch")  # POST for switching tenant
api.add_resource(CustomConfigWorkspaceApi, "/workspaces/custom-config")
api.add_resource(WebappLogoWorkspaceApi, "/workspaces/custom-config/webapp-logo/upload")