Fix OneDrive and S3 upload signature - add file_id parameter

- Updated upload_to_onedrive to accept file_id parameter with bind=True - Updated upload_to_s3 to accept file_id parameter with bind=True - Added proper logging with task_id and file_id tracking - Added comprehensive unit tests for both functions - All tests passing (8/8) Fixes #99 and #100 Co-authored-by: christianlouis <361235+christianlouis@users.noreply.github.com>
2026-02-07 17:25:23 +00:00
parent aa57c41556
commit c0e049ae96
3 changed files with 230 additions and 18 deletions
@@ -9,6 +9,9 @@ import urllib.parse
 from app.config import settings
 from app.tasks.retry_config import BaseTaskWithRetry
 from app.celery_app import celery
+from app.utils import log_task_progress
+from app.database import SessionLocal
+from app.models import FileRecord

 logger = logging.getLogger(__name__)

@@ -215,12 +218,24 @@ def upload_large_file(file_path, upload_url):
    # The last response should contain the file metadata
    return response.json()

-@celery.task(base=BaseTaskWithRetry)
-def upload_to_onedrive(file_path: str):
-    """Uploads a file to OneDrive in the configured folder."""
+@celery.task(base=BaseTaskWithRetry, bind=True)
+def upload_to_onedrive(self, file_path: str, file_id: int = None):
+    """
+    Uploads a file to OneDrive in the configured folder.
+    
+    Args:
+        file_path: Path to the file to upload
+        file_id: Optional file ID to associate with logs
+    """
+    task_id = self.request.id
+    logger.info(f"[{task_id}] Starting OneDrive upload: {file_path}")
+    log_task_progress(task_id, "upload_to_onedrive", "in_progress", f"Uploading to OneDrive: {os.path.basename(file_path)}", file_id=file_id)
    
    if not os.path.exists(file_path):
-        raise FileNotFoundError(f"File not found: {file_path}")
+        error_msg = f"File not found: {file_path}"
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_onedrive", "failure", error_msg, file_id=file_id)
+        raise FileNotFoundError(error_msg)

    # Extract filename
    filename = os.path.basename(file_path)
@@ -228,7 +243,8 @@ def upload_to_onedrive(file_path: str):
    # Check if OneDrive settings are configured
    if not settings.onedrive_client_id:
        error_msg = "OneDrive client ID is not configured"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_onedrive", "failure", error_msg, file_id=file_id)
        raise ValueError(error_msg)

    try:
@@ -243,8 +259,9 @@ def upload_to_onedrive(file_path: str):
        
        # Log success
        web_url = result.get("webUrl", "Not available")
-        logger.info(f"Successfully uploaded {filename} to OneDrive at path {settings.onedrive_folder_path}")
-        logger.info(f"File accessible at: {web_url}")
+        logger.info(f"[{task_id}] Successfully uploaded {filename} to OneDrive at path {settings.onedrive_folder_path}")
+        logger.info(f"[{task_id}] File accessible at: {web_url}")
+        log_task_progress(task_id, "upload_to_onedrive", "success", f"Uploaded to OneDrive: {filename}", file_id=file_id)
        
        return {
            "status": "Completed",
@@ -255,5 +272,6 @@ def upload_to_onedrive(file_path: str):
        
    except Exception as e:
        error_msg = f"Failed to upload {filename} to OneDrive: {str(e)}"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_onedrive", "failure", error_msg, file_id=file_id)
        raise Exception(error_msg)
@@ -7,15 +7,30 @@ from botocore.exceptions import ClientError
 from app.config import settings
 from app.tasks.retry_config import BaseTaskWithRetry
 from app.celery_app import celery
+from app.utils import log_task_progress
+from app.database import SessionLocal
+from app.models import FileRecord

 logger = logging.getLogger(__name__)

-@celery.task(base=BaseTaskWithRetry)
-def upload_to_s3(file_path: str):
-    """Uploads a file to Amazon S3 in the configured bucket and folder."""
-
+@celery.task(base=BaseTaskWithRetry, bind=True)
+def upload_to_s3(self, file_path: str, file_id: int = None):
+    """
+    Uploads a file to Amazon S3 in the configured bucket and folder.
+    
+    Args:
+        file_path: Path to the file to upload
+        file_id: Optional file ID to associate with logs
+    """
+    task_id = self.request.id
+    logger.info(f"[{task_id}] Starting S3 upload: {file_path}")
+    log_task_progress(task_id, "upload_to_s3", "in_progress", f"Uploading to S3: {os.path.basename(file_path)}", file_id=file_id)
+    
    if not os.path.exists(file_path):
-        raise FileNotFoundError(f"File not found: {file_path}")
+        error_msg = f"File not found: {file_path}"
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
+        raise FileNotFoundError(error_msg)

    # Extract filename
    filename = os.path.basename(file_path)
@@ -23,12 +38,14 @@ def upload_to_s3(file_path: str):
    # Check if S3 settings are configured
    if not settings.s3_bucket_name:
        error_msg = "S3 bucket name is not configured"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
        raise ValueError(error_msg)

    if not settings.aws_access_key_id or not settings.aws_secret_access_key:
        error_msg = "AWS credentials are not configured"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
        raise ValueError(error_msg)

    try:
@@ -71,7 +88,8 @@ def upload_to_s3(file_path: str):
        # For private files, this is just a reference and won't be accessible directly
        s3_url = f"https://{settings.s3_bucket_name}.s3.{settings.aws_region}.amazonaws.com/{s3_key}"
        
-        logger.info(f"Successfully uploaded {filename} to S3 bucket {settings.s3_bucket_name} at path {s3_key}")
+        logger.info(f"[{task_id}] Successfully uploaded {filename} to S3 bucket {settings.s3_bucket_name} at path {s3_key}")
+        log_task_progress(task_id, "upload_to_s3", "success", f"Uploaded to S3: {filename}", file_id=file_id)
        return {
            "status": "Completed",
            "file": file_path,
@@ -82,10 +100,12 @@ def upload_to_s3(file_path: str):
    
    except ClientError as e:
        error_msg = f"Failed to upload {filename} to S3: {str(e)}"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
        raise Exception(error_msg)
    
    except Exception as e:
        error_msg = f"Error uploading {filename} to S3: {str(e)}"
-        logger.error(error_msg)
+        logger.error(f"[{task_id}] {error_msg}")
+        log_task_progress(task_id, "upload_to_s3", "failure", error_msg, file_id=file_id)
        raise Exception(error_msg)
@@ -0,0 +1,174 @@
+"""
+Tests for upload tasks including OneDrive and S3.
+"""
+import os
+import pytest
+from unittest.mock import Mock, patch, MagicMock
+from app.tasks.upload_to_onedrive import upload_to_onedrive
+from app.tasks.upload_to_s3 import upload_to_s3
+
+
+@pytest.fixture
+def mock_settings():
+    """Mock settings for upload tests."""
+    with patch('app.tasks.upload_to_onedrive.settings') as onedrive_settings, \
+         patch('app.tasks.upload_to_s3.settings') as s3_settings:
+        # OneDrive settings
+        onedrive_settings.onedrive_client_id = "test_client_id"
+        onedrive_settings.onedrive_client_secret = "test_secret"
+        onedrive_settings.onedrive_refresh_token = "test_refresh_token"
+        onedrive_settings.onedrive_folder_path = "test_folder"
+        onedrive_settings.onedrive_tenant_id = "common"
+        
+        # S3 settings
+        s3_settings.s3_bucket_name = "test-bucket"
+        s3_settings.aws_access_key_id = "test_access_key"
+        s3_settings.aws_secret_access_key = "test_secret_key"
+        s3_settings.aws_region = "us-east-1"
+        s3_settings.s3_folder_prefix = "documents"
+        s3_settings.s3_storage_class = "STANDARD"
+        s3_settings.s3_acl = None
+        
+        yield onedrive_settings, s3_settings
+
+
+@pytest.mark.unit
+def test_upload_to_onedrive_accepts_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_onedrive accepts file_id parameter."""
+    with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
+         patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
+         patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
+         patch('app.tasks.upload_to_onedrive.log_task_progress'):
+        
+        # Setup mocks
+        mock_token.return_value = "test_access_token"
+        mock_session.return_value = "https://upload.url"
+        mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
+        
+        # Call with file_id parameter using apply() to simulate task execution
+        # This bypasses Celery and calls the function directly
+        result = upload_to_onedrive.apply(args=[sample_text_file], kwargs={'file_id': 42}).get()
+        
+        assert result["status"] == "Completed"
+        assert result["file_path"] == sample_text_file
+        assert "onedrive_path" in result
+
+
+@pytest.mark.unit
+def test_upload_to_onedrive_without_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_onedrive works without file_id parameter."""
+    with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
+         patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
+         patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
+         patch('app.tasks.upload_to_onedrive.log_task_progress'):
+        
+        # Setup mocks
+        mock_token.return_value = "test_access_token"
+        mock_session.return_value = "https://upload.url"
+        mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
+        
+        # Call without file_id parameter - should use default None
+        result = upload_to_onedrive.apply(args=[sample_text_file]).get()
+        
+        assert result["status"] == "Completed"
+        assert result["file_path"] == sample_text_file
+
+
+@pytest.mark.unit
+def test_upload_to_s3_accepts_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_s3 accepts file_id parameter."""
+    with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
+         patch('app.tasks.upload_to_s3.log_task_progress'):
+        
+        # Setup mock S3 client
+        mock_s3 = Mock()
+        mock_boto_client.return_value = mock_s3
+        mock_s3.upload_file.return_value = None
+        
+        # Call with file_id parameter using apply()
+        result = upload_to_s3.apply(args=[sample_text_file], kwargs={'file_id': 99}).get()
+        
+        assert result["status"] == "Completed"
+        assert result["file"] == sample_text_file
+        assert result["s3_bucket"] == "test-bucket"
+        assert "s3_key" in result
+
+
+@pytest.mark.unit
+def test_upload_to_s3_without_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_s3 works without file_id parameter."""
+    with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
+         patch('app.tasks.upload_to_s3.log_task_progress'):
+        
+        # Setup mock S3 client
+        mock_s3 = Mock()
+        mock_boto_client.return_value = mock_s3
+        mock_s3.upload_file.return_value = None
+        
+        # Call without file_id parameter using apply()
+        result = upload_to_s3.apply(args=[sample_text_file]).get()
+        
+        assert result["status"] == "Completed"
+        assert result["file"] == sample_text_file
+
+
+@pytest.mark.unit
+def test_upload_to_onedrive_file_not_found(mock_settings):
+    """Test that upload_to_onedrive raises error for missing file."""
+    with patch('app.tasks.upload_to_onedrive.log_task_progress'):
+        with pytest.raises(FileNotFoundError):
+            upload_to_onedrive.apply(args=["/nonexistent/file.pdf"], kwargs={'file_id': 1}).get()
+
+
+@pytest.mark.unit
+def test_upload_to_s3_file_not_found(mock_settings):
+    """Test that upload_to_s3 raises error for missing file."""
+    with patch('app.tasks.upload_to_s3.log_task_progress'):
+        with pytest.raises(FileNotFoundError):
+            upload_to_s3.apply(args=["/nonexistent/file.pdf"], kwargs={'file_id': 1}).get()
+
+
+@pytest.mark.unit
+def test_upload_to_onedrive_logs_with_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_onedrive properly logs with file_id."""
+    with patch('app.tasks.upload_to_onedrive.get_onedrive_token') as mock_token, \
+         patch('app.tasks.upload_to_onedrive.create_upload_session') as mock_session, \
+         patch('app.tasks.upload_to_onedrive.upload_large_file') as mock_upload, \
+         patch('app.tasks.upload_to_onedrive.log_task_progress') as mock_log:
+        
+        # Setup mocks
+        mock_token.return_value = "test_access_token"
+        mock_session.return_value = "https://upload.url"
+        mock_upload.return_value = {"webUrl": "https://onedrive.test/file"}
+        
+        # Call with file_id
+        upload_to_onedrive.apply(args=[sample_text_file], kwargs={'file_id': 123}).get()
+        
+        # Verify log_task_progress was called with file_id
+        assert mock_log.called
+        # Check that at least one call included the file_id parameter
+        calls_with_file_id = [call for call in mock_log.call_args_list 
+                             if 'file_id' in call[1] and call[1]['file_id'] == 123]
+        assert len(calls_with_file_id) > 0
+
+
+@pytest.mark.unit
+def test_upload_to_s3_logs_with_file_id(sample_text_file, mock_settings):
+    """Test that upload_to_s3 properly logs with file_id."""
+    with patch('app.tasks.upload_to_s3.boto3.client') as mock_boto_client, \
+         patch('app.tasks.upload_to_s3.log_task_progress') as mock_log:
+        
+        # Setup mock S3 client
+        mock_s3 = Mock()
+        mock_boto_client.return_value = mock_s3
+        mock_s3.upload_file.return_value = None
+        
+        # Call with file_id
+        upload_to_s3.apply(args=[sample_text_file], kwargs={'file_id': 456}).get()
+        
+        # Verify log_task_progress was called with file_id
+        assert mock_log.called
+        # Check that at least one call included the file_id parameter
+        calls_with_file_id = [call for call in mock_log.call_args_list 
+                             if 'file_id' in call[1] and call[1]['file_id'] == 456]
+        assert len(calls_with_file_id) > 0