Merge pull request #918 from tulibraries/ocrmypdf-file-updates-devo1302-devo-1339-devo-1340-devo-1301

Ocrmypdf file updates devo1302 devo 1339 devo 1340 devo 1301

Author: cdoyle-temple

centralized_metadata/ocrmypdf_processing_dag.py

@@ -1,5 +1,6 @@
 """Airflow DAG to run OCRmyPDF across a target directory."""
 import logging
+import shutil
 import subprocess
 from datetime import timedelta
 from pathlib import Path
@@ -14,6 +15,8 @@
 DEFAULT_PDF_DIRECTORY = str(Path(DEFAULT_SHARE_ROOT) / DEFAULT_RELATIVE_PATH)
 SHARE_ROOT_VARIABLE = "OCR_PDF_SHARE_ROOT"
 RELATIVE_PATH_VARIABLE = "OCR_PDF_RELATIVE_PATH"
+DEFAULT_SCHEDULE_INTERVAL = "@weekly"
+SCHEDULE_INTERVAL_VARIABLE = "OPTIMIZE_PDF_SCHEDULE_INTERVAL"
 
 DEFAULT_ARGS = {
     "owner": "airflow",
@@ -65,10 +68,11 @@ def process_pdfs(**context):
 
     if not pdf_files:
         logging.info("No PDF files found in %s", pdf_directory)
-        return "no-pdfs-found"
+        return []
 
+    processed_files = []
     for pdf_file in pdf_files:
-        output_pdf = pdf_file.with_name(f"{pdf_file.stem}_ocr.pdf")
+        output_pdf = pdf_file.with_name(f"{pdf_file.stem}_opti.pdf")
         command = [
             "ocrmypdf",
             "--optimize",
@@ -78,16 +82,55 @@ def process_pdfs(**context):
         ]
         logging.info("Running command: %s", " ".join(command))
         subprocess.run(command, check=True)  # Raises CalledProcessError on failure.
-
-    return f"processed-{len(pdf_files)}-pdfs"
+        processed_files.append(
+            {"original": str(pdf_file), "optimized": str(output_pdf)}
+        )
+
+    return processed_files
+
+
+def move_processed_pdfs(**context):
+    """Move optimized and original PDFs into their destination directories."""
+    ti = context["ti"]
+    processed_files = ti.xcom_pull(task_ids="run_ocrmypdf") or []
+    if not processed_files:
+        logging.info("No processed files to move")
+        return "no-files-to-move"
+
+    moved_count = 0
+    for file_info in processed_files:
+        original_pdf = Path(file_info["original"])
+        optimized_pdf = Path(file_info["optimized"])
+        pdf_directory = original_pdf.parent
+        optimized_directory = pdf_directory / "Optimized"
+        originals_directory = pdf_directory / "Originals"
+        optimized_directory.mkdir(parents=True, exist_ok=True)
+        originals_directory.mkdir(parents=True, exist_ok=True)
+
+        optimized_destination = optimized_directory / optimized_pdf.name
+        originals_destination = originals_directory / original_pdf.name
+        logging.info(
+            "Moving %s to %s and %s to %s",
+            optimized_pdf,
+            optimized_destination,
+            original_pdf,
+            originals_destination,
+        )
+        shutil.move(str(optimized_pdf), str(optimized_destination))
+        shutil.move(str(original_pdf), str(originals_destination))
+        moved_count += 1
+
+    return f"moved-{moved_count}-pdfs"
 
 
 DAG = airflow.DAG(
     "ocrmypdf_batch",
     default_args=DEFAULT_ARGS,
     catchup=False,
     max_active_runs=1,
-    schedule_interval=None,
+    schedule_interval=Variable.get(
+        SCHEDULE_INTERVAL_VARIABLE, default_var=DEFAULT_SCHEDULE_INTERVAL
+    ),
 )
 
 RUN_OCR = PythonOperator(
@@ -97,3 +140,12 @@ def process_pdfs(**context):
     params={"pdf_directory": DEFAULT_PDF_DIRECTORY},
     dag=DAG,
 )
+
+MOVE_PROCESSED_FILES = PythonOperator(
+    task_id="move_processed_pdfs",
+    python_callable=move_processed_pdfs,
+    provide_context=True,
+    dag=DAG,
+)
+
+RUN_OCR.set_downstream(MOVE_PROCESSED_FILES)

tests/ocrmypdf_processing_dag_test.py

@@ -15,30 +15,44 @@ def test_dag_loads(self):
         """Ensure the DAG is registered with the expected ID and task."""
         self.assertEqual(DAG.dag_id, "ocrmypdf_batch")
         task_ids = [task.task_id for task in DAG.tasks]
-        self.assertEqual(task_ids, ["run_ocrmypdf"])
+        self.assertEqual(task_ids, ["run_ocrmypdf", "move_processed_pdfs"])
 
     @mock.patch("centralized_metadata.ocrmypdf_processing_dag.subprocess.run")
     def test_process_pdfs_runs_command_per_file(self, mock_run):
         """Touch PDFs and verify we issue an OCR command for each."""
         with tempfile.TemporaryDirectory() as tmp_dir:
-            pdf_one = os.path.join(tmp_dir, "file_one.pdf")
-            pdf_two = os.path.join(tmp_dir, "file_two.pdf")
+            pdf_one = os.path.realpath(os.path.join(tmp_dir, "file_one.pdf"))
+            pdf_two = os.path.realpath(os.path.join(tmp_dir, "file_two.pdf"))
             open(pdf_one, "wb").close()
             open(pdf_two, "wb").close()
 
             result = process_pdfs(params={"pdf_directory": tmp_dir})
 
-            self.assertEqual(result, "processed-2-pdfs")
+            self.assertEqual(
+                result,
+                [
+                    {
+                        "original": pdf_one,
+                        "optimized": os.path.realpath(
+                            os.path.join(tmp_dir, "file_one_opti.pdf")
+                        ),
+                    },
+                    {
+                        "original": pdf_two,
+                        "optimized": os.path.realpath(
+                            os.path.join(tmp_dir, "file_two_opti.pdf")
+                        ),
+                    },
+                ],
+            )
             self.assertEqual(mock_run.call_count, 2)
 
             expected_files = [Path(pdf_one), Path(pdf_two)]
             for call_args, expected_file in zip(mock_run.call_args_list, expected_files):
                 command = call_args.args[0]
                 self.assertEqual(command[:3], ["ocrmypdf", "--optimize", "1"])
                 self.assertEqual(Path(command[3]).name, expected_file.name)
-                self.assertEqual(
-                    Path(command[4]).name, f"{expected_file.stem}_ocr.pdf"
-                )
+                self.assertEqual(Path(command[4]).name, f"{expected_file.stem}_opti.pdf")
                 self.assertTrue(call_args.kwargs.get("check"))
 
     @mock.patch("centralized_metadata.ocrmypdf_processing_dag.subprocess.run")
@@ -47,30 +61,42 @@ def test_process_pdfs_handles_empty_directory(self, mock_run):
         with tempfile.TemporaryDirectory() as tmp_dir:
             result = process_pdfs(params={"pdf_directory": tmp_dir})
 
-            self.assertEqual(result, "no-pdfs-found")
+            self.assertEqual(result, [])
             mock_run.assert_not_called()
 
     @mock.patch("centralized_metadata.ocrmypdf_processing_dag.subprocess.run")
     def test_process_pdfs_prefers_dag_run_conf(self, mock_run):
         """dag_run.conf should override params and defaults."""
         with tempfile.TemporaryDirectory() as tmp_dir:
-            pdf_path = os.path.join(tmp_dir, "file.pdf")
+            pdf_path = os.path.realpath(os.path.join(tmp_dir, "file.pdf"))
             open(pdf_path, "wb").close()
 
             dag_run_mock = mock.Mock()
             dag_run_mock.conf = {"pdf_directory": tmp_dir}
 
-            result = process_pdfs(dag_run=dag_run_mock, params={"pdf_directory": "/unused"})
-
-            self.assertEqual(result, "processed-1-pdfs")
+            result = process_pdfs(
+                dag_run=dag_run_mock, params={"pdf_directory": "/unused"}
+            )
+
+            self.assertEqual(
+                result,
+                [
+                    {
+                        "original": pdf_path,
+                        "optimized": os.path.realpath(
+                            os.path.join(tmp_dir, "file_opti.pdf")
+                        ),
+                    }
+                ],
+            )
             mock_run.assert_called_once()
 
     @mock.patch("centralized_metadata.ocrmypdf_processing_dag.Variable.get")
     @mock.patch("centralized_metadata.ocrmypdf_processing_dag.subprocess.run")
     def test_process_pdfs_uses_share_root_and_relative_path(self, mock_run, mock_variable_get):
         """Variable-based root + relative path should resolve to final directory."""
         with tempfile.TemporaryDirectory() as tmp_dir:
-            share_root = tmp_dir
+            share_root = os.path.realpath(tmp_dir)
             relative_path = "incoming"
             target_dir = Path(share_root) / relative_path
             target_dir.mkdir()
@@ -88,7 +114,15 @@ def fake_variable_get(key, default_var=None):
 
             result = process_pdfs()
 
-            self.assertEqual(result, "processed-1-pdfs")
+            self.assertEqual(
+                result,
+                [
+                    {
+                        "original": str(pdf_path.resolve()),
+                        "optimized": str((target_dir / "variable_opti.pdf").resolve()),
+                    }
+                ],
+            )
             command = mock_run.call_args.args[0]
             self.assertEqual(Path(command[3]).resolve(), pdf_path.resolve())