Cybersecurity-Projects/PROJECTS/beginner/metadata-scrubber-tool/src/services/batch_processor.py

290 lines
9.5 KiB
Python

"""
Batch processing service for metadata operations.
This module provides handler-agnostic batch processing that works with any
MetadataHandler subclass (images now, PDF/Office docs in future).
Supports concurrent processing via ThreadPoolExecutor for efficient handling
of large batches (1000+ files).
"""
import logging
import threading
from collections.abc import Callable, Iterable
from concurrent.futures import ThreadPoolExecutor, as_completed
from dataclasses import dataclass, field
from pathlib import Path
from rich.console import Console
from src.services.metadata_factory import MetadataFactory
log = logging.getLogger("metadata-scrubber")
console = Console()
@dataclass
class FileResult:
"""Result of processing a single file."""
filepath: Path
success: bool
action: str # "scrubbed", "skipped", "dry-run"
output_path: Path | None = None
error: str | None = None
@dataclass
class BatchSummary:
"""Aggregated statistics for batch processing."""
total: int = 0
success: int = 0
skipped: int = 0
failed: int = 0
dry_run: bool = False
output_dir: Path | None = None
results: list[FileResult] = field(default_factory = list)
class BatchProcessor:
"""
Handler-agnostic batch processor for metadata operations.
Works with any MetadataHandler subclass via MetadataFactory.
Supports dry-run mode, automatic duplicate suffix handling,
and concurrent processing via ThreadPoolExecutor.
"""
def __init__(
self,
output_dir: str | None = None,
dry_run: bool = False,
max_workers: int = 4,
):
"""
Initialize the batch processor.
Args:
output_dir: Directory to save processed files. Defaults to "./scrubbed".
dry_run: If True, preview what would be processed without writing files.
max_workers: Maximum number of concurrent worker threads. Defaults to 4.
"""
self.output_dir = Path(output_dir) if output_dir else Path("./scrubbed")
self.dry_run = dry_run
self.max_workers = max_workers
self.results: list[FileResult] = []
# Thread synchronization
self._path_lock = threading.Lock() # Protects unique path generation
self._results_lock = threading.Lock() # Protects results list
def process_file(self, file: Path) -> FileResult:
"""
Process a single file through the read→wipe→save pipeline.
Uses MetadataFactory to get the appropriate handler, so this method
automatically works with any file type that has a registered handler.
Args:
file: Path to the file to process.
Returns:
FileResult with success status and details.
"""
output_path: Path | None = None # Track reserved path for cleanup
try:
# Dry-run mode: just report what would happen
if self.dry_run:
# Verify the file can be handled (will raise if not)
MetadataFactory.get_handler(str(file))
output_path = self._get_unique_output_path(file, reserve = False)
result = FileResult(
filepath = file,
success = True,
action = "dry-run",
output_path = output_path,
)
self._append_result(result)
log.debug(f"[DRY-RUN] Would process: {file}")
return result
# Get handler from factory
handler = MetadataFactory.get_handler(str(file))
# Execute the read → wipe → save pipeline
handler.read()
handler.wipe()
output_path = self._get_unique_output_path(file)
handler.save(str(output_path))
result = FileResult(
filepath = file,
success = True,
action = "scrubbed",
output_path = output_path,
)
self._append_result(result)
if log.isEnabledFor(logging.DEBUG):
# if verbose mode is enabled, log the Info
log.info(f"✅ Scrubbed: {file.name}{output_path}")
return result
except Exception as e:
# Cleanup: remove empty placeholder file if reservation failed
self._cleanup_reserved_path(output_path)
result = FileResult(
filepath = file,
success = False,
action = "skipped",
error = str(e),
)
self._append_result(result)
if log.isEnabledFor(logging.DEBUG):
# if verbose mode is enabled, log the traceback
log.warning(f"⚠️ Skipped {file.name}: {e}")
return result
def process_batch(
self,
files: Iterable[Path],
progress_callback: Callable[[FileResult],
None] | None = None,
) -> list[FileResult]:
"""
Process multiple files concurrently using ThreadPoolExecutor.
Args:
files: Iterable of file paths to process.
progress_callback: Optional callback called after each file completes.
Receives the FileResult for progress updates.
Returns:
List of FileResult objects for all processed files.
"""
file_list = list(files)
if not file_list:
return self.results
# Used ThreadPoolExecutor for I/O-bound concurrent processing
with ThreadPoolExecutor(max_workers = self.max_workers) as executor:
# Submit all files for processing
future_to_file = {
executor.submit(self.process_file,
file): file
for file in file_list
}
# Collect results as they complete
for future in as_completed(future_to_file):
result = future.result()
if progress_callback:
progress_callback(result)
return self.results
def process_batch_sequential(self, files: Iterable[Path]) -> list[FileResult]:
"""
Process files sequentially (legacy behavior for debugging).
Args:
files: Iterable of file paths to process.
Returns:
List of FileResult objects for all processed files.
"""
for file in files:
self.process_file(file)
return self.results
def get_summary(self) -> BatchSummary:
"""
Return aggregated statistics for all processed files.
Returns:
BatchSummary with counts and result details.
"""
with self._results_lock:
results_copy = list(self.results)
summary = BatchSummary(
total = len(results_copy),
success = sum(
1 for r in results_copy if r.success and r.action == "scrubbed"
),
skipped = sum(1 for r in results_copy if not r.success),
dry_run = self.dry_run,
output_dir = self.output_dir,
results = results_copy,
)
# Count dry-run as separate from success for clarity
if self.dry_run:
summary.success = sum(1 for r in results_copy if r.success)
return summary
def _append_result(self, result: FileResult) -> None:
"""Thread-safe append to results list."""
with self._results_lock:
self.results.append(result)
def _cleanup_reserved_path(self, output_path: Path | None) -> None:
"""
Remove empty placeholder file created during path reservation.
Called when processing fails after _get_unique_output_path() reserved
a path via touch(). Only removes files that are empty (0 bytes) to
avoid deleting partially written data.
Args:
output_path: Path that was reserved, or None if not yet reserved.
"""
if output_path is None:
return
try:
if output_path.exists() and output_path.stat().st_size == 0:
output_path.unlink()
log.debug(f"Cleaned up empty placeholder: {output_path}")
except OSError:
# Best effort cleanup - don't fail if we can't delete
pass
def _get_unique_output_path(self, file: Path, reserve: bool = True) -> Path:
"""
Generate unique output path with suffix (_1, _2) if file exists.
Thread-safe: uses lock to prevent race conditions during concurrent processing.
Args:
file: Original file path.
reserve: If True, create placeholder file to reserve the path.
Set to False for dry-run mode.
Returns:
Unique path in output directory that doesn't conflict with existing files.
"""
with self._path_lock:
# creates the destination directory if it doesn't exist
self.output_dir.mkdir(parents = True, exist_ok = True)
base_name = file.stem
extension = file.suffix
output_path = self.output_dir / f"processed_{base_name}{extension}"
# If file exists, add incrementing suffix
counter = 1
while output_path.exists():
output_path = (
self.output_dir / f"processed_{base_name}_{counter}{extension}"
)
counter += 1
# Create empty placeholder to reserve the path (skip in dry-run)
if reserve:
output_path.touch()
return output_path