mirror of
https://github.com/blw1138/Zordon.git
synced 2026-06-09 13:39:24 -05:00
Job submission code and API cleanup (#127)
* Refactor add jobs and make add_job api only be one job (instead of a list) * Renamed to JobImportHandler and misc cleanup * Dont bury exceptions in server proxy post_job * Update code to create child jobs in a cleaner manner
This commit is contained in:
@@ -0,0 +1,145 @@
|
||||
#!/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()
|
||||
|
||||
|
||||
class JobImportHandler:
|
||||
|
||||
@classmethod
|
||||
def validate_job_data(cls, new_job_data, upload_directory, uploaded_file=None):
|
||||
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")
|
||||
elif 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 = os.path.splitext(referred_name)[0].replace(' ', '-')
|
||||
job_dir = os.path.join(upload_directory, '-'.join(
|
||||
[datetime.now().strftime("%Y.%m.%d_%H.%M.%S"), engine_name, 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_file and uploaded_file.filename:
|
||||
loaded_project_local_path = os.path.join(project_source_dir, secure_filename(uploaded_file.filename))
|
||||
uploaded_file.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]}")
|
||||
|
||||
if loaded_project_local_path.lower().endswith('.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):
|
||||
# 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
|
||||
|
||||
@staticmethod
|
||||
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(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 {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
|
||||
Reference in New Issue
Block a user