Zordon/src/api/add_job_helpers.py

#!/usr/bin/env python3
import logging
import os
import shutil
import tempfile
import zipfile
from datetime import datetime

import requests
from tqdm import tqdm
from werkzeug.utils import secure_filename

logger = logging.getLogger()


def handle_uploaded_project_files(request, jobs_list, upload_directory):
    """
    Handles the uploaded project files.

    This method takes a request with a file, a list of jobs, and an upload directory. It checks if the file was uploaded
    directly, if it needs to be downloaded from a URL, or if it's already present on the local file system. It then
    moves the file to the appropriate directory and returns the local path to the file and its name.

    Args:
        request (Request): The request object containing the file.
        jobs_list (list): A list of jobs. The first job in the list is used to get the file's URL and local path.
        upload_directory (str): The directory where the file should be uploaded.

    Raises:
        ValueError: If no valid project paths are found.

    Returns:
        tuple: A tuple containing the local path to the loaded project file and its name.
    """
    # Initialize default values
    loaded_project_local_path = None

    uploaded_project = request.files.get('file', None)
    project_url = jobs_list[0].get('url', None)
    local_path = jobs_list[0].get('local_path', None)
    renderer = jobs_list[0].get('renderer')
    downloaded_file_url = None

    if uploaded_project and uploaded_project.filename:
        referred_name = os.path.basename(uploaded_project.filename)
    elif project_url:
        referred_name, downloaded_file_url = download_project_from_url(project_url)
        if not referred_name:
            raise ValueError(f"Error downloading file from URL: {project_url}")
    elif local_path and os.path.exists(local_path):
        referred_name = os.path.basename(local_path)
    else:
        raise ValueError("Cannot find any valid project paths")

    # Prepare the local filepath
    cleaned_path_name = jobs_list[0].get('name', os.path.splitext(referred_name)[0]).replace(' ', '-')
    job_dir = os.path.join(upload_directory, '-'.join(
        [datetime.now().strftime("%Y.%m.%d_%H.%M.%S"), renderer, cleaned_path_name]))
    os.makedirs(job_dir, exist_ok=True)
    project_source_dir = os.path.join(job_dir, 'source')
    os.makedirs(project_source_dir, exist_ok=True)

    # Move projects to their work directories
    if uploaded_project and uploaded_project.filename:
        loaded_project_local_path = os.path.join(project_source_dir, secure_filename(uploaded_project.filename))
        uploaded_project.save(loaded_project_local_path)
        logger.info(f"Transfer complete for {loaded_project_local_path.split(upload_directory)[-1]}")
    elif project_url:
        loaded_project_local_path = os.path.join(project_source_dir, referred_name)
        shutil.move(downloaded_file_url, loaded_project_local_path)
        logger.info(f"Download complete for {loaded_project_local_path.split(upload_directory)[-1]}")
    elif local_path:
        loaded_project_local_path = os.path.join(project_source_dir, referred_name)
        shutil.copy(local_path, loaded_project_local_path)
        logger.info(f"Import complete for {loaded_project_local_path.split(upload_directory)[-1]}")

    return loaded_project_local_path, referred_name


def download_project_from_url(project_url):
    # This nested function is to handle downloading from a URL
    logger.info(f"Downloading project from url: {project_url}")
    referred_name = os.path.basename(project_url)

    try:
        response = requests.get(project_url, stream=True)
        if response.status_code == 200:
            # Get the total file size from the "Content-Length" header
            file_size = int(response.headers.get("Content-Length", 0))
            # Create a progress bar using tqdm
            progress_bar = tqdm(total=file_size, unit="B", unit_scale=True)
            # Open a file for writing in binary mode
            downloaded_file_url = os.path.join(tempfile.gettempdir(), referred_name)
            with open(downloaded_file_url, "wb") as file:
                for chunk in response.iter_content(chunk_size=1024):
                    if chunk:
                        # Write the chunk to the file
                        file.write(chunk)
                        # Update the progress bar
                        progress_bar.update(len(chunk))
            # Close the progress bar
            progress_bar.close()
            return referred_name, downloaded_file_url
    except Exception as e:
        logger.error(f"Error downloading file: {e}")
    return None, None


def process_zipped_project(zip_path):
    """
    Processes a zipped project.

    This method takes a path to a zip file, extracts its contents, and returns the path to the extracted project file.
    If the zip file contains more than one project file or none, an error is raised.

    Args:
        zip_path (str): The path to the zip file.

    Raises:
        ValueError: If there's more than 1 project file or none in the zip file.

    Returns:
        str: The path to the main project file.
    """
    work_path = os.path.dirname(zip_path)

    try:
        with zipfile.ZipFile(zip_path, 'r') as myzip:
            myzip.extractall(work_path)

        project_files = [x for x in os.listdir(work_path) if os.path.isfile(os.path.join(work_path, x))]
        project_files = [x for x in project_files if '.zip' not in x]

        logger.debug(f"Zip files: {project_files}")

        # supported_exts = RenderWorkerFactory.class_for_name(renderer).engine.supported_extensions
        # if supported_exts:
        #     project_files = [file for file in project_files if any(file.endswith(ext) for ext in supported_exts)]

        # If there's more than 1 project file or none, raise an error
        if len(project_files) != 1:
            raise ValueError(f'Cannot find a valid project file in {os.path.basename(zip_path)}')

        extracted_project_path = os.path.join(work_path, project_files[0])
        logger.info(f"Extracted zip file to {extracted_project_path}")

    except (zipfile.BadZipFile, zipfile.LargeZipFile) as e:
        logger.error(f"Error processing zip file: {e}")
        raise ValueError(f"Error processing zip file: {e}")
    return extracted_project_path