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.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.setup import setup_required from controllers.console.wraps import account_initialization_required, cloud_edition_billing_resource_check 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() extension = file.filename.split('.')[-1] if extension.lower() not in ['svg', 'png']: raise UnsupportedFileTypeError() try: upload_file = FileService.upload_file(file, current_user, True) 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')