#!/usr/bin/env python3 import logging import os import shutil import tempfile import zipfile from datetime import datetime from pathlib import Path import requests from tqdm import tqdm from werkzeug.utils import secure_filename from src.distributed_job_manager import DistributedJobManager logger = logging.getLogger() class JobImportHandler: """Handles job import operations for rendering projects. This class provides functionality to validate, download, and process job data and project files for the rendering queue system. """ @classmethod def create_jobs_from_processed_data(cls, processed_job_data: dict) -> list[dict]: """ Takes processed job data and creates new jobs Args: processed_job_data: Dictionary containing job information""" loaded_project_local_path = processed_job_data['__loaded_project_local_path'] # prepare child job data job_data_to_create = [] if processed_job_data.get("child_jobs"): for child_job_diffs in processed_job_data["child_jobs"]: processed_child_job_data = processed_job_data.copy() processed_child_job_data.pop("child_jobs") processed_child_job_data.update(child_job_diffs) processed_child_job_data['__use_output_subdir'] = True job_data_to_create.append(processed_child_job_data) else: job_data_to_create.append(processed_job_data) # create the jobs created_jobs = [] for job_data in job_data_to_create: new_job = DistributedJobManager.create_render_job(job_data, loaded_project_local_path) created_jobs.append(new_job) # Save notes to .txt if processed_job_data.get("notes"): parent_dir = Path(loaded_project_local_path).parent.parent notes_name = processed_job_data['name'] + "-notes.txt" with (Path(parent_dir) / notes_name).open("w") as f: f.write(processed_job_data["notes"]) return [x.json() for x in created_jobs] @classmethod def validate_job_data(cls, new_job_data: dict, upload_directory: Path, uploaded_file=None) -> dict: """Validates and prepares job data for import. This method validates the job data dictionary, handles project file acquisition (upload, download, or local copy), and prepares the job directory structure. Args: new_job_data: Dictionary containing job information including 'name', 'engine_name', and optionally 'url' or 'local_path'. upload_directory: Base directory for storing uploaded jobs. uploaded_file: Optional uploaded file object from the request. Returns: The validated job data dictionary with additional metadata. Raises: KeyError: If required fields 'name' or 'engine_name' are missing. FileNotFoundError: If no valid project file can be found. """ loaded_project_local_path = None # check for required keys job_name = new_job_data.get('name') engine_name = new_job_data.get('engine_name') if not job_name: raise KeyError("Missing job name") if not engine_name: raise KeyError("Missing engine name") project_url = new_job_data.get('url', None) local_path = new_job_data.get('local_path', None) downloaded_file_url = None if uploaded_file and uploaded_file.filename: referred_name = os.path.basename(uploaded_file.filename) elif project_url: referred_name, downloaded_file_url = cls.download_project_from_url(project_url) if not referred_name: raise FileNotFoundError(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 FileNotFoundError("Cannot find any valid project paths") # Prepare the local filepath cleaned_path_name = job_name.replace(' ', '-') folder_name = f"{cleaned_path_name}-{engine_name}-{datetime.now().strftime('%Y.%m.%d_%H.%M.%S')}" job_dir = Path(upload_directory) / folder_name os.makedirs(job_dir, exist_ok=True) project_source_dir = Path(job_dir) / 'source' os.makedirs(project_source_dir, exist_ok=True) # Move projects to their work directories if uploaded_file and uploaded_file.filename: # Handle file uploading filename = secure_filename(uploaded_file.filename) loaded_project_local_path = Path(project_source_dir) / filename uploaded_file.save(str(loaded_project_local_path)) logger.info(f"Transfer complete for {loaded_project_local_path.relative_to(upload_directory)}") elif project_url: # Handle downloading project from a URL loaded_project_local_path = Path(project_source_dir) / referred_name shutil.move(str(downloaded_file_url), str(loaded_project_local_path)) logger.info(f"Download complete for {loaded_project_local_path.relative_to(upload_directory)}") elif local_path: # Handle local files loaded_project_local_path = Path(project_source_dir) / referred_name shutil.copy(str(local_path), str(loaded_project_local_path)) logger.info(f"Import complete for {loaded_project_local_path.relative_to(upload_directory)}") if loaded_project_local_path.suffix == ".zip": loaded_project_local_path = cls.process_zipped_project(loaded_project_local_path) new_job_data["__loaded_project_local_path"] = loaded_project_local_path return new_job_data @staticmethod def download_project_from_url(project_url: str): """Downloads a project file from the given URL. Downloads the file from the specified URL to a temporary directory with progress tracking. Returns the filename and temporary path. Args: project_url: The URL to download the project file from. Returns: A tuple of (filename, temp_file_path) if successful, (None, None) if download fails. """ # 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, timeout=300) 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 @staticmethod def process_zipped_project(zip_path: Path) -> 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 (Path): The path to the zip file. Raises: ValueError: If there's more than 1 project file or none in the zip file. Returns: Path: The path to the main project file. """ work_path = zip_path.parent try: with zipfile.ZipFile(zip_path, 'r') as myzip: myzip.extractall(str(work_path)) project_files = [p for p in work_path.iterdir() if p.is_file() and p.suffix.lower() != ".zip"] logger.debug(f"Zip files: {project_files}") # supported_exts = RenderWorkerFactory.class_for_name(engine).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 {zip_path.name}') extracted_project_path = 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}') from e return extracted_project_path