refactor: use EnumText for ApiToken.type (#33961)

api/controllers/console/apikey.py

@@ -9,6 +9,7 @@ from extensions.ext_database import db
 from libs.helper import TimestampField
 from libs.login import current_account_with_tenant, login_required
 from models.dataset import Dataset
+from models.enums import ApiTokenType
 from models.model import ApiToken, App
 from services.api_token_service import ApiTokenCache
 
@@ -47,7 +48,7 @@ def _get_resource(resource_id, tenant_id, resource_model):
 class BaseApiKeyListResource(Resource):
     method_decorators = [account_initialization_required, login_required, setup_required]
 
-    resource_type: str | None = None
+    resource_type: ApiTokenType | None = None
     resource_model: type | None = None
     resource_id_field: str | None = None
     token_prefix: str | None = None
@@ -91,6 +92,7 @@ class BaseApiKeyListResource(Resource):
             )
 
         key = ApiToken.generate_api_key(self.token_prefix or "", 24)
+        assert self.resource_type is not None, "resource_type must be set"
         api_token = ApiToken()
         setattr(api_token, self.resource_id_field, resource_id)
         api_token.tenant_id = current_tenant_id
@@ -104,7 +106,7 @@ class BaseApiKeyListResource(Resource):
 class BaseApiKeyResource(Resource):
     method_decorators = [account_initialization_required, login_required, setup_required]
 
-    resource_type: str | None = None
+    resource_type: ApiTokenType | None = None
     resource_model: type | None = None
     resource_id_field: str | None = None
 
@@ -159,7 +161,7 @@ class AppApiKeyListResource(BaseApiKeyListResource):
         """Create a new API key for an app"""
         return super().post(resource_id)
 
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = App
     resource_id_field = "app_id"
     token_prefix = "app-"
@@ -175,7 +177,7 @@ class AppApiKeyResource(BaseApiKeyResource):
         """Delete an API key for an app"""
         return super().delete(resource_id, api_key_id)
 
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = App
     resource_id_field = "app_id"
 
@@ -199,7 +201,7 @@ class DatasetApiKeyListResource(BaseApiKeyListResource):
         """Create a new API key for a dataset"""
         return super().post(resource_id)
 
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
     resource_model = Dataset
     resource_id_field = "dataset_id"
     token_prefix = "ds-"
@@ -215,6 +217,6 @@ class DatasetApiKeyResource(BaseApiKeyResource):
         """Delete an API key for a dataset"""
         return super().delete(resource_id, api_key_id)
 
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
     resource_model = Dataset
     resource_id_field = "dataset_id"

api/controllers/console/datasets/datasets.py

@@ -54,7 +54,7 @@ from fields.document_fields import document_status_fields
 from libs.login import current_account_with_tenant, login_required
 from models import ApiToken, Dataset, Document, DocumentSegment, UploadFile
 from models.dataset import DatasetPermission, DatasetPermissionEnum
-from models.enums import SegmentStatus
+from models.enums import ApiTokenType, SegmentStatus
 from models.provider_ids import ModelProviderID
 from services.api_token_service import ApiTokenCache
 from services.dataset_service import DatasetPermissionService, DatasetService, DocumentService
@@ -777,7 +777,7 @@ class DatasetIndexingStatusApi(Resource):
 class DatasetApiKeyApi(Resource):
     max_keys = 10
     token_prefix = "dataset-"
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
 
     @console_ns.doc("get_dataset_api_keys")
     @console_ns.doc(description="Get dataset API keys")
@@ -826,7 +826,7 @@ class DatasetApiKeyApi(Resource):
 
 @console_ns.route("/datasets/api-keys/<uuid:api_key_id>")
 class DatasetApiDeleteApi(Resource):
-    resource_type = "dataset"
+    resource_type = ApiTokenType.DATASET
 
     @console_ns.doc("delete_dataset_api_key")
     @console_ns.doc(description="Delete dataset API key")

api/models/enums.py

@@ -323,3 +323,10 @@ class ProviderQuotaType(StrEnum):
             if member.value == value:
                 return member
         raise ValueError(f"No matching enum found for value '{value}'")
+
+
+class ApiTokenType(StrEnum):
+    """API Token type"""
+
+    APP = "app"
+    DATASET = "dataset"

api/models/model.py

@@ -31,6 +31,7 @@ from .account import Account, Tenant
 from .base import Base, TypeBase, gen_uuidv4_string
 from .engine import db
 from .enums import (
+    ApiTokenType,
     AppMCPServerStatus,
     AppStatus,
     BannerStatus,
@@ -2095,7 +2096,7 @@ class ApiToken(Base):  # bug: this uses setattr so idk the field.
     id = mapped_column(StringUUID, default=lambda: str(uuid4()))
     app_id = mapped_column(StringUUID, nullable=True)
     tenant_id = mapped_column(StringUUID, nullable=True)
-    type = mapped_column(String(16), nullable=False)
+    type: Mapped[ApiTokenType] = mapped_column(EnumText(ApiTokenType, length=16), nullable=False)
     token: Mapped[str] = mapped_column(String(255), nullable=False)
     last_used_at = mapped_column(sa.DateTime, nullable=True)
     created_at = mapped_column(sa.DateTime, nullable=False, server_default=func.current_timestamp())

api/tests/integration_tests/libs/test_api_token_cache_integration.py

@@ -13,6 +13,7 @@ from unittest.mock import patch
 import pytest
 
 from extensions.ext_redis import redis_client
+from models.enums import ApiTokenType
 from models.model import ApiToken
 from services.api_token_service import ApiTokenCache, CachedApiToken
 
@@ -279,7 +280,7 @@ class TestEndToEndCacheFlow:
         test_token = ApiToken()
         test_token.id = "test-e2e-id"
         test_token.token = test_token_value
-        test_token.type = test_scope
+        test_token.type = ApiTokenType.APP
         test_token.app_id = "test-app"
         test_token.tenant_id = "test-tenant"
         test_token.last_used_at = None

api/tests/unit_tests/controllers/console/test_apikey.py

@@ -8,6 +8,7 @@ from controllers.console.apikey import (
     BaseApiKeyResource,
     _get_resource,
 )
+from models.enums import ApiTokenType
 
 
 @pytest.fixture
@@ -45,14 +46,14 @@ def bypass_permissions():
 
 
 class DummyApiKeyListResource(BaseApiKeyListResource):
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = MagicMock()
     resource_id_field = "app_id"
     token_prefix = "app-"
 
 
 class DummyApiKeyResource(BaseApiKeyResource):
-    resource_type = "app"
+    resource_type = ApiTokenType.APP
     resource_model = MagicMock()
     resource_id_field = "app_id"