video-query/backend/chunked_upload.py

import os
import uuid
import json
from flask import Blueprint, request, jsonify, current_app
from werkzeug.utils import secure_filename
import logging
from auth import require_auth

logger = logging.getLogger('video_query')

# Create blueprint for handling chunked uploads
chunked_upload_bp = Blueprint('chunked_upload', __name__)

# Track upload sessions
active_uploads = {}

@chunked_upload_bp.route('/api/init-upload', methods=['POST'])
@require_auth
def init_upload():
    """Initialize a new chunked upload session"""
    if not request.is_json:
        return jsonify({"success": False, "message": "Request must be JSON"}), 400
    
    data = request.get_json()
    filename = data.get('filename')
    total_size = data.get('size')
    
    if not filename or not total_size:
        return jsonify({"success": False, "message": "Filename and size are required"}), 400
    
    # Generate a unique ID for this upload
    upload_id = str(uuid.uuid4())
    
    # Generate a unique filename
    original_filename = secure_filename(filename)
    unique_filename = f"{upload_id}_{original_filename}"
    upload_path = os.path.join(current_app.config['UPLOAD_FOLDER'], unique_filename)
    
    # Create/ensure the upload folder exists
    os.makedirs(current_app.config['UPLOAD_FOLDER'], exist_ok=True)
    
    # Initialize an empty file
    with open(upload_path, 'wb') as f:
        pass
    
    # Store upload info
    active_uploads[upload_id] = {
        'path': upload_path,
        'original_filename': original_filename,
        'total_size': total_size,
        'uploaded_size': 0,
        'complete': False
    }
    
    logger.info(f"Initialized upload session {upload_id} for {filename} ({total_size} bytes)")
    
    return jsonify({
        "success": True,
        "upload_id": upload_id
    })

@chunked_upload_bp.route('/api/upload-chunk/<upload_id>', methods=['POST'])
@require_auth
def upload_chunk(upload_id):
    """Handle a chunk of file data"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400
    
    upload = active_uploads[upload_id]
    
    # Check if we have the file chunk
    if 'chunk' not in request.files:
        return jsonify({"success": False, "message": "No chunk in request"}), 400
    
    chunk = request.files['chunk']
    chunk_number = request.form.get('chunk_number', 0, type=int)
    
    # Update the file with this chunk
    with open(upload['path'], 'ab') as f:
        chunk_data = chunk.read()
        chunk_size = len(chunk_data)
        f.write(chunk_data)
    
    # Update upload state
    upload['uploaded_size'] += chunk_size
    progress = min(100, round((upload['uploaded_size'] / upload['total_size']) * 100))
    
    logger.info(f"Received chunk {chunk_number} for upload {upload_id} - Progress: {progress}%")
    
    # Check if upload is complete
    if upload['uploaded_size'] >= upload['total_size']:
        upload['complete'] = True
        logger.info(f"Upload complete for {upload_id}")
    
    return jsonify({
        "success": True,
        "upload_id": upload_id,
        "chunk_number": chunk_number,
        "bytes_received": chunk_size,
        "total_received": upload['uploaded_size'],
        "progress": progress,
        "complete": upload['complete']
    })

@chunked_upload_bp.route('/api/complete-upload/<upload_id>', methods=['POST'])
@require_auth
def complete_upload(upload_id):
    """Mark an upload as complete and return the file path for processing"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400
    
    upload = active_uploads[upload_id]
    
    # Verify the upload is actually complete
    if not upload['complete']:
        # Check the file size
        if os.path.exists(upload['path']):
            actual_size = os.path.getsize(upload['path'])
            if actual_size >= upload['total_size']:
                upload['complete'] = True
                upload['uploaded_size'] = actual_size
                logger.info(f"Manually verified upload complete for {upload_id}")
            else:
                logger.warning(f"Upload not complete for {upload_id}: {actual_size}/{upload['total_size']} bytes")
                return jsonify({
                    "success": False, 
                    "message": f"Upload not complete: {actual_size}/{upload['total_size']} bytes"
                }), 400
        else:
            logger.error(f"Upload file not found for {upload_id}")
            return jsonify({"success": False, "message": "Upload file not found"}), 500
    
    logger.info(f"Upload {upload_id} marked as complete: {upload['original_filename']}")
    
    return jsonify({
        "success": True,
        "upload_id": upload_id,
        "file_path": upload['path'],
        "filename": upload['original_filename'],
        "size": upload['uploaded_size']
    })

@chunked_upload_bp.route('/api/cancel-upload/<upload_id>', methods=['POST'])
@require_auth
def cancel_upload(upload_id):
    """Cancel an upload and delete the partial file"""
    if upload_id not in active_uploads:
        return jsonify({"success": False, "message": "Invalid upload ID"}), 400
    
    upload = active_uploads[upload_id]
    
    # Delete the partial file
    if os.path.exists(upload['path']):
        try:
            os.remove(upload['path'])
            logger.info(f"Deleted partial upload for {upload_id}")
        except Exception as e:
            logger.error(f"Error deleting partial upload for {upload_id}: {str(e)}")
    
    # Remove from active uploads
    del active_uploads[upload_id]
    
    return jsonify({
        "success": True,
        "message": "Upload cancelled"
    })