refactor: use EnumText for TidbAuthBinding.status and MessageFile.type (#33975)

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -33,6 +33,7 @@ from core.rag.models.document import Document
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import Dataset, TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 if TYPE_CHECKING:
     from qdrant_client import grpc  # noqa
@@ -452,7 +453,7 @@ class TidbOnQdrantVectorFactory(AbstractVectorFactory):
                             password=new_cluster["password"],
                             tenant_id=dataset.tenant_id,
                             active=True,
-                            status="ACTIVE",
+                            status=TidbAuthBindingStatus.ACTIVE,
                         )
                         db.session.add(new_tidb_auth_binding)
                         db.session.commit()

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_service.py

@@ -9,6 +9,7 @@ from configs import dify_config
 from extensions.ext_database import db
 from extensions.ext_redis import redis_client
 from models.dataset import TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 
 class TidbService:
@@ -170,7 +171,7 @@ class TidbService:
                 userPrefix = item["userPrefix"]
                 if state == "ACTIVE" and len(userPrefix) > 0:
                     cluster_info = tidb_serverless_list_map[item["clusterId"]]
-                    cluster_info.status = "ACTIVE"
+                    cluster_info.status = TidbAuthBindingStatus.ACTIVE
                     cluster_info.account = f"{userPrefix}.root"
                     db.session.add(cluster_info)
             db.session.commit()

api/models/dataset.py

@@ -45,6 +45,7 @@ from .enums import (
     SegmentStatus,
     SegmentType,
     SummaryStatus,
+    TidbAuthBindingStatus,
 )
 from .model import App, Tag, TagBinding, UploadFile
 from .types import AdjustedJSON, BinaryData, EnumText, LongText, StringUUID, adjusted_json_index
@@ -1242,7 +1243,9 @@ class TidbAuthBinding(TypeBase):
     cluster_id: Mapped[str] = mapped_column(String(255), nullable=False)
     cluster_name: Mapped[str] = mapped_column(String(255), nullable=False)
     active: Mapped[bool] = mapped_column(sa.Boolean, nullable=False, server_default=sa.text("false"))
-    status: Mapped[str] = mapped_column(sa.String(255), nullable=False, server_default=sa.text("'CREATING'"))
+    status: Mapped[TidbAuthBindingStatus] = mapped_column(
+        EnumText(TidbAuthBindingStatus, length=255), nullable=False, server_default=sa.text("'CREATING'")
+    )
     account: Mapped[str] = mapped_column(String(255), nullable=False)
     password: Mapped[str] = mapped_column(String(255), nullable=False)
     created_at: Mapped[datetime] = mapped_column(

api/models/model.py

@@ -21,7 +21,7 @@ from configs import dify_config
 from constants import DEFAULT_FILE_NUMBER_LIMITS
 from core.tools.signature import sign_tool_file
 from dify_graph.enums import WorkflowExecutionStatus
-from dify_graph.file import FILE_MODEL_IDENTITY, File, FileTransferMethod
+from dify_graph.file import FILE_MODEL_IDENTITY, File, FileTransferMethod, FileType
 from dify_graph.file import helpers as file_helpers
 from extensions.storage.storage_type import StorageType
 from libs.helper import generate_string  # type: ignore[import-not-found]
@@ -1785,7 +1785,7 @@ class MessageFile(TypeBase):
         StringUUID, insert_default=lambda: str(uuid4()), default_factory=lambda: str(uuid4()), init=False
     )
     message_id: Mapped[str] = mapped_column(StringUUID, nullable=False)
-    type: Mapped[str] = mapped_column(String(255), nullable=False)
+    type: Mapped[FileType] = mapped_column(EnumText(FileType, length=255), nullable=False)
     transfer_method: Mapped[FileTransferMethod] = mapped_column(
         EnumText(FileTransferMethod, length=255), nullable=False
     )

api/schedule/create_tidb_serverless_task.py

@@ -8,6 +8,7 @@ from configs import dify_config
 from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService
 from extensions.ext_database import db
 from models.dataset import TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 
 @app.celery.task(queue="dataset")
@@ -57,7 +58,7 @@ def create_clusters(batch_size):
                 account=new_cluster["account"],
                 password=new_cluster["password"],
                 active=False,
-                status="CREATING",
+                status=TidbAuthBindingStatus.CREATING,
             )
             db.session.add(tidb_auth_binding)
         db.session.commit()

api/schedule/update_tidb_serverless_status_task.py

@@ -9,6 +9,7 @@ from configs import dify_config
 from core.rag.datasource.vdb.tidb_on_qdrant.tidb_service import TidbService
 from extensions.ext_database import db
 from models.dataset import TidbAuthBinding
+from models.enums import TidbAuthBindingStatus
 
 
 @app.celery.task(queue="dataset")
@@ -18,7 +19,10 @@ def update_tidb_serverless_status_task():
     try:
         # check the number of idle tidb serverless
         tidb_serverless_list = db.session.scalars(
-            select(TidbAuthBinding).where(TidbAuthBinding.active == False, TidbAuthBinding.status == "CREATING")
+            select(TidbAuthBinding).where(
+                TidbAuthBinding.active == False,
+                TidbAuthBinding.status == TidbAuthBindingStatus.CREATING,
+            )
         ).all()
         if len(tidb_serverless_list) == 0:
             return

api/tests/test_containers_integration_tests/services/test_messages_clean_service.py

@@ -8,6 +8,7 @@ import pytest
 from faker import Faker
 from sqlalchemy.orm import Session
 
+from dify_graph.file.enums import FileType
 from enums.cloud_plan import CloudPlan
 from extensions.ext_redis import redis_client
 from models.account import Account, Tenant, TenantAccountJoin, TenantAccountRole
@@ -253,7 +254,7 @@ class TestMessagesCleanServiceIntegration:
         # MessageFile
         file = MessageFile(
             message_id=message.id,
-            type="image",
+            type=FileType.IMAGE,
             transfer_method="local_file",
             url="http://example.com/test.jpg",
             belongs_to=MessageFileBelongsTo.USER,

api/tests/unit_tests/core/app/task_pipeline/test_easy_ui_message_end_files.py

@@ -21,7 +21,7 @@ from sqlalchemy.orm import Session
 
 from core.app.entities.task_entities import MessageEndStreamResponse
 from core.app.task_pipeline.easy_ui_based_generate_task_pipeline import EasyUIBasedGenerateTaskPipeline
-from dify_graph.file.enums import FileTransferMethod
+from dify_graph.file.enums import FileTransferMethod, FileType
 from models.model import MessageFile, UploadFile
 
 
@@ -51,7 +51,7 @@ class TestMessageEndStreamResponseFiles:
         message_file.transfer_method = FileTransferMethod.LOCAL_FILE
         message_file.upload_file_id = str(uuid.uuid4())
         message_file.url = None
-        message_file.type = "image"
+        message_file.type = FileType.IMAGE
         return message_file
 
     @pytest.fixture
@@ -63,7 +63,7 @@ class TestMessageEndStreamResponseFiles:
         message_file.transfer_method = FileTransferMethod.REMOTE_URL
         message_file.upload_file_id = None
         message_file.url = "https://example.com/image.jpg"
-        message_file.type = "image"
+        message_file.type = FileType.IMAGE
         return message_file
 
     @pytest.fixture
@@ -75,7 +75,7 @@ class TestMessageEndStreamResponseFiles:
         message_file.transfer_method = FileTransferMethod.TOOL_FILE
         message_file.upload_file_id = None
         message_file.url = "tool_file_123.png"
-        message_file.type = "image"
+        message_file.type = FileType.IMAGE
         return message_file
 
     @pytest.fixture