gh-christianlouis-docuelevate/app/tasks/upload_to_onedrive.py

#!/usr/bin/env python3

import os
import time
import logging
import requests
import msal
from app.config import settings
from app.tasks.retry_config import BaseTaskWithRetry
from app.celery_app import celery

logger = logging.getLogger(__name__)

def get_onedrive_token():
    """
    Get an access token for Microsoft Graph API using the appropriate flow.
    For personal accounts, uses refresh token flow.
    For organizational accounts, uses client credentials flow if refresh token isn't provided.
    """
    # Check for required settings
    if not settings.onedrive_client_id or not settings.onedrive_client_secret:
        raise ValueError("OneDrive client ID and client secret must be configured")
    
    # Use refresh token flow (works for both personal and org accounts)
    if settings.onedrive_refresh_token:
        # Use MSAL to get token from refresh token
        app = msal.PublicClientApplication(settings.onedrive_client_id)
        
        # Request new token using refresh token
        token_response = app.acquire_token_by_refresh_token(
            refresh_token=settings.onedrive_refresh_token,
            scopes=["https://graph.microsoft.com/Files.ReadWrite"]
        )
        
        if "access_token" not in token_response:
            error = token_response.get("error", "")
            error_desc = token_response.get("error_description", "Unknown error")
            raise ValueError(f"Failed to get access token: {error} - {error_desc}")
            
        return token_response["access_token"]
    
    # No refresh token - try client credentials (only works for org accounts)
    elif settings.onedrive_tenant_id and settings.onedrive_tenant_id != "common":
        authority = f"https://login.microsoftonline.com/{settings.onedrive_tenant_id}"
        app = msal.ConfidentialClientApplication(
            client_id=settings.onedrive_client_id,
            client_credential=settings.onedrive_client_secret,
            authority=authority
        )
        
        # Acquire token for application
        token_response = app.acquire_token_for_client(
            scopes=["https://graph.microsoft.com/.default"]
        )
        
        if "access_token" not in token_response:
            error = token_response.get("error", "")
            error_desc = token_response.get("error_description", "Unknown error")
            raise ValueError(f"Failed to get access token: {error} - {error_desc}")
            
        return token_response["access_token"]
    
    else:
        raise ValueError("For personal Microsoft accounts, ONEDRIVE_REFRESH_TOKEN must be configured")

def create_upload_session(filename, folder_path, access_token):
    """Creates an upload session for large files in Microsoft Graph API."""
    # Construct the API endpoint
    base_url = "https://graph.microsoft.com/v1.0/me/drive"
    
    # Format the folder path correctly
    if folder_path:
        # Remove leading/trailing slashes
        folder_path = folder_path.strip('/')
        # Replace spaces with %20
        folder_path = folder_path.replace(' ', '%20')
        item_path = f"/root:/{folder_path}/{filename}:/createUploadSession"
    else:
        item_path = f"/root:/{filename}:/createUploadSession"
    
    url = f"{base_url}{item_path}"
    
    headers = {
        "Authorization": f"Bearer {access_token}",
        "Content-Type": "application/json"
    }
    
    response = requests.post(url, headers=headers)
    
    if response.status_code == 200:
        return response.json().get("uploadUrl")
    else:
        error_msg = f"Failed to create upload session: {response.status_code} - {response.text}"
        logger.error(error_msg)
        raise Exception(error_msg)

def upload_large_file(file_path, upload_url):
    """
    Upload a large file to OneDrive using the upload session URL.
    Uses chunked upload for reliability.
    """
    # Get file size
    file_size = os.path.getsize(file_path)
    
    # Define chunk size (10 MB)
    chunk_size = 10 * 1024 * 1024
    
    # Open and read file in chunks
    with open(file_path, 'rb') as f:
        # Process file in chunks
        chunk_number = 0
        while True:
            chunk = f.read(chunk_size)
            if not chunk:
                break
                
            # Get the position in the file
            chunk_start = chunk_number * chunk_size
            chunk_end = chunk_start + len(chunk) - 1
            
            # Prepare content range header
            content_range = f"bytes {chunk_start}-{chunk_end}/{file_size}"
            
            # Upload chunk
            headers = {
                "Content-Length": str(len(chunk)),
                "Content-Range": content_range
            }
            
            # Try to upload chunk with retries
            max_retries = 3
            retry_delay = 2  # seconds
            
            for attempt in range(max_retries):
                try:
                    response = requests.put(
                        upload_url,
                        headers=headers,
                        data=chunk
                    )
                    
                    # Check if successful
                    if response.status_code in (201, 202):
                        # 201 = Created (final chunk), 202 = Accepted (more chunks coming)
                        break
                    else:
                        logger.warning(f"Chunk upload failed (attempt {attempt+1}): {response.status_code}")
                        if attempt < max_retries - 1:
                            time.sleep(retry_delay * (attempt + 1))
                except Exception as e:
                    logger.warning(f"Chunk upload error (attempt {attempt+1}): {str(e)}")
                    if attempt < max_retries - 1:
                        time.sleep(retry_delay * (attempt + 1))
            
            if response.status_code not in (201, 202):
                raise Exception(f"Failed to upload chunk after {max_retries} attempts: {response.status_code} - {response.text}")
            
            # Move to next chunk
            chunk_number += 1
            
    # If we get here, all chunks were uploaded successfully
    # 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."""
    
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"File not found: {file_path}")

    # Extract filename
    filename = os.path.basename(file_path)

    # Check if OneDrive settings are configured
    if not settings.onedrive_client_id:
        error_msg = "OneDrive client ID is not configured"
        logger.error(error_msg)
        raise ValueError(error_msg)

    try:
        # Get access token
        access_token = get_onedrive_token()
        
        # Create upload session
        upload_url = create_upload_session(filename, settings.onedrive_folder_path, access_token)
        
        # Upload the file
        result = upload_large_file(file_path, upload_url)
        
        # 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}")
        
        return {
            "status": "Completed",
            "file_path": file_path,
            "onedrive_path": f"{settings.onedrive_folder_path}/{filename}",
            "web_url": web_url
        }
        
    except Exception as e:
        error_msg = f"Failed to upload {filename} to OneDrive: {str(e)}"
        logger.error(error_msg)
        raise Exception(error_msg)