refactor(api): migrate console extension endpoint from api.model to BaseModel (#35189)

Co-authored-by: ai-hpc <ai-hpc@users.noreply.github.com>

api/controllers/console/extension.py

@@ -1,15 +1,18 @@
+from datetime import datetime
+from typing import Any
+
 from flask import request
-from flask_restx import Resource, fields, marshal_with
+from flask_restx import Resource
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, TypeAdapter, field_validator
 
 from constants import HIDDEN_VALUE
-from fields.api_based_extension_fields import api_based_extension_fields
+from fields.base import ResponseModel
 from libs.login import current_account_with_tenant, login_required
 from models.api_based_extension import APIBasedExtension
 from services.api_based_extension_service import APIBasedExtensionService
 from services.code_based_extension_service import CodeBasedExtensionService
 
-from ..common.schema import register_schema_models
+from ..common.schema import DEFAULT_REF_TEMPLATE_SWAGGER_2_0, register_schema_models
 from . import console_ns
 from .wraps import account_initialization_required, setup_required
 
@@ -24,12 +27,52 @@ class APIBasedExtensionPayload(BaseModel):
     api_key: str = Field(description="API key for authentication")
 
 
-register_schema_models(console_ns, APIBasedExtensionPayload)
+class CodeBasedExtensionResponse(ResponseModel):
+    module: str = Field(description="Module name")
+    data: Any = Field(description="Extension data")
 
 
-api_based_extension_model = console_ns.model("ApiBasedExtensionModel", api_based_extension_fields)
+def _mask_api_key(api_key: str) -> str:
+    if not api_key:
+        return api_key
+    if len(api_key) <= 8:
+        return api_key[0] + "******" + api_key[-1]
+    return api_key[:3] + "******" + api_key[-3:]
 
-api_based_extension_list_model = fields.List(fields.Nested(api_based_extension_model))
+
+def _to_timestamp(value: datetime | int | None) -> int | None:
+    if isinstance(value, datetime):
+        return int(value.timestamp())
+    return value
+
+
+class APIBasedExtensionResponse(ResponseModel):
+    id: str
+    name: str
+    api_endpoint: str
+    api_key: str
+    created_at: int | None = None
+
+    @field_validator("api_key", mode="before")
+    @classmethod
+    def _normalize_api_key(cls, value: str) -> str:
+        return _mask_api_key(value)
+
+    @field_validator("created_at", mode="before")
+    @classmethod
+    def _normalize_created_at(cls, value: datetime | int | None) -> int | None:
+        return _to_timestamp(value)
+
+
+register_schema_models(console_ns, APIBasedExtensionPayload, CodeBasedExtensionResponse, APIBasedExtensionResponse)
+console_ns.schema_model(
+    "APIBasedExtensionListResponse",
+    TypeAdapter(list[APIBasedExtensionResponse]).json_schema(ref_template=DEFAULT_REF_TEMPLATE_SWAGGER_2_0),
+)
+
+
+def _serialize_api_based_extension(extension: APIBasedExtension) -> dict[str, Any]:
+    return APIBasedExtensionResponse.model_validate(extension, from_attributes=True).model_dump(mode="json")
 
 
 @console_ns.route("/code-based-extension")
@@ -40,10 +83,7 @@ class CodeBasedExtensionAPI(Resource):
     @console_ns.response(
         200,
         "Success",
-        console_ns.model(
+        console_ns.models[CodeBasedExtensionResponse.__name__],
-            "CodeBasedExtensionResponse",
-            {"module": fields.String(description="Module name"), "data": fields.Raw(description="Extension data")},
-        ),
     )
     @setup_required
     @login_required
@@ -51,30 +91,34 @@ class CodeBasedExtensionAPI(Resource):
     def get(self):
         query = CodeBasedExtensionQuery.model_validate(request.args.to_dict(flat=True))  # type: ignore
 
-        return {"module": query.module, "data": CodeBasedExtensionService.get_code_based_extension(query.module)}
+        return CodeBasedExtensionResponse(
+            module=query.module,
+            data=CodeBasedExtensionService.get_code_based_extension(query.module),
+        ).model_dump(mode="json")
 
 
 @console_ns.route("/api-based-extension")
 class APIBasedExtensionAPI(Resource):
     @console_ns.doc("get_api_based_extensions")
     @console_ns.doc(description="Get all API-based extensions for current tenant")
-    @console_ns.response(200, "Success", api_based_extension_list_model)
+    @console_ns.response(200, "Success", console_ns.models["APIBasedExtensionListResponse"])
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(api_based_extension_model)
     def get(self):
         _, tenant_id = current_account_with_tenant()
-        return APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
+        return [
+            _serialize_api_based_extension(extension)
+            for extension in APIBasedExtensionService.get_all_by_tenant_id(tenant_id)
+        ]
 
     @console_ns.doc("create_api_based_extension")
     @console_ns.doc(description="Create a new API-based extension")
     @console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
-    @console_ns.response(201, "Extension created successfully", api_based_extension_model)
+    @console_ns.response(201, "Extension created successfully", console_ns.models[APIBasedExtensionResponse.__name__])
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(api_based_extension_model)
     def post(self):
         payload = APIBasedExtensionPayload.model_validate(console_ns.payload or {})
         _, current_tenant_id = current_account_with_tenant()
@@ -86,7 +130,7 @@ class APIBasedExtensionAPI(Resource):
             api_key=payload.api_key,
         )
 
-        return APIBasedExtensionService.save(extension_data)
+        return _serialize_api_based_extension(APIBasedExtensionService.save(extension_data))
 
 
 @console_ns.route("/api-based-extension/<uuid:id>")
@@ -94,26 +138,26 @@ class APIBasedExtensionDetailAPI(Resource):
     @console_ns.doc("get_api_based_extension")
     @console_ns.doc(description="Get API-based extension by ID")
     @console_ns.doc(params={"id": "Extension ID"})
-    @console_ns.response(200, "Success", api_based_extension_model)
+    @console_ns.response(200, "Success", console_ns.models[APIBasedExtensionResponse.__name__])
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(api_based_extension_model)
     def get(self, id):
         api_based_extension_id = str(id)
         _, tenant_id = current_account_with_tenant()
 
-        return APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+        return _serialize_api_based_extension(
+            APIBasedExtensionService.get_with_tenant_id(tenant_id, api_based_extension_id)
+        )
 
     @console_ns.doc("update_api_based_extension")
     @console_ns.doc(description="Update API-based extension")
     @console_ns.doc(params={"id": "Extension ID"})
     @console_ns.expect(console_ns.models[APIBasedExtensionPayload.__name__])
-    @console_ns.response(200, "Extension updated successfully", api_based_extension_model)
+    @console_ns.response(200, "Extension updated successfully", console_ns.models[APIBasedExtensionResponse.__name__])
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(api_based_extension_model)
     def post(self, id):
         api_based_extension_id = str(id)
         _, current_tenant_id = current_account_with_tenant()
@@ -128,7 +172,7 @@ class APIBasedExtensionDetailAPI(Resource):
         if payload.api_key != HIDDEN_VALUE:
             extension_data_from_db.api_key = payload.api_key
 
-        return APIBasedExtensionService.save(extension_data_from_db)
+        return _serialize_api_based_extension(APIBasedExtensionService.save(extension_data_from_db))
 
     @console_ns.doc("delete_api_based_extension")
     @console_ns.doc(description="Delete API-based extension")