Refator: tasks for less duplicated code
This commit is contained in:
parent
11ce69eaca
commit
c42c3b8e4a
GPG key ID:
583642324A1D2070
1 changed files with 128 additions and 41 deletions
|
|
@ -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)}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue