diff --git a/backend/apps/files/tasks.py b/backend/apps/files/tasks.py index 7444eb6..9050964 100644 --- a/backend/apps/files/tasks.py +++ b/backend/apps/files/tasks.py @@ -1,7 +1,8 @@ import os import io -import hashlib import subprocess +from pathlib import Path +from typing import Optional, Tuple from django.db import transaction @@ -11,12 +12,91 @@ from PIL import Image as PillowImage import blurhash from .models import PostFileModel +from utils.hash import compute_file_hash_blake3, compute_md5_hash, compute_blur_hash -from utils.hash import compute_file_hash_blake3, compute_blur_hash + +class ThumbnailGenerationError(Exception): + """Custom exception for thumbnail generation errors.""" + + pass + + +def _setup_output_path(file_hash: str, prefix: str = "thumbnail") -> Tuple[str, str]: + """ + Set up the output directory and generate a unique filename. + + Args: + file_hash (str): Hash to use in the filename + prefix (str): Prefix for the filename + + Returns: + Tuple[str, str]: Output directory path and full file path + """ + output_dir = "/tmp/thumbgen/" + os.makedirs(output_dir, exist_ok=True) + + filename = f"{prefix}_{file_hash}.png" + filepath = os.path.join(output_dir, filename) + + return output_dir, filepath + + +def _update_file_model( + file_model: PostFileModel, thumbnail_path: str, thumbnail_filename: str +) -> None: + """ + Update the PostFileModel with the new thumbnail and related hashes. + + Args: + file_model (PostFileModel): The model to update + thumbnail_path (str): Path to the generated thumbnail + thumbnail_filename (str): Filename for the saved thumbnail + """ + # Compute the hash for the generated thumbnail + thumbnail_hash_blake3 = compute_file_hash_blake3(thumbnail_path) + + # Update the PostFileModel's thumbnail field with the new file + with open(thumbnail_path, "rb") as file: + file_model.thumbnail.save(thumbnail_filename, file) + + # Set the thumbnail hash + file_model.thumbnail_hash_blake3 = thumbnail_hash_blake3 + + # Generate and set the blur hash for the thumbnail + file_model.thumbnail_blur_hash = compute_blur_hash(thumbnail_path) + + # Save the model + file_model.save() + + +def _handle_task_error(e: Exception, file_id: int, process_name: str): + """ + Handle errors in thumbnail generation tasks. + + Args: + e (Exception): The exception that occurred + file_id (int): ID of the file being processed + process_name (str): Name of the process for error reporting + + Raises: + Retry: To trigger Celery retry mechanism + """ + error_message = f"Error in {process_name} for file {file_id}: {str(e)}" + print(error_message) + raise Retry(exc=e) @shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5) -def generate_blur_hash_PostFile(file_id): +def generate_blur_hash_PostFile(file_id: int) -> str: + """ + Generate and save a blur hash for an image stored in PostFileModel. + + Args: + file_id (int): ID of the PostFileModel instance + + Returns: + str: Success message + """ try: with transaction.atomic(): img = PostFileModel.objects.select_for_update().get(id=file_id) @@ -28,47 +108,56 @@ def generate_blur_hash_PostFile(file_id): img.refresh_from_db() img.blur_hash = blurhash_string img.save() - return f"Successfully generated blur hash for file {file_id}" # Success message + + return f"Successfully generated blur hash for file {file_id}" + except Exception as e: - error_message = f"Error generating blur hash for file {file_id}: {e}" - print(error_message) - raise Retry(exc=e) # Retry on exception - return error_message # This ensures the error message is stored in the results + _handle_task_error(e, file_id, "blur hash generation") @shared_task(autoretry_for=(Exception,), retry_backoff=True, max_retries=5) -def generate_md5_hash_PostFile(file_id): +def generate_md5_hash_PostFile(file_id: int) -> str: + """ + Generate and save an MD5 hash for a file stored in PostFileModel. + + Args: + file_id (int): ID of the PostFileModel instance + + Returns: + str: Success message + """ try: with transaction.atomic(): pstfile = PostFileModel.objects.select_for_update().get(id=file_id) - hash_md5 = hashlib.md5() - with open(pstfile.file.path, "rb") as f: - for chunk in iter(lambda: f.read(4096), b""): - hash_md5.update(chunk) - md5_hash = hash_md5.hexdigest() + # Compute the MD5 hash + md5_hash = compute_md5_hash(pstfile.file.path) + + # Save the computed hash pstfile.refresh_from_db() pstfile.hash_md5 = md5_hash pstfile.save() - return f"Successfully generated MD5 hash for file {file_id}" # Success message + + return f"Successfully generated MD5 hash for file {file_id}" + except Exception as e: - error_message = f"Error generating MD5 hash for file {file_id}: {e}" - print(error_message) - raise Retry(exc=e) # Retry on exception - return error_message # This ensures the error message is stored in the results + _handle_task_error(e, file_id, "MD5 hash generation") @shared_task(name="generate_video_thumbnail") def generate_video_thumbnail( - file_id: int, size: int = 0, timestamp=None, movie_strip: bool = False -): + file_id: int, + size: int = 0, + timestamp: Optional[float] = None, + movie_strip: bool = False, +) -> str: """ Generate video thumbnails using ffmpegthumbnailer and update the PostFileModel instance. Args: file_id (int): ID of the PostFileModel instance - size (int): Desired thumbnail width or height defulats to video size - timestamp (float): Timestamp(s) in seconds where the thumbnail should be extracted + size (int): Desired thumbnail width or height (defaults to video size) + timestamp (float): Timestamp in seconds where the thumbnail should be extracted movie_strip (bool): Create a movie strip overlay Returns: @@ -84,14 +173,13 @@ def generate_video_thumbnail( video_path = pstfile.file.path - # Create output directory if it doesn't exist - output_dir = "/tmp/thumbgen/" - os.makedirs(output_dir, exist_ok=True) - - thumbnail_filename = f"thumbnail_{pstfile.hash_blake3}.png" - - thumbnail_file_path = os.path.join(output_dir, thumbnail_filename) + # Setup output path + _, thumbnail_file_path = _setup_output_path( + pstfile.hash_blake3, "video_thumbnail" + ) + thumbnail_filename = Path(thumbnail_file_path).name + # Build command cmd = [ "ffmpegthumbnailer", "-i", @@ -110,25 +198,24 @@ def generate_video_thumbnail( if timestamp is not None: cmd.extend(["-t", f"{timestamp}"]) + # Execute command subprocess.run(cmd, check=True) - thumbnail_hash_blake3 = compute_file_hash_blake3(thumbnail_file_path) + # Update model with new thumbnail + _update_file_model(pstfile, thumbnail_file_path, thumbnail_filename) - # Update the PostFileModel's thumbnail field with the new file - with open(thumbnail_file_path, "rb") as file: - pstfile.thumbnail.save(thumbnail_filename, file) + # Clean up temporary file + os.remove(thumbnail_file_path) - pstfile.thumbnail_hash_blake3 = thumbnail_hash_blake3 + return f"Video thumbnail generated successfully for file {file_id}" - pstfile.thumbnail_blur_hash = compute_blur_hash(thumbnail_file_path) + except subprocess.CalledProcessError as e: + _handle_task_error(e, file_id, "video thumbnail generation") + except Exception as e: + _handle_task_error(e, file_id, "video thumbnail generation") - pstfile.save() os.remove(thumbnail_file_path) - return f"Thumbnail generated and saved to {thumbnail_file_path}" - except subprocess.CalledProcessError as e: - return f"Error generating thumbnail: {str(e)}" except Exception as e: - return f"Unexpected error: {str(e)}"