machinepainting-ComfyUI_DriveSendNode/gdrive_upload.py

"""
Google Drive Upload Module
Handles file uploads to Google Drive with integrity verification
"""

import os
import hashlib
from pathlib import Path

from googleapiclient.http import MediaFileUpload
from googleapiclient.errors import HttpError

from .gdrive_auth_manager import get_drive_service, get_folder_id


# MIME type mapping
MIME_TYPES = {
    '.png': 'image/png',
    '.jpg': 'image/jpeg',
    '.jpeg': 'image/jpeg',
    '.webp': 'image/webp',
    '.gif': 'image/gif',
    '.mp4': 'video/mp4',
    '.avi': 'video/x-msvideo',
    '.mov': 'video/quicktime',
    '.enc': 'application/octet-stream',
}


def get_mime_type(file_path):
    """Get MIME type based on file extension."""
    ext = Path(file_path).suffix.lower()
    return MIME_TYPES.get(ext, 'application/octet-stream')


def calculate_sha256(file_path):
    """Calculate SHA256 hash of a file."""
    sha256_hash = hashlib.sha256()
    with open(file_path, 'rb') as f:
        for chunk in iter(lambda: f.read(8192), b''):
            sha256_hash.update(chunk)
    return sha256_hash.hexdigest()


def upload_file(file_path, folder_id=None, service=None, auth_method='oauth', **kwargs):
    """
    Upload a file to Google Drive.
    
    Args:
        file_path: Path to the file to upload
        folder_id: Google Drive folder ID (optional, uses env var or root if not provided)
        service: Existing Drive service object (optional)
        auth_method: Authentication method ('oauth' or 'service_account')
        **kwargs: Additional arguments (for compatibility)
    
    Returns:
        dict with upload result including file ID and verification status
    """
    file_path = Path(file_path)
    
    if not file_path.exists():
        return {'success': False, 'error': f'File not found: {file_path}'}
    
    # Get or create service
    if service is None:
        try:
            service = get_drive_service(auth_method)
        except Exception as e:
            return {'success': False, 'error': f'Authentication failed: {e}'}
    
    # Get folder ID
    if folder_id is None:
        folder_id = get_folder_id()
    
    # Calculate local hash for verification
    local_hash = calculate_sha256(file_path)
    
    # Prepare file metadata
    file_metadata = {
        'name': file_path.name,
        'description': f'Uploaded by ComfyUI DriveSend. SHA256: {local_hash}'
    }
    
    if folder_id:
        file_metadata['parents'] = [folder_id]
    
    # Prepare media
    mime_type = get_mime_type(file_path)
    media = MediaFileUpload(
        str(file_path),
        mimetype=mime_type,
        resumable=True
    )
    
    try:
        # Upload file
        print(f"[DriveSend] Uploading: {file_path.name}")
        
        file = service.files().create(
            body=file_metadata,
            media_body=media,
            fields='id, name, size, md5Checksum',
            supportsAllDrives=True
        ).execute()
        
        file_id = file.get('id')
        remote_size = int(file.get('size', 0))
        local_size = file_path.stat().st_size
        
        # Verify upload
        size_match = remote_size == local_size
        
        if size_match:
            print(f"[DriveSend] ✓ Upload verified: {file_path.name} (ID: {file_id})")
            return {
                'success': True,
                'file_id': file_id,
                'file_name': file.get('name'),
                'size': remote_size,
                'local_hash': local_hash,
                'verified': True
            }
        else:
            print(f"[DriveSend] ⚠ Upload size mismatch: local={local_size}, remote={remote_size}")
            return {
                'success': True,
                'file_id': file_id,
                'file_name': file.get('name'),
                'size': remote_size,
                'local_hash': local_hash,
                'verified': False,
                'warning': 'Size mismatch - file may be corrupted'
            }
    
    except HttpError as e:
        error_msg = str(e)
        
        # Provide helpful error messages
        if 'storageQuotaExceeded' in error_msg:
            print(f"[DriveSend] ✗ Storage quota exceeded!")
            print(f"[DriveSend]   If using service account with personal Gmail, this won't work.")
            print(f"[DriveSend]   Service accounts require Google Workspace (paid).")
            print(f"[DriveSend]   For personal Gmail, use OAuth authentication instead.")
            return {
                'success': False, 
                'error': 'Storage quota exceeded. Service accounts do not work with personal Gmail. Use OAuth instead.'
            }
        
        print(f"[DriveSend] ✗ Google Drive API error: {error_msg}")
        return {'success': False, 'error': f'Google Drive API error: {error_msg}'}
    
    except Exception as e:
        error_msg = f'Upload failed: {e}'
        print(f"[DriveSend] ✗ {error_msg}")
        return {'success': False, 'error': error_msg}


def create_folder(folder_name, parent_folder_id=None, service=None, auth_method='oauth'):
    """
    Create a folder in Google Drive.
    
    Args:
        folder_name: Name of the folder to create
        parent_folder_id: Parent folder ID (optional)
        service: Existing Drive service object (optional)
        auth_method: Authentication method
    
    Returns:
        dict with folder creation result including folder ID
    """
    if service is None:
        try:
            service = get_drive_service(auth_method)
        except Exception as e:
            return {'success': False, 'error': f'Authentication failed: {e}'}
    
    file_metadata = {
        'name': folder_name,
        'mimeType': 'application/vnd.google-apps.folder'
    }
    
    if parent_folder_id:
        file_metadata['parents'] = [parent_folder_id]
    
    try:
        folder = service.files().create(
            body=file_metadata,
            fields='id, name',
            supportsAllDrives=True
        ).execute()
        
        print(f"[DriveSend] Created folder: {folder_name} (ID: {folder.get('id')})")
        return {
            'success': True,
            'folder_id': folder.get('id'),
            'folder_name': folder.get('name')
        }
    
    except HttpError as e:
        return {'success': False, 'error': f'Failed to create folder: {e}'}


def list_files(folder_id=None, service=None, auth_method='oauth', max_results=100):
    """
    List files in a Google Drive folder.
    
    Args:
        folder_id: Folder ID to list (optional, lists root if not provided)
        service: Existing Drive service object (optional)
        auth_method: Authentication method
        max_results: Maximum number of files to return
    
    Returns:
        dict with list of files or error
    """
    if service is None:
        try:
            service = get_drive_service(auth_method)
        except Exception as e:
            return {'success': False, 'error': f'Authentication failed: {e}'}
    
    if folder_id is None:
        folder_id = get_folder_id()
    
    query = f"'{folder_id}' in parents and trashed=false" if folder_id else "trashed=false"
    
    try:
        results = service.files().list(
            q=query,
            fields='files(id, name, size, mimeType, createdTime)',
            pageSize=max_results,
            supportsAllDrives=True,
            includeItemsFromAllDrives=True
        ).execute()
        
        files = results.get('files', [])
        return {
            'success': True,
            'files': files,
            'count': len(files)
        }
    
    except HttpError as e:
        return {'success': False, 'error': f'Failed to list files: {e}'}