Update to move files as they get processed.

Author: dkinzer

centralized_metadata/optimize_pdf_dag.py

@@ -175,6 +175,31 @@ def run_and_stream(command, prefix=None):
         raise subprocess.CalledProcessError(return_code, command)
 
 
+def _move_processed_pair(original_pdf: Path, optimized_pdf: Path):
+    """
+    Move the optimized and source PDFs into their destination folders.
+    Returns paths to the moved files.
+    """
+    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))
+    return originals_destination, optimized_destination
+
+
 def process_pdfs(**context):
     """Iterate through the directory and run OCRmyPDF with optimize level 1."""
     pdf_directory = _resolve_pdf_directory(context)
@@ -210,48 +235,20 @@ def process_pdfs(**context):
         logging.info("Running command: %s", " ".join(command))
 
         run_and_stream(command, prefix=pdf_file.name)
+        originals_destination, optimized_destination = _move_processed_pair(
+            pdf_file, output_pdf
+        )
 
         processed_files.append(
-            {"original": str(pdf_file), "optimized": str(output_pdf)}
+            {
+                "original": str(originals_destination),
+                "optimized": str(optimized_destination),
+            }
         )
 
     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,
@@ -270,19 +267,11 @@ def move_processed_pdfs(**context):
     dag=DAG,
 )
 
-MOVE_PROCESSED_FILES = PythonOperator(
-    task_id="move_processed_pdfs",
-    python_callable=move_processed_pdfs,
-    provide_context=True,
-    dag=DAG,
-)
-
 SUCCESS = EmptyOperator(
     task_id='success',
     on_success_callback=[teamspostonsuccess],
     trigger_rule="none_failed_min_one_success",
     dag=DAG,
 )
 
-RUN_OCR.set_downstream(MOVE_PROCESSED_FILES)
-MOVE_PROCESSED_FILES.set_downstream(SUCCESS)
+RUN_OCR.set_downstream(SUCCESS)

tests/optimize_pdf_dag_test.py

@@ -15,7 +15,7 @@ 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", "move_processed_pdfs", "success"])
+        self.assertEqual(task_ids, ["run_ocrmypdf", "success"])
 
     @mock.patch("centralized_metadata.optimize_pdf_dag.run_and_stream")
     def test_process_pdfs_runs_command_per_file(self, mock_run_and_stream):
@@ -26,21 +26,30 @@ def test_process_pdfs_runs_command_per_file(self, mock_run_and_stream):
             open(pdf_one, "wb").close()
             open(pdf_two, "wb").close()
 
+            def fake_run(command, prefix=None):
+                Path(command[5]).write_bytes(b"optimized")
+
+            mock_run_and_stream.side_effect = fake_run
+
             result = process_pdfs(params={"pdf_directory": tmp_dir})
 
             self.assertEqual(
                 result,
                 [
                     {
-                        "original": pdf_one,
+                        "original": os.path.realpath(
+                            os.path.join(tmp_dir, "Originals", "file_one.pdf")
+                        ),
                         "optimized": os.path.realpath(
-                            os.path.join(tmp_dir, "file_one_opti.pdf")
+                            os.path.join(tmp_dir, "Optimized", "file_one_opti.pdf")
                         ),
                     },
                     {
-                        "original": pdf_two,
+                        "original": os.path.realpath(
+                            os.path.join(tmp_dir, "Originals", "file_two.pdf")
+                        ),
                         "optimized": os.path.realpath(
-                            os.path.join(tmp_dir, "file_two_opti.pdf")
+                            os.path.join(tmp_dir, "Optimized", "file_two_opti.pdf")
                         ),
                     },
                 ],
@@ -61,6 +70,13 @@ def test_process_pdfs_runs_command_per_file(self, mock_run_and_stream):
                 )
                 self.assertEqual(call_args.kwargs.get("prefix"), expected_file.name)
 
+            originals_dir = Path(tmp_dir) / "Originals"
+            optimized_dir = Path(tmp_dir) / "Optimized"
+            self.assertTrue((originals_dir / "file_one.pdf").exists())
+            self.assertTrue((originals_dir / "file_two.pdf").exists())
+            self.assertTrue((optimized_dir / "file_one_opti.pdf").exists())
+            self.assertTrue((optimized_dir / "file_two_opti.pdf").exists())
+
     @mock.patch("centralized_metadata.optimize_pdf_dag.run_and_stream")
     def test_process_pdfs_handles_empty_directory(self, mock_run_and_stream):
         """Confirm we short-circuit gracefully when no PDFs exist."""
@@ -77,6 +93,11 @@ def test_process_pdfs_prefers_dag_run_conf(self, mock_run_and_stream):
             pdf_path = os.path.realpath(os.path.join(tmp_dir, "file.pdf"))
             open(pdf_path, "wb").close()
 
+            def fake_run(command, prefix=None):
+                Path(command[5]).write_text("optimized")
+
+            mock_run_and_stream.side_effect = fake_run
+
             dag_run_mock = mock.Mock()
             dag_run_mock.conf = {"pdf_directory": tmp_dir}
 
@@ -88,9 +109,11 @@ def test_process_pdfs_prefers_dag_run_conf(self, mock_run_and_stream):
                 result,
                 [
                     {
-                        "original": pdf_path,
+                        "original": os.path.realpath(
+                            os.path.join(tmp_dir, "Originals", "file.pdf")
+                        ),
                         "optimized": os.path.realpath(
-                            os.path.join(tmp_dir, "file_opti.pdf")
+                            os.path.join(tmp_dir, "Optimized", "file_opti.pdf")
                         ),
                     }
                 ],
@@ -111,6 +134,11 @@ def test_process_pdfs_uses_share_root_and_relative_path(
             pdf_path = target_dir / "variable.pdf"
             pdf_path.touch()
 
+            def fake_run(command, prefix=None):
+                Path(command[5]).write_bytes(b"optimized")
+
+            mock_run_and_stream.side_effect = fake_run
+
             def fake_variable_get(key, default_var=None):
                 if key == "OCR_PDF_SHARE_ROOT":
                     return share_root
@@ -126,8 +154,8 @@ def fake_variable_get(key, default_var=None):
                 result,
                 [
                     {
-                        "original": str(pdf_path.resolve()),
-                        "optimized": str((target_dir / "variable_opti.pdf").resolve()),
+                        "original": str((target_dir / "Originals" / "variable.pdf").resolve()),
+                        "optimized": str((target_dir / "Optimized" / "variable_opti.pdf").resolve()),
                     }
                 ],
             )