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

#!/usr/bin/env python3

import logging
import os
import time
import urllib.parse

import msal
import requests

from app.celery_app import celery
from app.config import settings
from app.tasks.retry_config import UploadTaskWithRetry
from app.utils import log_task_progress

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")

    # Log more details about the configuration
    tenant = settings.onedrive_tenant_id or "common"
    logger.info(f"Using OneDrive tenant: {tenant}")

    # Define scopes consistently
    scopes = ["https://graph.microsoft.com/.default"]

    # Use refresh token flow (works for both personal and org accounts)
    if settings.onedrive_refresh_token:
        # Use MSAL's ConfidentialClientApplication instead of PublicClientApplication
        app = msal.ConfidentialClientApplication(
            client_id=settings.onedrive_client_id,
            client_credential=settings.onedrive_client_secret,
            authority=f"https://login.microsoftonline.com/{tenant}",
        )

        # Request new token using refresh token
        logger.info("Attempting to acquire token using refresh token")
        token_response = app.acquire_token_by_refresh_token(
            refresh_token=settings.onedrive_refresh_token, scopes=scopes
        )

        if "access_token" not in token_response:
            error = token_response.get("error", "")
            error_desc = token_response.get("error_description", "Unknown error")

            # Log more details about the error
            logger.error("Failed to get access token using refresh token")
            logger.error(f"Error code: {error}")
            logger.error(f"Error description: {error_desc}")

            if error == "invalid_grant":
                logger.error("The refresh token appears to be expired or revoked")
                logger.error("A new authorization flow is required to obtain a fresh token")

            raise ValueError(f"Failed to get access token: {error} - {error_desc}")

        # Check if we received a new refresh token and update it
        if "refresh_token" in token_response:
            new_refresh_token = token_response["refresh_token"]
            logger.info("Received new refresh token from Microsoft")

            # Update the refresh token in memory
            settings.onedrive_refresh_token = new_refresh_token
            logger.info("Updated refresh token in memory")

        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=scopes)

        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 and properly encode for URL
    if folder_path:
        # Remove leading/trailing slashes
        folder_path = folder_path.strip("/")

        # URL encode the path components separately
        path_components = folder_path.split("/")
        encoded_path = "/".join(urllib.parse.quote(component) for component in path_components)

        # Also encode the filename
        encoded_filename = urllib.parse.quote(filename)
        item_path = f"/root:/{encoded_path}/{encoded_filename}:/createUploadSession"
    else:
        # Just encode the filename
        encoded_filename = urllib.parse.quote(filename)
        item_path = f"/root:/{encoded_filename}:/createUploadSession"

    url = f"{base_url}{item_path}"

    # Add required request body (can be empty JSON object)
    request_body = {"item": {"@microsoft.graph.conflictBehavior": "replace"}}

    headers = {"Authorization": f"Bearer {access_token}", "Content-Type": "application/json"}

    logger.info(f"Creating upload session for {filename} at path {folder_path}")

    response = requests.post(url, headers=headers, json=request_body, timeout=settings.http_request_timeout)

    if response.status_code == 200:
        upload_url = response.json().get("uploadUrl")
        logger.info(f"Upload session created successfully for {filename}")
        return upload_url
    else:
        error_msg = f"Failed to create upload session: {response.status_code} - {response.text}"
        logger.error(error_msg)
        logger.error(f"Request URL was: {url}")
        logger.error(f"Request headers: {headers}")
        logger.error(f"Request body: {request_body}")
        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, timeout=settings.http_request_timeout
                    )

                    # 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=UploadTaskWithRetry, bind=True)
def upload_to_onedrive(self, file_path: str, file_id: int = None, folder_override: str = 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):
        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)

    # Check if OneDrive settings are configured
    if not settings.onedrive_client_id:
        error_msg = "OneDrive client ID is not configured"
        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:
        # Get access token
        access_token = get_onedrive_token()

        onedrive_folder = folder_override if folder_override is not None else settings.onedrive_folder_path

        # Create upload session
        upload_url = create_upload_session(filename, onedrive_folder, 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"[{task_id}] Successfully uploaded {filename} to OneDrive at path {onedrive_folder}")
        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",
            "file_path": file_path,
            "onedrive_path": f"{onedrive_folder}/{filename}",
            "web_url": web_url,
        }

    except Exception as e:
        error_msg = f"Failed to upload {filename} to OneDrive: {str(e)}"
        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)