refactor(api): migrate console.datasets.data_source to BaseModel (#36624)

api/controllers/console/datasets/data_source.py

@@ -1,33 +1,29 @@
 import json
 from collections.abc import Generator
+from datetime import datetime
 from typing import Any, Literal, cast
 from uuid import UUID
 
 from flask import request
-from flask_restx import Resource, fields, marshal_with
-from pydantic import BaseModel, Field
+from flask_restx import Resource
+from pydantic import BaseModel, Field, field_serializer
 from sqlalchemy import select
 from sqlalchemy.orm import sessionmaker
 from werkzeug.exceptions import NotFound
 
 from controllers.common.fields import SimpleResultResponse, TextContentResponse
-from controllers.common.schema import get_or_create_model, register_response_schema_models, register_schema_model
+from controllers.common.schema import query_params_from_model, register_response_schema_models, register_schema_models
 from core.datasource.entities.datasource_entities import DatasourceProviderType, OnlineDocumentPagesMessage
 from core.datasource.online_document.online_document_plugin import OnlineDocumentDatasourcePlugin
+from core.entities.knowledge_entities import IndexingEstimate
 from core.indexing_runner import IndexingRunner
 from core.rag.extractor.entity.datasource_type import DatasourceType
 from core.rag.extractor.entity.extract_setting import ExtractSetting, NotionInfo
 from core.rag.extractor.notion_extractor import NotionExtractor
 from extensions.ext_database import db
-from fields.data_source_fields import (
-    integrate_fields,
-    integrate_icon_fields,
-    integrate_list_fields,
-    integrate_notion_info_list_fields,
-    integrate_page_fields,
-    integrate_workspace_fields,
-)
+from fields.base import ResponseModel
 from libs.datetime_utils import naive_utc_now
+from libs.helper import dump_response, to_timestamp
 from libs.login import current_account_with_tenant, login_required
 from models import DataSourceOauthBinding, Document
 from services.dataset_service import DatasetService, DocumentService
@@ -54,50 +50,74 @@ class DataSourceNotionPreviewQuery(BaseModel):
     credential_id: str = Field(..., description="Credential ID", min_length=1)
 
 
-register_schema_model(console_ns, NotionEstimatePayload)
-register_response_schema_models(console_ns, SimpleResultResponse, TextContentResponse)
+class DataSourceIntegrateIconResponse(ResponseModel):
+    type: str | None = None
+    url: str | None = None
+    emoji: str | None = None
 
 
-integrate_icon_model = get_or_create_model("DataSourceIntegrateIcon", integrate_icon_fields)
+class DataSourceIntegratePageResponse(ResponseModel):
+    page_name: str
+    page_id: str
+    page_icon: DataSourceIntegrateIconResponse | None
+    parent_id: str
+    type: str
 
-integrate_page_fields_copy = integrate_page_fields.copy()
-integrate_page_fields_copy["page_icon"] = fields.Nested(integrate_icon_model, allow_null=True)
-integrate_page_model = get_or_create_model("DataSourceIntegratePage", integrate_page_fields_copy)
 
-integrate_workspace_fields_copy = integrate_workspace_fields.copy()
-integrate_workspace_fields_copy["pages"] = fields.List(fields.Nested(integrate_page_model))
-integrate_workspace_model = get_or_create_model("DataSourceIntegrateWorkspace", integrate_workspace_fields_copy)
+class DataSourceIntegrateWorkspaceResponse(ResponseModel):
+    workspace_name: str | None
+    workspace_id: str | None
+    workspace_icon: str | None
+    pages: list[DataSourceIntegratePageResponse]
+    total: int
 
-integrate_fields_copy = integrate_fields.copy()
-integrate_fields_copy["source_info"] = fields.Nested(integrate_workspace_model)
-integrate_model = get_or_create_model("DataSourceIntegrate", integrate_fields_copy)
 
-integrate_list_fields_copy = integrate_list_fields.copy()
-integrate_list_fields_copy["data"] = fields.List(fields.Nested(integrate_model))
-integrate_list_model = get_or_create_model("DataSourceIntegrateList", integrate_list_fields_copy)
+class DataSourceIntegrateResponse(ResponseModel):
+    id: str | None
+    provider: str
+    created_at: datetime | int | None
+    is_bound: bool
+    disabled: bool | None
+    link: str
+    source_info: DataSourceIntegrateWorkspaceResponse | None
 
-notion_page_fields = {
-    "page_name": fields.String,
-    "page_id": fields.String,
-    "page_icon": fields.Nested(integrate_icon_model, allow_null=True),
-    "is_bound": fields.Boolean,
-    "parent_id": fields.String,
-    "type": fields.String,
-}
-notion_page_model = get_or_create_model("NotionIntegratePage", notion_page_fields)
+    @field_serializer("created_at")
+    def serialize_created_at(self, value: datetime | int | None) -> int | None:
+        return to_timestamp(value)
 
-notion_workspace_fields = {
-    "workspace_name": fields.String,
-    "workspace_id": fields.String,
-    "workspace_icon": fields.String,
-    "pages": fields.List(fields.Nested(notion_page_model)),
-}
-notion_workspace_model = get_or_create_model("NotionIntegrateWorkspace", notion_workspace_fields)
 
-integrate_notion_info_list_fields_copy = integrate_notion_info_list_fields.copy()
-integrate_notion_info_list_fields_copy["notion_info"] = fields.List(fields.Nested(notion_workspace_model))
-integrate_notion_info_list_model = get_or_create_model(
-    "NotionIntegrateInfoList", integrate_notion_info_list_fields_copy
+class DataSourceIntegrateListResponse(ResponseModel):
+    data: list[DataSourceIntegrateResponse]
+
+
+class NotionIntegratePageResponse(ResponseModel):
+    page_name: str
+    page_id: str
+    page_icon: DataSourceIntegrateIconResponse | None
+    parent_id: str | None
+    type: str
+    is_bound: bool
+
+
+class NotionIntegrateWorkspaceResponse(ResponseModel):
+    workspace_name: str | None
+    workspace_id: str | None
+    workspace_icon: str | None
+    pages: list[NotionIntegratePageResponse]
+
+
+class NotionIntegrateInfoListResponse(ResponseModel):
+    notion_info: list[NotionIntegrateWorkspaceResponse]
+
+
+register_schema_models(console_ns, NotionEstimatePayload)
+register_response_schema_models(
+    console_ns,
+    DataSourceIntegrateListResponse,
+    IndexingEstimate,
+    NotionIntegrateInfoListResponse,
+    SimpleResultResponse,
+    TextContentResponse,
 )
 
 
@@ -109,8 +129,8 @@ class DataSourceApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(integrate_list_model)
-    def get(self):
+    @console_ns.response(200, "Success", console_ns.models[DataSourceIntegrateListResponse.__name__])
+    def get(self) -> tuple[dict[str, Any], int]:
         _, current_tenant_id = current_account_with_tenant()
 
         # get workspace data source integrates
@@ -154,19 +174,19 @@ class DataSourceApi(Resource):
                         "link": f"{base_url}{data_source_oauth_base_path}/{provider}",
                     }
                 )
-        return {"data": integrate_data}, 200
+        return dump_response(DataSourceIntegrateListResponse, {"data": integrate_data}), 200
 
     @setup_required
     @login_required
     @account_initialization_required
     @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
-    def patch(self, binding_id, action: Literal["enable", "disable"]):
+    def patch(self, binding_id: UUID, action: Literal["enable", "disable"]) -> tuple[dict[str, str], int]:
         _, current_tenant_id = current_account_with_tenant()
-        binding_id = str(binding_id)
+        binding_id_str = str(binding_id)
         with sessionmaker(db.engine, expire_on_commit=False).begin() as session:
             data_source_binding = session.execute(
                 select(DataSourceOauthBinding).where(
-                    DataSourceOauthBinding.id == binding_id, DataSourceOauthBinding.tenant_id == current_tenant_id
+                    DataSourceOauthBinding.id == binding_id_str, DataSourceOauthBinding.tenant_id == current_tenant_id
                 )
             ).scalar_one_or_none()
         if data_source_binding is None:
@@ -198,12 +218,12 @@ class DataSourceNotionListApi(Resource):
     @setup_required
     @login_required
     @account_initialization_required
-    @marshal_with(integrate_notion_info_list_model)
-    def get(self):
+    @console_ns.doc(params=query_params_from_model(DataSourceNotionListQuery))
+    @console_ns.response(200, "Success", console_ns.models[NotionIntegrateInfoListResponse.__name__])
+    def get(self) -> tuple[dict[str, Any], int]:
         current_user, current_tenant_id = current_account_with_tenant()
 
-        query = DataSourceNotionListQuery.model_validate(request.args.to_dict())
-
+        query = DataSourceNotionListQuery.model_validate(request.args.to_dict(flat=True))
         datasource_provider_service = DatasourceProviderService()
         credential = datasource_provider_service.get_datasource_credentials(
             tenant_id=current_tenant_id,
@@ -278,22 +298,23 @@ class DataSourceNotionListApi(Resource):
                             pages.append(page_info)
             except Exception as e:
                 raise e
-            return {"notion_info": {**workspace_info, "pages": pages}}, 200
+            notion_info = [{**workspace_info, "pages": pages}] if workspace_info else []
+            return dump_response(NotionIntegrateInfoListResponse, {"notion_info": notion_info}), 200
 
 
-@console_ns.route(
-    "/notion/pages/<uuid:page_id>/<string:page_type>/preview",
-    "/datasets/notion-indexing-estimate",
-)
-class DataSourceNotionApi(Resource):
+@console_ns.route("/notion/pages/<uuid:page_id>/<string:page_type>/preview")
+class DataSourceNotionPreviewApi(Resource):
+    """Preview one authorized Notion page through the datasource credential."""
+
     @setup_required
     @login_required
     @account_initialization_required
+    @console_ns.doc(params=query_params_from_model(DataSourceNotionPreviewQuery))
     @console_ns.response(200, "Success", console_ns.models[TextContentResponse.__name__])
-    def get(self, page_id: UUID, page_type: str):
+    def get(self, page_id: UUID, page_type: str) -> tuple[dict[str, str], int]:
         _, current_tenant_id = current_account_with_tenant()
 
-        query = DataSourceNotionPreviewQuery.model_validate(request.args.to_dict())
+        query = DataSourceNotionPreviewQuery.model_validate(request.args.to_dict(flat=True))
 
         datasource_provider_service = DatasourceProviderService()
         credential = datasource_provider_service.get_datasource_credentials(
@@ -316,11 +337,17 @@ class DataSourceNotionApi(Resource):
         text_docs = extractor.extract()
         return {"content": "\n".join([doc.page_content for doc in text_docs])}, 200
 
+
+@console_ns.route("/datasets/notion-indexing-estimate")
+class DataSourceNotionIndexingEstimateApi(Resource):
+    """Estimate indexing work for selected Notion pages."""
+
     @setup_required
     @login_required
     @account_initialization_required
     @console_ns.expect(console_ns.models[NotionEstimatePayload.__name__])
-    def post(self):
+    @console_ns.response(200, "Success", console_ns.models[IndexingEstimate.__name__])
+    def post(self) -> tuple[dict[str, Any], int]:
         _, current_tenant_id = current_account_with_tenant()
 
         payload = NotionEstimatePayload.model_validate(console_ns.payload or {})
@@ -355,7 +382,7 @@ class DataSourceNotionApi(Resource):
             args["doc_form"],
             args["doc_language"],
         )
-        return response.model_dump(), 200
+        return dump_response(IndexingEstimate, response), 200
 
 
 @console_ns.route("/datasets/<uuid:dataset_id>/notion/sync")
@@ -364,7 +391,7 @@ class DataSourceNotionDatasetSyncApi(Resource):
     @login_required
     @account_initialization_required
     @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
-    def get(self, dataset_id: UUID):
+    def get(self, dataset_id: UUID) -> tuple[dict[str, str], int]:
         dataset_id_str = str(dataset_id)
         dataset = DatasetService.get_dataset(dataset_id_str)
         if dataset is None:
@@ -382,7 +409,7 @@ class DataSourceNotionDocumentSyncApi(Resource):
     @login_required
     @account_initialization_required
     @console_ns.response(200, "Success", console_ns.models[SimpleResultResponse.__name__])
-    def get(self, dataset_id: UUID, document_id: UUID):
+    def get(self, dataset_id: UUID, document_id: UUID) -> tuple[dict[str, str], int]:
         dataset_id_str = str(dataset_id)
         document_id_str = str(document_id)
         dataset = DatasetService.get_dataset(dataset_id_str)

api/fields/data_source_fields.py

@@ -1,55 +0,0 @@
-from flask_restx import fields
-
-from libs.helper import TimestampField
-
-integrate_icon_fields = {"type": fields.String, "url": fields.String, "emoji": fields.String}
-
-integrate_page_fields = {
-    "page_name": fields.String,
-    "page_id": fields.String,
-    "page_icon": fields.Nested(integrate_icon_fields, allow_null=True),
-    "is_bound": fields.Boolean,
-    "parent_id": fields.String,
-    "type": fields.String,
-}
-
-integrate_workspace_fields = {
-    "workspace_name": fields.String,
-    "workspace_id": fields.String,
-    "workspace_icon": fields.String,
-    "pages": fields.List(fields.Nested(integrate_page_fields)),
-}
-
-integrate_notion_info_list_fields = {
-    "notion_info": fields.List(fields.Nested(integrate_workspace_fields)),
-}
-
-integrate_page_fields = {
-    "page_name": fields.String,
-    "page_id": fields.String,
-    "page_icon": fields.Nested(integrate_icon_fields, allow_null=True),
-    "parent_id": fields.String,
-    "type": fields.String,
-}
-
-integrate_workspace_fields = {
-    "workspace_name": fields.String,
-    "workspace_id": fields.String,
-    "workspace_icon": fields.String,
-    "pages": fields.List(fields.Nested(integrate_page_fields)),
-    "total": fields.Integer,
-}
-
-integrate_fields = {
-    "id": fields.String,
-    "provider": fields.String,
-    "created_at": TimestampField,
-    "is_bound": fields.Boolean,
-    "disabled": fields.Boolean,
-    "link": fields.String,
-    "source_info": fields.Nested(integrate_workspace_fields),
-}
-
-integrate_list_fields = {
-    "data": fields.List(fields.Nested(integrate_fields)),
-}

api/openapi/markdown/console-swagger.md

@@ -4326,9 +4326,9 @@ Get compliance document download link
 #### GET
 ##### Responses
 
-| Code | Description |
-| ---- | ----------- |
-| 200 | Success |
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Success | [DataSourceIntegrateListResponse](#datasourceintegratelistresponse) |
 
 #### PATCH
 ##### Responses
@@ -4349,9 +4349,9 @@ Get compliance document download link
 
 ##### Responses
 
-| Code | Description |
-| ---- | ----------- |
-| 200 | Success |
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Success | [DataSourceIntegrateListResponse](#datasourceintegratelistresponse) |
 
 #### PATCH
 ##### Parameters
@@ -4662,13 +4662,6 @@ Initialize dataset with documents
 
 ### /datasets/notion-indexing-estimate
 
-#### GET
-##### Responses
-
-| Code | Description | Schema |
-| ---- | ----------- | ------ |
-| 200 | Success | [TextContentResponse](#textcontentresponse) |
-
 #### POST
 ##### Parameters
 
@@ -4678,9 +4671,9 @@ Initialize dataset with documents
 
 ##### Responses
 
-| Code | Description |
-| ---- | ----------- |
-| 200 | Success |
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Success | [IndexingEstimate](#indexingestimate) |
 
 ### /datasets/process-rule
 
@@ -6652,6 +6645,7 @@ Mark a notification as dismissed for the current user.
 | ---- | ---------- | ----------- | -------- | ------ |
 | page_id | path |  | Yes | string |
 | page_type | path |  | Yes | string |
+| credential_id | query | Credential ID | Yes | string |
 
 ##### Responses
 
@@ -6659,29 +6653,21 @@ Mark a notification as dismissed for the current user.
 | ---- | ----------- | ------ |
 | 200 | Success | [TextContentResponse](#textcontentresponse) |
 
-#### POST
+### /notion/pre-import/pages
+
+#### GET
 ##### Parameters
 
 | Name | Located in | Description | Required | Schema |
 | ---- | ---------- | ----------- | -------- | ------ |
-| page_id | path |  | Yes | string |
-| page_type | path |  | Yes | string |
-| payload | body |  | Yes | [NotionEstimatePayload](#notionestimatepayload) |
+| credential_id | query | Credential ID | Yes | string |
+| dataset_id | query | Dataset ID | No | string |
 
 ##### Responses
 
-| Code | Description |
-| ---- | ----------- |
-| 200 | Success |
-
-### /notion/pre-import/pages
-
-#### GET
-##### Responses
-
-| Code | Description |
-| ---- | ----------- |
-| 200 | Success |
+| Code | Description | Schema |
+| ---- | ----------- | ------ |
+| 200 | Success | [NotionIntegrateInfoListResponse](#notionintegrateinfolistresponse) |
 
 ### /oauth/authorize/{provider}
 
@@ -12456,19 +12442,7 @@ Condition detail
 | ---- | ---- | ----------- | -------- |
 | info_list | [InfoList](#infolist) |  | Yes |
 
-#### DataSourceIntegrate
-
-| Name | Type | Description | Required |
-| ---- | ---- | ----------- | -------- |
-| created_at | object |  | No |
-| disabled | boolean |  | No |
-| id | string |  | No |
-| is_bound | boolean |  | No |
-| link | string |  | No |
-| provider | string |  | No |
-| source_info | [DataSourceIntegrateWorkspace](#datasourceintegrateworkspace) |  | No |
-
-#### DataSourceIntegrateIcon
+#### DataSourceIntegrateIconResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
@@ -12476,31 +12450,43 @@ Condition detail
 | type | string |  | No |
 | url | string |  | No |
 
-#### DataSourceIntegrateList
+#### DataSourceIntegrateListResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| data | [ [DataSourceIntegrate](#datasourceintegrate) ] |  | No |
+| data | [ [DataSourceIntegrateResponse](#datasourceintegrateresponse) ] |  | Yes |
 
-#### DataSourceIntegratePage
+#### DataSourceIntegratePageResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| page_icon | [DataSourceIntegrateIcon](#datasourceintegrateicon) |  | No |
-| page_id | string |  | No |
-| page_name | string |  | No |
-| parent_id | string |  | No |
-| type | string |  | No |
+| page_icon | [DataSourceIntegrateIconResponse](#datasourceintegrateiconresponse) |  | Yes |
+| page_id | string |  | Yes |
+| page_name | string |  | Yes |
+| parent_id | string |  | Yes |
+| type | string |  | Yes |
 
-#### DataSourceIntegrateWorkspace
+#### DataSourceIntegrateResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| pages | [ [DataSourceIntegratePage](#datasourceintegratepage) ] |  | No |
-| total | integer |  | No |
-| workspace_icon | string |  | No |
-| workspace_id | string |  | No |
-| workspace_name | string |  | No |
+| created_at | integer |  | Yes |
+| disabled | boolean |  | Yes |
+| id | string |  | Yes |
+| is_bound | boolean |  | Yes |
+| link | string |  | Yes |
+| provider | string |  | Yes |
+| source_info | [DataSourceIntegrateWorkspaceResponse](#datasourceintegrateworkspaceresponse) |  | Yes |
+
+#### DataSourceIntegrateWorkspaceResponse
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| pages | [ [DataSourceIntegratePageResponse](#datasourceintegratepageresponse) ] |  | Yes |
+| total | integer |  | Yes |
+| workspace_icon | string |  | Yes |
+| workspace_id | string |  | Yes |
+| workspace_name | string |  | Yes |
 
 #### DatasetAndDocumentResponse
 
@@ -13906,6 +13892,14 @@ Request payload for bulk downloading documents as a zip archive.
 | ---- | ---- | ----------- | -------- |
 | include_secret | string | Whether to include secret values in the exported DSL | No |
 
+#### IndexingEstimate
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| preview | [ [PreviewDetail](#previewdetail) ] |  | Yes |
+| qa_preview | [ [QAPreviewDetail](#qapreviewdetail) ] |  | No |
+| total_segments | integer |  | Yes |
+
 #### IndexingEstimatePayload
 
 | Name | Type | Description | Required |
@@ -14462,31 +14456,31 @@ Enum class for model type.
 | pages | [ [NotionPage](#notionpage) ] |  | Yes |
 | workspace_id | string |  | Yes |
 
-#### NotionIntegrateInfoList
+#### NotionIntegrateInfoListResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| notion_info | [ [NotionIntegrateWorkspace](#notionintegrateworkspace) ] |  | No |
+| notion_info | [ [NotionIntegrateWorkspaceResponse](#notionintegrateworkspaceresponse) ] |  | Yes |
 
-#### NotionIntegratePage
+#### NotionIntegratePageResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| is_bound | boolean |  | No |
-| page_icon | [DataSourceIntegrateIcon](#datasourceintegrateicon) |  | No |
-| page_id | string |  | No |
-| page_name | string |  | No |
-| parent_id | string |  | No |
-| type | string |  | No |
+| is_bound | boolean |  | Yes |
+| page_icon | [DataSourceIntegrateIconResponse](#datasourceintegrateiconresponse) |  | Yes |
+| page_id | string |  | Yes |
+| page_name | string |  | Yes |
+| parent_id | string |  | Yes |
+| type | string |  | Yes |
 
-#### NotionIntegrateWorkspace
+#### NotionIntegrateWorkspaceResponse
 
 | Name | Type | Description | Required |
 | ---- | ---- | ----------- | -------- |
-| pages | [ [NotionIntegratePage](#notionintegratepage) ] |  | No |
-| workspace_icon | string |  | No |
-| workspace_id | string |  | No |
-| workspace_name | string |  | No |
+| pages | [ [NotionIntegratePageResponse](#notionintegratepageresponse) ] |  | Yes |
+| workspace_icon | string |  | Yes |
+| workspace_id | string |  | Yes |
+| workspace_name | string |  | Yes |
 
 #### NotionPage
 
@@ -15018,6 +15012,14 @@ Shared permission levels for resources (datasets, credentials, etc.)
 | enabled | boolean |  | Yes |
 | id | string |  | Yes |
 
+#### PreviewDetail
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| child_chunks | [ string ] |  | No |
+| content | string |  | Yes |
+| summary | string |  | No |
+
 #### ProcessRule
 
 | Name | Type | Description | Required |
@@ -15044,6 +15046,13 @@ Shared permission levels for resources (datasets, credentials, etc.)
 | response_mode | string | *Enum:* `"blocking"`, `"streaming"` | No |
 | start_node_id | string |  | Yes |
 
+#### QAPreviewDetail
+
+| Name | Type | Description | Required |
+| ---- | ---- | ----------- | -------- |
+| answer | string |  | Yes |
+| question | string |  | Yes |
+
 #### Quota
 
 | Name | Type | Description | Required |

api/tests/test_containers_integration_tests/controllers/console/datasets/test_data_source.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from datetime import UTC, datetime
 from unittest.mock import MagicMock, PropertyMock, patch
 
 import pytest
@@ -11,12 +12,14 @@ from werkzeug.exceptions import NotFound
 from controllers.console.datasets import data_source
 from controllers.console.datasets.data_source import (
     DataSourceApi,
-    DataSourceNotionApi,
     DataSourceNotionDatasetSyncApi,
     DataSourceNotionDocumentSyncApi,
+    DataSourceNotionIndexingEstimateApi,
     DataSourceNotionListApi,
+    DataSourceNotionPreviewApi,
 )
 from core.rag.index_processor.constant.index_type import IndexStructureType
+from models import DataSourceOauthBinding
 
 
 def unwrap(func):
@@ -59,13 +62,29 @@ class TestDataSourceApi:
         api = DataSourceApi()
         method = unwrap(api.get)
 
-        binding = MagicMock(
-            id="b1",
+        binding = DataSourceOauthBinding(
+            tenant_id="tenant-1",
+            access_token="token",
             provider="notion",
-            created_at="now",
-            disabled=False,
-            source_info={},
+            source_info={
+                "workspace_name": "Workspace",
+                "workspace_id": "workspace-1",
+                "workspace_icon": None,
+                "total": 1,
+                "pages": [
+                    {
+                        "page_id": "page-1",
+                        "page_name": "Page",
+                        "page_icon": {"type": "emoji", "emoji": "P", "url": None},
+                        "parent_id": "parent-1",
+                        "type": "page",
+                    }
+                ],
+            },
         )
+        binding.id = "b1"
+        binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
+        binding.disabled = False
 
         with (
             app.test_request_context("/"),
@@ -77,7 +96,29 @@ class TestDataSourceApi:
             response, status = method(api)
 
         assert status == 200
-        assert response["data"][0]["is_bound"] is True
+        assert response["data"][0] == {
+            "id": "b1",
+            "provider": "notion",
+            "created_at": 1779670923,
+            "is_bound": True,
+            "disabled": False,
+            "source_info": {
+                "workspace_name": "Workspace",
+                "workspace_id": "workspace-1",
+                "workspace_icon": None,
+                "pages": [
+                    {
+                        "page_name": "Page",
+                        "page_id": "page-1",
+                        "page_icon": {"type": "emoji", "url": None, "emoji": "P"},
+                        "parent_id": "parent-1",
+                        "type": "page",
+                    }
+                ],
+                "total": 1,
+            },
+            "link": "http://localhost/console/api/oauth/data-source/notion",
+        }
 
     def test_get_no_bindings(self, app: Flask, patch_tenant):
         api = DataSourceApi()
@@ -322,13 +363,13 @@ class TestDataSourceNotionListApi:
                 method(api)
 
 
-class TestDataSourceNotionApi:
+class TestDataSourceNotionPreviewApi:
     @pytest.fixture
     def app(self, flask_app_with_containers: Flask):
         return flask_app_with_containers
 
     def test_get_preview_success(self, app: Flask, patch_tenant):
-        api = DataSourceNotionApi()
+        api = DataSourceNotionPreviewApi()
         method = unwrap(api.get)
 
         extractor = MagicMock(extract=lambda: [MagicMock(page_content="hello")])
@@ -348,8 +389,14 @@ class TestDataSourceNotionApi:
 
         assert status == 200
 
+
+class TestDataSourceNotionIndexingEstimateApi:
+    @pytest.fixture
+    def app(self, flask_app_with_containers: Flask):
+        return flask_app_with_containers
+
     def test_post_indexing_estimate_success(self, app: Flask, patch_tenant):
-        api = DataSourceNotionApi()
+        api = DataSourceNotionIndexingEstimateApi()
         method = unwrap(api.post)
 
         payload = {

api/tests/unit_tests/controllers/console/datasets/test_data_source.py

@@ -0,0 +1,173 @@
+from __future__ import annotations
+
+import inspect
+from collections.abc import Callable
+from datetime import UTC, datetime
+from typing import cast
+from unittest.mock import MagicMock, PropertyMock, patch
+
+import pytest
+from flask import Flask
+
+from controllers.console.datasets import data_source as module
+from controllers.console.datasets.data_source import DataSourceApi, DataSourceNotionListApi
+from models import DataSourceOauthBinding
+
+ControllerMethod = Callable[..., tuple[dict[str, object], int]]
+
+
+def unwrap(func: object) -> ControllerMethod:
+    return cast(ControllerMethod, inspect.unwrap(cast(Callable[..., object], func)))
+
+
+@pytest.fixture
+def flask_app() -> Flask:
+    app = Flask(__name__)
+    app.config["TESTING"] = True
+    return app
+
+
+@pytest.fixture
+def tenant_context() -> tuple[MagicMock, str]:
+    return MagicMock(id="user-1"), "tenant-1"
+
+
+def test_get_data_source_integrates_serializes_orm_binding(
+    flask_app: Flask, tenant_context: tuple[MagicMock, str]
+) -> None:
+    binding = DataSourceOauthBinding(
+        tenant_id="tenant-1",
+        access_token="token",
+        provider="notion",
+        source_info={
+            "workspace_name": "Workspace",
+            "workspace_id": "workspace-1",
+            "workspace_icon": None,
+            "total": 1,
+            "pages": [
+                {
+                    "page_id": "page-1",
+                    "page_name": "Page",
+                    "page_icon": {"type": "emoji", "emoji": "P", "url": None},
+                    "parent_id": "parent-1",
+                    "type": "page",
+                }
+            ],
+        },
+    )
+    binding.id = "binding-1"
+    binding.created_at = datetime(2026, 5, 25, 1, 2, 3, tzinfo=UTC)
+    binding.disabled = False
+
+    with (
+        flask_app.test_request_context("/"),
+        patch.object(module, "current_account_with_tenant", return_value=tenant_context),
+        patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [binding])),
+    ):
+        response, status = unwrap(DataSourceApi().get)(DataSourceApi())
+
+    assert status == 200
+    assert response == {
+        "data": [
+            {
+                "id": "binding-1",
+                "provider": "notion",
+                "created_at": 1779670923,
+                "is_bound": True,
+                "disabled": False,
+                "source_info": {
+                    "workspace_name": "Workspace",
+                    "workspace_id": "workspace-1",
+                    "workspace_icon": None,
+                    "pages": [
+                        {
+                            "page_name": "Page",
+                            "page_id": "page-1",
+                            "page_icon": {"type": "emoji", "url": None, "emoji": "P"},
+                            "parent_id": "parent-1",
+                            "type": "page",
+                        }
+                    ],
+                    "total": 1,
+                },
+                "link": "http://localhost/console/api/oauth/data-source/notion",
+            }
+        ]
+    }
+
+
+def test_get_data_source_integrates_preserves_empty_list_when_no_binding(
+    flask_app: Flask, tenant_context: tuple[MagicMock, str]
+) -> None:
+    with (
+        flask_app.test_request_context("/"),
+        patch.object(module, "current_account_with_tenant", return_value=tenant_context),
+        patch.object(module.db.session, "scalars", return_value=MagicMock(all=lambda: [])),
+    ):
+        response, status = unwrap(DataSourceApi().get)(DataSourceApi())
+
+    assert status == 200
+    assert response == {"data": []}
+
+
+def test_notion_pre_import_pages_serializes_frontend_list_shape(
+    flask_app: Flask, tenant_context: tuple[MagicMock, str]
+) -> None:
+    page = MagicMock(
+        page_id="page-1",
+        page_name="Page",
+        type="page",
+        parent_id="parent-1",
+        page_icon={"type": "emoji", "emoji": "P", "url": None},
+    )
+    online_document_message = MagicMock(
+        result=[
+            MagicMock(
+                workspace_id="workspace-1",
+                workspace_name="Workspace",
+                workspace_icon=None,
+                pages=[page],
+            )
+        ]
+    )
+    runtime = MagicMock(
+        get_online_document_pages=MagicMock(return_value=iter([online_document_message])),
+        datasource_provider_type=MagicMock(return_value="online_document"),
+    )
+
+    with (
+        flask_app.test_request_context("/?credential_id=credential-1"),
+        patch.object(module, "current_account_with_tenant", return_value=tenant_context),
+        patch.object(
+            module.DatasourceProviderService,
+            "get_datasource_credentials",
+            return_value={"token": "token"},
+        ),
+        patch.object(type(module.db), "engine", new_callable=PropertyMock, return_value=MagicMock()),
+        patch.object(module, "sessionmaker"),
+        patch("core.datasource.datasource_manager.DatasourceManager.get_datasource_runtime", return_value=runtime),
+    ):
+        response, status = unwrap(DataSourceNotionListApi().get)(DataSourceNotionListApi())
+
+    assert status == 200
+    assert response == {
+        "notion_info": [
+            {
+                "workspace_name": "Workspace",
+                "workspace_id": "workspace-1",
+                "workspace_icon": None,
+                "pages": [
+                    {
+                        "page_name": "Page",
+                        "page_id": "page-1",
+                        "page_icon": {"type": "emoji", "url": None, "emoji": "P"},
+                        "parent_id": "parent-1",
+                        "type": "page",
+                        "is_bound": False,
+                    }
+                ],
+            }
+        ]
+    }
+    runtime.get_online_document_pages.assert_called_once()
+    assert runtime.get_online_document_pages.call_args.kwargs["datasource_parameters"] == {}