Skip to content

Commit

Permalink
feat: add license. (#10403)
Browse files Browse the repository at this point in the history
  • Loading branch information
GarfieldDai authored Nov 15, 2024
1 parent 2fed55a commit c2ce2f8
Show file tree
Hide file tree
Showing 7 changed files with 60 additions and 5 deletions.
3 changes: 3 additions & 0 deletions api/controllers/console/app/app.py
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@
from controllers.console.wraps import (
account_initialization_required,
cloud_edition_billing_resource_check,
enterprise_license_required,
setup_required,
)
from core.ops.ops_trace_manager import OpsTraceManager
Expand All @@ -28,6 +29,7 @@ class AppListApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
def get(self):
"""Get app list"""

Expand Down Expand Up @@ -149,6 +151,7 @@ class AppApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
@get_app_model
@marshal_with(app_detail_fields_with_site)
def get(self, app_model):
Expand Down
3 changes: 2 additions & 1 deletion api/controllers/console/datasets/datasets.py
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
from controllers.console.apikey import api_key_fields, api_key_list
from controllers.console.app.error import ProviderNotInitializeError
from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError, IndexingEstimateError
from controllers.console.wraps import account_initialization_required, setup_required
from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
from core.indexing_runner import IndexingRunner
from core.model_runtime.entities.model_entities import ModelType
Expand Down Expand Up @@ -44,6 +44,7 @@ class DatasetListApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
def get(self):
page = request.args.get("page", default=1, type=int)
limit = request.args.get("limit", default=20, type=int)
Expand Down
6 changes: 6 additions & 0 deletions api/controllers/console/error.py
Original file line number Diff line number Diff line change
Expand Up @@ -86,3 +86,9 @@ class NoFileUploadedError(BaseHTTPException):
error_code = "no_file_uploaded"
description = "Please upload your file."
code = 400


class UnauthorizedAndForceLogout(BaseHTTPException):
error_code = "unauthorized_and_force_logout"
description = "Unauthorized and force logout."
code = 401
3 changes: 2 additions & 1 deletion api/controllers/console/workspace/account.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
InvalidInvitationCodeError,
RepeatPasswordNotMatchError,
)
from controllers.console.wraps import account_initialization_required, setup_required
from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
from extensions.ext_database import db
from fields.member_fields import account_fields
from libs.helper import TimestampField, timezone
Expand Down Expand Up @@ -79,6 +79,7 @@ class AccountProfileApi(Resource):
@login_required
@account_initialization_required
@marshal_with(account_fields)
@enterprise_license_required
def get(self):
return current_user

Expand Down
3 changes: 2 additions & 1 deletion api/controllers/console/workspace/tool_providers.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

from configs import dify_config
from controllers.console import api
from controllers.console.wraps import account_initialization_required, setup_required
from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
from core.model_runtime.utils.encoders import jsonable_encoder
from libs.helper import alphanumeric, uuid_value
from libs.login import login_required
Expand Down Expand Up @@ -549,6 +549,7 @@ class ToolLabelsApi(Resource):
@setup_required
@login_required
@account_initialization_required
@enterprise_license_required
def get(self):
return jsonable_encoder(ToolLabelsService.list_tool_labels())

Expand Down
16 changes: 14 additions & 2 deletions api/controllers/console/wraps.py
Original file line number Diff line number Diff line change
Expand Up @@ -8,10 +8,10 @@
from configs import dify_config
from controllers.console.workspace.error import AccountNotInitializedError
from models.model import DifySetup
from services.feature_service import FeatureService
from services.feature_service import FeatureService, LicenseStatus
from services.operation_service import OperationService

from .error import NotInitValidateError, NotSetupError
from .error import NotInitValidateError, NotSetupError, UnauthorizedAndForceLogout


def account_initialization_required(view):
Expand Down Expand Up @@ -142,3 +142,15 @@ def decorated(*args, **kwargs):
return view(*args, **kwargs)

return decorated


def enterprise_license_required(view):
@wraps(view)
def decorated(*args, **kwargs):
settings = FeatureService.get_system_features()
if settings.license.status in [LicenseStatus.INACTIVE, LicenseStatus.EXPIRED, LicenseStatus.LOST]:
raise UnauthorizedAndForceLogout("Your license is invalid. Please contact your administrator.")

return view(*args, **kwargs)

return decorated
31 changes: 31 additions & 0 deletions api/services/feature_service.py
Original file line number Diff line number Diff line change
@@ -1,3 +1,5 @@
from enum import Enum

from pydantic import BaseModel, ConfigDict

from configs import dify_config
Expand All @@ -20,6 +22,20 @@ class LimitationModel(BaseModel):
limit: int = 0


class LicenseStatus(str, Enum):
NONE = "none"
INACTIVE = "inactive"
ACTIVE = "active"
EXPIRING = "expiring"
EXPIRED = "expired"
LOST = "lost"


class LicenseModel(BaseModel):
status: LicenseStatus = LicenseStatus.NONE
expired_at: str = ""


class FeatureModel(BaseModel):
billing: BillingModel = BillingModel()
members: LimitationModel = LimitationModel(size=0, limit=1)
Expand Down Expand Up @@ -47,6 +63,7 @@ class SystemFeatureModel(BaseModel):
enable_social_oauth_login: bool = False
is_allow_register: bool = False
is_allow_create_workspace: bool = False
license: LicenseModel = LicenseModel()


class FeatureService:
Expand Down Expand Up @@ -131,17 +148,31 @@ def _fulfill_params_from_enterprise(cls, features):

if "sso_enforced_for_signin" in enterprise_info:
features.sso_enforced_for_signin = enterprise_info["sso_enforced_for_signin"]

if "sso_enforced_for_signin_protocol" in enterprise_info:
features.sso_enforced_for_signin_protocol = enterprise_info["sso_enforced_for_signin_protocol"]

if "sso_enforced_for_web" in enterprise_info:
features.sso_enforced_for_web = enterprise_info["sso_enforced_for_web"]

if "sso_enforced_for_web_protocol" in enterprise_info:
features.sso_enforced_for_web_protocol = enterprise_info["sso_enforced_for_web_protocol"]

if "enable_email_code_login" in enterprise_info:
features.enable_email_code_login = enterprise_info["enable_email_code_login"]

if "enable_email_password_login" in enterprise_info:
features.enable_email_password_login = enterprise_info["enable_email_password_login"]

if "is_allow_register" in enterprise_info:
features.is_allow_register = enterprise_info["is_allow_register"]

if "is_allow_create_workspace" in enterprise_info:
features.is_allow_create_workspace = enterprise_info["is_allow_create_workspace"]

if "license" in enterprise_info:
if "status" in enterprise_info["license"]:
features.license.status = enterprise_info["license"]["status"]

if "expired_at" in enterprise_info["license"]:
features.license.expired_at = enterprise_info["license"]["expired_at"]

0 comments on commit c2ce2f8

Please sign in to comment.