dify

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

@@ -0,0 +1,268 @@
+import logging
+
+from flask import request
+from flask_restx import Resource, fields, inputs, marshal, marshal_with, reqparse
+from sqlalchemy import select
+from werkzeug.exceptions import Unauthorized
+
+import services
+from controllers.common.errors import (
+    FilenameNotExistsError,
+    FileTooLargeError,
+    NoFileUploadedError,
+    TooManyFilesError,
+    UnsupportedFileTypeError,
+)
+from controllers.console import console_ns
+from controllers.console.admin import admin_required
+from controllers.console.error import AccountNotLinkTenantError
+from controllers.console.wraps import (
+    account_initialization_required,
+    cloud_edition_billing_resource_check,
+    setup_required,
+)
+from enums.cloud_plan import CloudPlan
+from extensions.ext_database import db
+from libs.helper import TimestampField
+from libs.login import current_account_with_tenant, login_required
+from models.account import Tenant, TenantStatus
+from services.account_service import TenantService
+from services.feature_service import FeatureService
+from services.file_service import FileService
+from services.workspace_service import WorkspaceService
+
+logger = logging.getLogger(__name__)
+
+
+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}
+
+
+@console_ns.route("/workspaces")
+class TenantListApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self):
+        current_user, current_tenant_id = current_account_with_tenant()
+        tenants = TenantService.get_join_tenants(current_user)
+        tenant_dicts = []
+
+        for tenant in tenants:
+            features = FeatureService.get_features(tenant.id)
+
+            # Create a dictionary with tenant attributes
+            tenant_dict = {
+                "id": tenant.id,
+                "name": tenant.name,
+                "status": tenant.status,
+                "created_at": tenant.created_at,
+                "plan": features.billing.subscription.plan if features.billing.enabled else CloudPlan.SANDBOX,
+                "current": tenant.id == current_tenant_id if current_tenant_id else False,
+            }
+
+            tenant_dicts.append(tenant_dict)
+
+        return {"workspaces": marshal(tenant_dicts, tenants_fields)}, 200
+
+
+@console_ns.route("/all-workspaces")
+class WorkspaceListApi(Resource):
+    @setup_required
+    @admin_required
+    def get(self):
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("page", type=inputs.int_range(1, 99999), required=False, default=1, location="args")
+            .add_argument("limit", type=inputs.int_range(1, 100), required=False, default=20, location="args")
+        )
+        args = parser.parse_args()
+
+        stmt = select(Tenant).order_by(Tenant.created_at.desc())
+        tenants = db.paginate(select=stmt, page=args["page"], per_page=args["limit"], error_out=False)
+        has_more = False
+
+        if tenants.has_next:
+            has_more = True
+
+        return {
+            "data": marshal(tenants.items, workspace_fields),
+            "has_more": has_more,
+            "limit": args["limit"],
+            "page": args["page"],
+            "total": tenants.total,
+        }, 200
+
+
+@console_ns.route("/workspaces/current", endpoint="workspaces_current")
+@console_ns.route("/info", endpoint="info")  # Deprecated
+class TenantApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @marshal_with(tenant_fields)
+    def post(self):
+        if request.path == "/info":
+            logger.warning("Deprecated URL /info was used.")
+
+        current_user, _ = current_account_with_tenant()
+        tenant = current_user.current_tenant
+        if not tenant:
+            raise ValueError("No 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
+
+
+parser_switch = reqparse.RequestParser().add_argument("tenant_id", type=str, required=True, location="json")
+
+
+@console_ns.route("/workspaces/switch")
+class SwitchWorkspaceApi(Resource):
+    @console_ns.expect(parser_switch)
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def post(self):
+        current_user, _ = current_account_with_tenant()
+        args = parser_switch.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
+        if new_tenant is None:
+            raise ValueError("Tenant not found")
+
+        return {"result": "success", "new_tenant": marshal(WorkspaceService.get_tenant_info(new_tenant), tenant_fields)}
+
+
+@console_ns.route("/workspaces/custom-config")
+class CustomConfigWorkspaceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @cloud_edition_billing_resource_check("workspace_custom")
+    def post(self):
+        _, current_tenant_id = current_account_with_tenant()
+        parser = (
+            reqparse.RequestParser()
+            .add_argument("remove_webapp_brand", type=bool, location="json")
+            .add_argument("replace_webapp_logo", type=str, location="json")
+        )
+        args = parser.parse_args()
+        tenant = db.get_or_404(Tenant, current_tenant_id)
+
+        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)}
+
+
+@console_ns.route("/workspaces/custom-config/webapp-logo/upload")
+class WebappLogoWorkspaceApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @cloud_edition_billing_resource_check("workspace_custom")
+    def post(self):
+        current_user, _ = current_account_with_tenant()
+        # check file
+        if "file" not in request.files:
+            raise NoFileUploadedError()
+
+        if len(request.files) > 1:
+            raise TooManyFilesError()
+
+        # get file from request
+        file = request.files["file"]
+        if not file.filename:
+            raise FilenameNotExistsError
+
+        extension = file.filename.split(".")[-1]
+        if extension.lower() not in {"svg", "png"}:
+            raise UnsupportedFileTypeError()
+
+        try:
+            upload_file = FileService(db.engine).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
+
+
+parser_info = reqparse.RequestParser().add_argument("name", type=str, required=True, location="json")
+
+
+@console_ns.route("/workspaces/info")
+class WorkspaceInfoApi(Resource):
+    @console_ns.expect(parser_info)
+    @setup_required
+    @login_required
+    @account_initialization_required
+    # Change workspace name
+    def post(self):
+        _, current_tenant_id = current_account_with_tenant()
+        args = parser_info.parse_args()
+
+        if not current_tenant_id:
+            raise ValueError("No current tenant")
+        tenant = db.get_or_404(Tenant, current_tenant_id)
+        tenant.name = args["name"]
+        db.session.commit()
+
+        return {"result": "success", "tenant": marshal(WorkspaceService.get_tenant_info(tenant), tenant_fields)}