diff --git a/app/tasks/process_document.py b/app/tasks/process_document.py index 05cfa74b..8b9f5921 100644 --- a/app/tasks/process_document.py +++ b/app/tasks/process_document.py @@ -9,7 +9,9 @@ import PyPDF2 # Replace fitz with PyPDF2 from app.config import settings from app.tasks.retry_config import BaseTaskWithRetry -from app.tasks.process_with_azure_document_intelligence import process_with_azure_document_intelligence +from app.tasks.process_with_azure_document_intelligence import ( + process_with_azure_document_intelligence, +) from app.tasks.extract_metadata_with_gpt import extract_metadata_with_gpt from app.celery_app import celery from app.database import SessionLocal @@ -33,7 +35,12 @@ def process_document(self, original_local_file: str): """ task_id = self.request.id logger.info(f"[{task_id}] Starting document processing: {original_local_file}") - log_task_progress(task_id, "process_document", "in_progress", f"Processing file: {original_local_file}") + log_task_progress( + task_id, + "process_document", + "in_progress", + f"Processing file: {original_local_file}", + ) if not os.path.exists(original_local_file): logger.error(f"[{task_id}] File {original_local_file} not found.") @@ -49,25 +56,42 @@ def process_document(self, original_local_file: str): mime_type, _ = mimetypes.guess_type(original_local_file) if not mime_type: mime_type = "application/octet-stream" - - logger.info(f"[{task_id}] File hash: {filehash[:10]}..., Size: {file_size} bytes, MIME: {mime_type}") - log_task_progress(task_id, "hash_file", "success", f"Hash: {filehash[:10]}..., Size: {file_size} bytes") + + logger.info( + f"[{task_id}] File hash: {filehash[:10]}..., Size: {file_size} bytes, MIME: {mime_type}" + ) + log_task_progress( + task_id, + "hash_file", + "success", + f"Hash: {filehash[:10]}..., Size: {file_size} bytes", + ) # Acquire DB session in the task with SessionLocal() as db: existing = db.query(FileRecord).filter_by(filehash=filehash).one_or_none() if existing: - logger.info(f"[{task_id}] Duplicate file detected (hash={filehash[:10]}...) Skipping processing.") - log_task_progress(task_id, "process_document", "success", "Duplicate file detected, skipping", file_id=existing.id) + logger.info( + f"[{task_id}] Duplicate file detected (hash={filehash[:10]}...) Skipping processing." + ) + log_task_progress( + task_id, + "process_document", + "success", + "Duplicate file detected, skipping", + file_id=existing.id, + ) return { "status": "duplicate_file", "file_id": existing.id, - "detail": "File already processed." + "detail": "File already processed.", } # Not a duplicate -> insert a new record logger.info(f"[{task_id}] Creating new file record in database") - log_task_progress(task_id, "create_file_record", "in_progress", "Creating file record") + log_task_progress( + task_id, "create_file_record", "in_progress", "Creating file record" + ) new_record = FileRecord( filehash=filehash, original_filename=original_filename, @@ -79,7 +103,13 @@ def process_document(self, original_local_file: str): db.commit() db.refresh(new_record) logger.info(f"[{task_id}] File record created with ID: {new_record.id}") - log_task_progress(task_id, "create_file_record", "success", f"File record ID: {new_record.id}", file_id=new_record.id) + log_task_progress( + task_id, + "create_file_record", + "success", + f"File record ID: {new_record.id}", + file_id=new_record.id, + ) # 1. Generate a UUID-based filename and place it in /workdir/tmp file_ext = os.path.splitext(original_local_file)[1] @@ -91,19 +121,40 @@ def process_document(self, original_local_file: str): new_local_path = os.path.join(tmp_dir, new_filename) logger.info(f"[{task_id}] Copying file to: {new_local_path}") - log_task_progress(task_id, "copy_file", "in_progress", f"Copying file to {new_filename}", file_id=new_record.id) + log_task_progress( + task_id, + "copy_file", + "in_progress", + f"Copying file to {new_filename}", + file_id=new_record.id, + ) # Copy the file instead of moving it shutil.copy(original_local_file, new_local_path) - log_task_progress(task_id, "copy_file", "success", f"File copied to {new_filename}", file_id=new_record.id) + log_task_progress( + task_id, + "copy_file", + "success", + f"File copied to {new_filename}", + file_id=new_record.id, + ) # Update the DB with final local filename new_record.local_filename = new_local_path db.commit() + # Store file_id before session closes to avoid DetachedInstanceError + file_id = new_record.id + # 2. Check for embedded text (outside the DB session to avoid long open transactions) logger.info(f"[{task_id}] Checking for embedded text in PDF") - log_task_progress(task_id, "check_text", "in_progress", "Checking for embedded text", file_id=new_record.id) - with open(new_local_path, 'rb') as file: + log_task_progress( + task_id, + "check_text", + "in_progress", + "Checking for embedded text", + file_id=file_id, + ) + with open(new_local_path, "rb") as file: pdf_reader = PyPDF2.PdfReader(file) has_text = False for page in pdf_reader.pages: @@ -112,30 +163,74 @@ def process_document(self, original_local_file: str): break if has_text: - logger.info(f"[{task_id}] PDF {original_local_file} contains embedded text. Processing locally.") - log_task_progress(task_id, "check_text", "success", "Embedded text found, extracting locally", file_id=new_record.id) + logger.info( + f"[{task_id}] PDF {original_local_file} contains embedded text. Processing locally." + ) + log_task_progress( + task_id, + "check_text", + "success", + "Embedded text found, extracting locally", + file_id=file_id, + ) # Extract text locally logger.info(f"[{task_id}] Extracting text from PDF") - log_task_progress(task_id, "extract_text", "in_progress", "Extracting text locally", file_id=new_record.id) + log_task_progress( + task_id, + "extract_text", + "in_progress", + "Extracting text locally", + file_id=file_id, + ) extracted_text = "" - with open(new_local_path, 'rb') as file: + with open(new_local_path, "rb") as file: pdf_reader = PyPDF2.PdfReader(file) for page in pdf_reader.pages: extracted_text += page.extract_text() + "\n" - + logger.info(f"[{task_id}] Extracted {len(extracted_text)} characters") - log_task_progress(task_id, "extract_text", "success", f"Extracted {len(extracted_text)} characters", file_id=new_record.id) + log_task_progress( + task_id, + "extract_text", + "success", + f"Extracted {len(extracted_text)} characters", + file_id=file_id, + ) # Call metadata extraction directly logger.info(f"[{task_id}] Queueing metadata extraction") - log_task_progress(task_id, "process_document", "success", "Queued for metadata extraction", file_id=new_record.id) - extract_metadata_with_gpt.delay(new_filename, extracted_text, new_record.id) - return {"file": new_local_path, "status": "Text extracted locally", "file_id": new_record.id} + log_task_progress( + task_id, + "process_document", + "success", + "Queued for metadata extraction", + file_id=file_id, + ) + extract_metadata_with_gpt.delay(new_filename, extracted_text, file_id) + return { + "file": new_local_path, + "status": "Text extracted locally", + "file_id": file_id, + } # 3. If no embedded text, queue Azure Document Intelligence processing - logger.info(f"[{task_id}] No embedded text found. Queueing Azure Document Intelligence processing") - log_task_progress(task_id, "check_text", "success", "No embedded text, queuing OCR", file_id=new_record.id) - log_task_progress(task_id, "process_document", "success", "Queued for OCR processing", file_id=new_record.id) - process_with_azure_document_intelligence.delay(new_filename, new_record.id) - return {"file": new_local_path, "status": "Queued for OCR", "file_id": new_record.id} + logger.info( + f"[{task_id}] No embedded text found. Queueing Azure Document Intelligence processing" + ) + log_task_progress( + task_id, + "check_text", + "success", + "No embedded text, queuing OCR", + file_id=file_id, + ) + log_task_progress( + task_id, + "process_document", + "success", + "Queued for OCR processing", + file_id=file_id, + ) + process_with_azure_document_intelligence.delay(new_filename, file_id) + return {"file": new_local_path, "status": "Queued for OCR", "file_id": file_id} diff --git a/tests/test_process_document.py b/tests/test_process_document.py new file mode 100644 index 00000000..84112867 --- /dev/null +++ b/tests/test_process_document.py @@ -0,0 +1,244 @@ +""" +Unit tests for the process_document task. + +These tests verify that the process_document task correctly handles file processing +and doesn't cause DetachedInstanceError when accessing database objects. +""" + +import os +import pytest +from unittest.mock import patch, MagicMock +from sqlalchemy.orm import Session + +from app.tasks.process_document import process_document +from app.models import FileRecord + + +@pytest.mark.unit +@pytest.mark.requires_db +def test_process_document_stores_file_id_before_session_closes(db_session, tmp_path): + """ + Test that process_document stores file_id before the database session closes. + This prevents DetachedInstanceError when accessing the file_id after the session ends. + """ + # Create a test PDF file with embedded text + test_pdf = tmp_path / "test.pdf" + pdf_content = b"""%PDF-1.4 +1 0 obj +<< +/Type /Catalog +/Pages 2 0 R +>> +endobj +2 0 obj +<< +/Type /Pages +/Kids [3 0 R] +/Count 1 +>> +endobj +3 0 obj +<< +/Type /Page +/Parent 2 0 R +/MediaBox [0 0 612 792] +/Resources << +/Font << +/F1 << +/Type /Font +/Subtype /Type1 +/BaseFont /Helvetica +>> +>> +>> +/Contents 4 0 R +>> +endobj +4 0 obj +<< +/Length 44 +>> +stream +BT +/F1 12 Tf +100 700 Td +(Test content) Tj +ET +endstream +endobj +xref +0 5 +0000000000 65535 f +0000000009 00000 n +0000000058 00000 n +0000000115 00000 n +0000000306 00000 n +trailer +<< +/Size 5 +/Root 1 0 R +>> +startxref +399 +%%EOF +""" + test_pdf.write_bytes(pdf_content) + + # Mock environment and dependencies + with patch("app.tasks.process_document.SessionLocal") as mock_session_local, patch( + "app.tasks.process_document.settings" + ) as mock_settings, patch("app.tasks.process_document.log_task_progress"), patch( + "app.tasks.process_document.extract_metadata_with_gpt" + ) as mock_extract: + + # Setup mocks + mock_settings.workdir = str(tmp_path) + mock_session_local.return_value.__enter__.return_value = db_session + mock_session_local.return_value.__exit__.return_value = None + mock_extract.delay = MagicMock() + + # Get the actual function from the task (skip decorators) + # The task function signature is: def process_document(self, original_local_file: str) + task_run_func = process_document.run + + # Call the task's run method directly + result = task_run_func(str(test_pdf)) + + # Verify that the task completed successfully + assert "file_id" in result + assert result["status"] == "Text extracted locally" + + # Verify that a FileRecord was created + file_record = db_session.query(FileRecord).first() + assert file_record is not None + assert file_record.original_filename == "test.pdf" + + # Verify that extract_metadata_with_gpt was called with the file_id + # This would fail if file_id wasn't extracted before the session closed + mock_extract.delay.assert_called_once() + call_args = mock_extract.delay.call_args + assert call_args[0][2] == file_record.id # Third argument should be file_id + + +@pytest.mark.unit +@pytest.mark.requires_db +def test_process_document_duplicate_file(db_session, tmp_path): + """ + Test that duplicate files are detected and processing is skipped. + """ + # Create a test PDF file + test_pdf = tmp_path / "test.pdf" + test_pdf.write_bytes(b"test content") + + # Pre-create a FileRecord with the same hash + from app.utils import hash_file + + filehash = hash_file(str(test_pdf)) + + existing_record = FileRecord( + filehash=filehash, + original_filename="existing.pdf", + local_filename="/tmp/existing.pdf", + file_size=100, + mime_type="application/pdf", + ) + db_session.add(existing_record) + db_session.commit() + existing_id = existing_record.id + + # Mock environment and dependencies + with patch("app.tasks.process_document.SessionLocal") as mock_session_local, patch( + "app.tasks.process_document.log_task_progress" + ): + + # Setup mocks + mock_session_local.return_value.__enter__.return_value = db_session + mock_session_local.return_value.__exit__.return_value = None + + # Call the task's run method directly + result = process_document.run(str(test_pdf)) + + # Verify that duplicate was detected + assert result["status"] == "duplicate_file" + assert result["file_id"] == existing_id + + # Verify only one FileRecord exists + assert db_session.query(FileRecord).count() == 1 + + +@pytest.mark.unit +@pytest.mark.requires_db +def test_process_document_queues_ocr_for_image_pdf(db_session, tmp_path): + """ + Test that PDFs without embedded text are queued for OCR processing. + """ + # Create a test PDF file without text + test_pdf = tmp_path / "test_image.pdf" + pdf_content = b"""%PDF-1.4 +1 0 obj +<< +/Type /Catalog +/Pages 2 0 R +>> +endobj +2 0 obj +<< +/Type /Pages +/Kids [3 0 R] +/Count 1 +>> +endobj +3 0 obj +<< +/Type /Page +/Parent 2 0 R +/MediaBox [0 0 612 792] +>> +endobj +xref +0 4 +0000000000 65535 f +0000000009 00000 n +0000000058 00000 n +0000000115 00000 n +trailer +<< +/Size 4 +/Root 1 0 R +>> +startxref +197 +%%EOF +""" + test_pdf.write_bytes(pdf_content) + + # Mock environment and dependencies + with patch("app.tasks.process_document.SessionLocal") as mock_session_local, patch( + "app.tasks.process_document.settings" + ) as mock_settings, patch("app.tasks.process_document.log_task_progress"), patch( + "app.tasks.process_document.process_with_azure_document_intelligence" + ) as mock_azure: + + # Setup mocks + mock_settings.workdir = str(tmp_path) + mock_session_local.return_value.__enter__.return_value = db_session + mock_session_local.return_value.__exit__.return_value = None + mock_azure.delay = MagicMock() + + # Call the task's run method directly + result = process_document.run(str(test_pdf)) + + # Verify that OCR was queued + assert result["status"] == "Queued for OCR" + assert "file_id" in result + + # Verify that process_with_azure_document_intelligence was called with file_id + mock_azure.delay.assert_called_once() + call_args = mock_azure.delay.call_args + + # Verify a FileRecord was created + file_record = db_session.query(FileRecord).first() + assert file_record is not None + + # The second argument should be the file_id + assert call_args[0][1] == file_record.id