🐛 Fix the issue of commercial edition database creation failure. #772

Author: Phinease

backend/apps/data_process_app.py

@@ -1,6 +1,7 @@
 import logging
 from contextlib import asynccontextmanager
-from fastapi import HTTPException, APIRouter, Form, File, UploadFile
+from fastapi import HTTPException, APIRouter, Form, File, UploadFile, Header
+from typing import Optional
 import base64
 import io
 import time
@@ -35,10 +36,10 @@ async def lifespan(app: APIRouter):
 
 
 @router.post("", response_model=TaskResponse, status_code=201)
-async def create_task(request: TaskRequest):
+async def create_task(request: TaskRequest, authorization: Optional[str] = Header(None)):
     """
     Create a new data processing task (Process → Forward chain)
-    
+
     Returns task ID immediately. Processing happens in the background.
     Tasks are forwarded to Elasticsearch when complete.
     """
@@ -50,7 +51,8 @@ async def create_task(request: TaskRequest):
         source_type=request.source_type,
         chunking_strategy=request.chunking_strategy,
         index_name=request.index_name,
-        original_filename=request.original_filename
+        original_filename=request.original_filename,
+        authorization=authorization
     )
 
     return TaskResponse(task_id=task_result.id)
@@ -114,7 +116,7 @@ async def process_sync_endpoint(
 
 
 @router.post("/batch", response_model=BatchTaskResponse, status_code=201)
-async def create_batch_tasks(request: BatchTaskRequest):
+async def create_batch_tasks(request: BatchTaskRequest, authorization: Optional[str] = Header(None)):
     """
     Create multiple data processing tasks at once (individual Process → Forward chains)
     
@@ -147,7 +149,8 @@ async def create_batch_tasks(request: BatchTaskRequest):
                 source_type=source_type,
                 chunking_strategy=chunking_strategy,
                 index_name=index_name,
-                original_filename=original_filename
+                original_filename=original_filename,
+                authorization=authorization
             )
 
             task_ids.append(task_result.id)

backend/apps/file_management_app.py

@@ -150,7 +150,8 @@ async def process_files(
         files: List[dict] = Body(..., description="List of file details to process, including path_or_url and filename"),
         chunking_strategy: Optional[str] = Body("basic"),
         index_name: str = Body(...),
-        destination: str = Body(...)
+        destination: str = Body(...),
+        authorization: Optional[str] = Header(None)
 ):
     """
     Trigger data processing for a list of uploaded files.
@@ -162,7 +163,8 @@ async def process_files(
     process_params = ProcessParams(
         chunking_strategy=chunking_strategy,
         source_type=destination,
-        index_name=index_name
+        index_name=index_name,
+        authorization=authorization
     )
 
     process_result = await trigger_data_process(files, process_params)

backend/consts/model.py

@@ -191,6 +191,7 @@ class ProcessParams(BaseModel):
     chunking_strategy: Optional[str] = "basic"
     source_type: str
     index_name: str
+    authorization: Optional[str] = None
 
 
 class OpinionRequest(BaseModel):

backend/data_process/tasks.py

@@ -280,7 +280,7 @@ def process(self, source: str, source_type: str,
 @app.task(bind=True, base=LoggingTask, name='data_process.tasks.forward', queue='forward_q')
 def forward(self, processed_data: Dict, index_name: str, 
             source: str, source_type: str = 'minio', 
-            original_filename: Optional[str] = None) -> Dict:
+            original_filename: Optional[str] = None, authorization: Optional[str] = None) -> Dict:
     """
     Vectorize and store processed chunks in Elasticsearch
     
@@ -290,6 +290,7 @@ def forward(self, processed_data: Dict, index_name: str,
         source: Original source path (for metadata)
         source_type: The type of the source("local", "minio")
         original_filename: The original name of the file
+        authorization: Authorization header for API calls
         
     Returns:
         Dict containing storage results and metadata
@@ -381,6 +382,8 @@ async def index_documents():
             route_url = f"/indices/{original_index_name}/documents"
             full_url = elasticsearch_url + route_url
             headers = {"Content-Type": "application/json"}
+            if authorization:
+                headers["Authorization"] = authorization
 
             max_retries = 5
             retry_delay = 5
@@ -540,8 +543,9 @@ async def index_documents():
         raise
 
 @app.task(bind=True, base=LoggingTask, name='data_process.tasks.process_and_forward')
-def process_and_forward(self, source: str, source_type: str, 
-                        chunking_strategy: str, index_name: Optional[str] = None, original_filename: Optional[str] = None) -> str:
+def process_and_forward(self, source: str, source_type: str,
+                        chunking_strategy: str, index_name: Optional[str] = None,
+                        original_filename: Optional[str] = None, authorization: Optional[str] = None) -> str:
     """
     Combined task that chains processing and forwarding
     
@@ -553,6 +557,7 @@ def process_and_forward(self, source: str, source_type: str,
         chunking_strategy: Strategy for chunking the document
         index_name: Name of the index to store documents
         original_filename: The original name of the file
+        authorization: Authorization header for API calls
         **params: Additional parameters
         
     Returns:
@@ -573,7 +578,8 @@ def process_and_forward(self, source: str, source_type: str,
             index_name=index_name,
             source=source,
             source_type=source_type,
-            original_filename=original_filename
+            original_filename=original_filename,
+            authorization=authorization
         ).set(queue='forward_q')
     )
 

backend/utils/file_management_utils.py

@@ -33,6 +33,11 @@ async def trigger_data_process(files: List[dict], process_params: ProcessParams)
         if not files:
             return None
 
+        # Build headers with authorization
+        headers = {
+            "Authorization": f"Bearer {process_params.authorization}"
+        }
+
         # Build source data list
         if len(files) == 1:
             # Single file request
@@ -47,7 +52,7 @@ async def trigger_data_process(files: List[dict], process_params: ProcessParams)
 
             try:
                 async with httpx.AsyncClient() as client:
-                    response = await client.post(f"{DATA_PROCESS_SERVICE}/tasks", json=payload, timeout=30.0)
+                    response = await client.post(f"{DATA_PROCESS_SERVICE}/tasks", headers=headers, json=payload, timeout=30.0)
 
                 if response.status_code == 201:
                     return response.json()
@@ -79,7 +84,7 @@ async def trigger_data_process(files: List[dict], process_params: ProcessParams)
 
             try:
                 async with httpx.AsyncClient() as client:
-                    response = await client.post(f"{DATA_PROCESS_SERVICE}/tasks/batch", json=payload, timeout=30.0)
+                    response = await client.post(f"{DATA_PROCESS_SERVICE}/tasks/batch", headers=headers, json=payload, timeout=30.0)
 
                 if response.status_code == 201:
                     return response.json()

test/backend/app/test_data_process_app.py

@@ -8,8 +8,8 @@
 from unittest.mock import patch, MagicMock, AsyncMock
 import os
 import sys
-from fastapi import FastAPI, HTTPException, Form, Body
-from typing import Dict, Any
+from fastapi import FastAPI, HTTPException, Form, Body, Header
+from typing import Dict, Any, Optional
 from fastapi.testclient import TestClient
 from PIL import Image
 import pytest
@@ -127,9 +127,16 @@ def setup_mock_routes(self):
         """
         @self.app.post("/tasks")
         @pytest.mark.asyncio
-        async def create_task(request: Dict[str, Any] = Body(None)):
-            # Simulate task creation
-            task_result = self.process_and_forward.delay.return_value
+        async def create_task(request: Dict[str, Any] = Body(None), authorization: Optional[str] = Header(None)):
+            # Simulate task creation by actually calling the mock
+            task_result = self.process_and_forward.delay(
+                source=request.get("source"),
+                source_type=request.get("source_type"),
+                chunking_strategy=request.get("chunking_strategy"),
+                index_name=request.get("index_name"),
+                original_filename=request.get("original_filename"),
+                authorization=authorization
+            )
             return {"task_id": task_result.id}
 
         @self.app.post("/tasks/process")
@@ -302,16 +309,67 @@ def test_create_task(self):
             "source_type": "file",
             "chunking_strategy": "basic",
             "index_name": self.index_name,
+            "original_filename": "test-file.pdf",
             "additional_params": {"param1": "value1"}
         }
 
-        # Execute request
+        # Execute request with authorization header
+        response = self.client.post(
+            "/tasks", 
+            json=request_data,
+            headers={"Authorization": "Bearer test-token"}
+        )
+
+        # Assert expectations
+        self.assertEqual(response.status_code, 200)  # Mock route defaults to 200
+        self.assertEqual(response.json(), {"task_id": self.task_id})
+
+        # Verify that process_and_forward.delay was called with authorization parameter
+        self.process_and_forward.delay.assert_called_once_with(
+            source=self.source,
+            source_type="file",
+            chunking_strategy="basic",
+            index_name=self.index_name,
+            original_filename="test-file.pdf",
+            authorization="Bearer test-token"
+        )
+
+    def test_create_task_without_authorization(self):
+        """
+        Test creating a new task without authorization header.
+        Verifies that the endpoint works when authorization is not provided.
+        """
+        # Set up mock
+        mock_task = MagicMock()
+        mock_task.id = self.task_id
+        self.process_and_forward.delay.return_value = mock_task
+
+        # Test data
+        request_data = {
+            "source": self.source,
+            "source_type": "file",
+            "chunking_strategy": "basic",
+            "index_name": self.index_name,
+            "original_filename": "test-file.pdf"
+        }
+
+        # Execute request without authorization header
         response = self.client.post("/tasks", json=request_data)
 
         # Assert expectations
-        self.assertEqual(response.status_code, 200)  # FastAPI test client defaults to 200
+        self.assertEqual(response.status_code, 200)  # Mock route defaults to 200
         self.assertEqual(response.json(), {"task_id": self.task_id})
 
+        # Verify that process_and_forward.delay was called with None authorization
+        self.process_and_forward.delay.assert_called_once_with(
+            source=self.source,
+            source_type="file",
+            chunking_strategy="basic",
+            index_name=self.index_name,
+            original_filename="test-file.pdf",
+            authorization=None
+        )
+
     def test_process_sync_endpoint_success(self):
         """
         Test synchronous processing endpoint success case.

test/backend/app/test_file_management_app.py

@@ -122,20 +122,45 @@ async def test_upload_files_success(mock_files):
         assert len(response.json()["uploaded_filenames"]) == 2
 
 @pytest.mark.asyncio
-async def test_upload_files_processing_error(mock_files):
-    with patch("backend.apps.file_management_app.save_upload_file") as mock_save, \
-         patch("backend.apps.file_management_app.trigger_data_process") as mock_trigger, \
-         patch("backend.apps.file_management_app.upload_dir", Path("test_uploads")):
+async def test_process_files_success(mock_files):
+    with patch("backend.apps.file_management_app.trigger_data_process") as mock_trigger:
 
-        # Configure mocks
-        mock_save.return_value = asyncio.Future()
-        mock_save.return_value.set_result(True)
+        # Configure mock for successful processing - return result directly
+        mock_trigger.return_value = {
+            "task_id": "task_123",
+            "status": "processing"
+        }
+        
+        # Create test client
+        with TestClient(app) as client:
+            response = client.post(
+                "/file/process",
+                json={
+                    "files": [
+                        {"path_or_url": "/test/path/test.txt", "filename": "test.txt"}
+                    ],
+                    "chunking_strategy": "basic",
+                    "index_name": "test_index",
+                    "destination": "local"
+                },
+                headers={"authorization": "Bearer test_token"}
+            )
 
-        # Use AsyncMock to properly handle async function mocking
-        mock_trigger.side_effect = AsyncMock(return_value={
+        # Assertions
+        assert response.status_code == 201
+        assert "message" in response.json()
+        assert "process_tasks" in response.json()
+        assert response.json()["message"] == "Files processing triggered successfully"
+
+@pytest.mark.asyncio
+async def test_process_files_processing_error(mock_files):
+    with patch("backend.apps.file_management_app.trigger_data_process") as mock_trigger:
+        
+        # Configure mock for processing error - return result directly
+        mock_trigger.return_value = {
             "status": "error",
             "message": "Processing failed"
-        })
+        }
 
         # Create test client
         with TestClient(app) as client:
@@ -145,15 +170,44 @@ async def test_upload_files_processing_error(mock_files):
                     "files": [
                         {"path_or_url": "/test/path/test.txt", "filename": "test.txt"}
                     ],
-                    "chunking_strategy": "paragraph", 
+                    "chunking_strategy": "basic",
                     "index_name": "test_index",
                     "destination": "local"
-                }
+                },
+                headers={"authorization": "Bearer test_token"}
+            )
+        
+        # Assertions
+        assert response.status_code == 500
+        assert "error" in response.json()
+        assert response.json()["error"] == "Processing failed"
+
+@pytest.mark.asyncio
+async def test_process_files_none_result(mock_files):
+    with patch("backend.apps.file_management_app.trigger_data_process") as mock_trigger:
+        
+        # Configure mock to return None directly
+        mock_trigger.return_value = None
+        
+        # Create test client
+        with TestClient(app) as client:
+            response = client.post(
+                "/file/process",
+                json={
+                    "files": [
+                        {"path_or_url": "/test/path/test.txt", "filename": "test.txt"}
+                    ],
+                    "chunking_strategy": "by_title",
+                    "index_name": "test_index",
+                    "destination": "minio"
+                },
+                headers={"authorization": "Bearer test_token"}
             )
 
         # Assertions
         assert response.status_code == 500
         assert "error" in response.json()
+        assert response.json()["error"] == "Data process service failed"
 
 @pytest.mark.asyncio
 async def test_upload_files_no_files(mock_files):