FLYPIE-270 Add record count to publish task

Author: sensei100

funcake_dags/scripts/index.sh

@@ -27,6 +27,15 @@ rm lib/$INDEXER.rb.bak
 # grab list of items from designated aws bucket (creds are envvars), then index each item
 TEMPFILE=$(mktemp /tmp/index-output.XXXXXX)
 PUBLISH_TASK_REPORT=$AIRFLOW_HOME/dags/funcake_dags/scripts/publish_task_report.rb
+
+solr_curl() {
+  if [ -n "${SOLR_AUTH_USER:-}" ]; then
+    curl -fsS -u "${SOLR_AUTH_USER}:${SOLR_AUTH_PASSWORD:-}" "$@"
+  else
+    curl -fsS "$@"
+  fi
+}
+
 report_and_cleanup() {
   rc=$?
   cat "$TEMPFILE" | ruby "$PUBLISH_TASK_REPORT" || true
@@ -70,3 +79,20 @@ do
     exit 1
   fi
 done
+
+SOLR_BASE_URL="${FUNCAKE_OAI_SOLR_URL%/}"
+if ! solr_curl -X POST "${SOLR_BASE_URL}/update?commit=true" >/dev/null; then
+  echo "ERROR: unable to commit Solr collection after publish"
+  exit 1
+fi
+
+SOLR_COUNT_RESPONSE=$(solr_curl "${SOLR_BASE_URL}/select?q=*:*&rows=0&wt=json")
+SOLR_PUBLISHED_COUNT=$(printf '%s' "$SOLR_COUNT_RESPONSE" | jq -r '.response.numFound // empty' 2>/dev/null || true)
+
+if ! [[ "${SOLR_PUBLISHED_COUNT:-}" =~ ^[0-9]+$ ]]; then
+  echo "ERROR: unable to parse published record count from Solr response"
+  exit 1
+fi
+
+export SOLR_PUBLISHED_COUNT
+echo "Published Record Count: $SOLR_PUBLISHED_COUNT"

funcake_dags/scripts/publish_task_report.rb

@@ -7,4 +7,11 @@ def collect_push_totals(output_lines)
     .sum(&:to_i)
 end
 
-puts "{ 'published': '#{collect_push_totals($stdin)}' }"
+def published_total(output_lines)
+  solr_count = ENV["SOLR_PUBLISHED_COUNT"]
+  return solr_count if solr_count&.match?(/\A\d+\z/)
+
+  collect_push_totals(output_lines)
+end
+
+puts "{ 'published': '#{published_total($stdin)}' }"

tests/index_script_test.py

@@ -30,12 +30,39 @@ def test_index_mode_with_empty_data_reports_missing_record_sets(self):
         self.assertNotIn("DATA: unbound variable", result.stdout + result.stderr)
         self.assertIn("ERROR: no record sets provided in DATA", result.stdout)
 
-    def _run_script(self, data: str | None = None):
+    def test_publish_succeeds_when_solr_commit_and_count_succeed(self):
+        result = self._run_script(
+            data='["set1.xml"]',
+            bundle_output="finished Traject::Indexer#process: 2 records in 0.1 seconds\n",
+            solr_count="2",
+        )
+
+        self.assertEqual(result.returncode, 0)
+        self.assertIn("Published Record Count: 2", result.stdout)
+
+    def test_publish_fails_when_solr_commit_fails(self):
+        result = self._run_script(
+            data='["set1.xml"]',
+            bundle_output="finished Traject::Indexer#process: 2 records in 0.1 seconds\n",
+            curl_commit_exit=1,
+        )
+
+        self.assertNotEqual(result.returncode, 0)
+        self.assertIn("ERROR: unable to commit Solr collection after publish", result.stdout)
+
+    def _run_script(
+        self,
+        data: str | None = None,
+        bundle_output: str = "",
+        solr_count: str = "0",
+        curl_commit_exit: int = 0,
+    ):
         tempdir = tempfile.mkdtemp()
         self.addCleanup(shutil.rmtree, tempdir, ignore_errors=True)
         tempdir_path = Path(tempdir)
         bin_dir = tempdir_path / "bin"
         bin_dir.mkdir()
+        (tempdir_path / ".bashrc").write_text("", encoding="utf-8")
 
         report_dir = tempdir_path / "dags" / "funcake_dags" / "scripts"
         report_dir.mkdir(parents=True)
@@ -53,8 +80,55 @@ def _run_script(self, data: str | None = None):
 """,
         )
         self._write_executable(bin_dir / "gem", "#!/usr/bin/env bash\nexit 0\n")
-        self._write_executable(bin_dir / "bundle", "#!/usr/bin/env bash\nexit 0\n")
-        self._write_executable(bin_dir / "aws", "#!/usr/bin/env bash\nexit 0\n")
+        self._write_executable(
+            bin_dir / "bundle",
+            f"""#!/usr/bin/env bash
+if [ "$1" = "exec" ]; then
+  printf %s {bundle_output!r}
+fi
+exit 0
+""",
+        )
+        self._write_executable(
+            bin_dir / "aws",
+            """#!/usr/bin/env bash
+if [ "$1" = "s3" ] && [ "$2" = "presign" ]; then
+  printf "http://example.test/presigned\\n"
+fi
+exit 0
+""",
+        )
+        self._write_executable(
+            bin_dir / "curl",
+            f"""#!/usr/bin/env bash
+args="$*"
+if [[ "$args" == *"/update?commit=true"* ]]; then
+  exit {curl_commit_exit}
+fi
+if [[ "$args" == *"/select?q=*:*&rows=0&wt=json"* ]]; then
+  printf '{{"response":{{"numFound":{solr_count}}}}}\\n'
+  exit 0
+fi
+exit 0
+""",
+        )
+        self._write_executable(
+            bin_dir / "jq",
+            """#!/usr/bin/env python3
+import json
+import sys
+
+payload = json.load(sys.stdin)
+query = sys.argv[-1]
+if query == ".[]":
+    for item in payload:
+        print(item)
+elif query == ".response.numFound // empty":
+    print(payload["response"]["numFound"])
+else:
+    raise SystemExit(f"unsupported jq query: {query}")
+""",
+        )
         self._write_executable(
             bin_dir / "ruby",
             """#!/usr/bin/env bash

tests/publish_task_report_test.py

@@ -0,0 +1,41 @@
+import os
+import subprocess
+import unittest
+from pathlib import Path
+
+
+REPO_ROOT = Path(__file__).resolve().parents[1]
+SCRIPT_PATH = REPO_ROOT / "funcake_dags" / "scripts" / "publish_task_report.rb"
+
+
+class PublishTaskReportTest(unittest.TestCase):
+    def test_falls_back_to_traject_batch_totals(self):
+        result = subprocess.run(
+            ["ruby", str(SCRIPT_PATH)],
+            input=(
+                "finished Traject::Indexer#process: 2 records in 0.1 seconds\n"
+                "finished Traject::Indexer#process: 3 records in 0.1 seconds\n"
+            ),
+            text=True,
+            capture_output=True,
+            check=False,
+        )
+
+        self.assertEqual(result.returncode, 0)
+        self.assertEqual(result.stdout.strip(), "{ 'published': '5' }")
+
+    def test_prefers_solr_published_count_when_present(self):
+        env = os.environ.copy()
+        env["SOLR_PUBLISHED_COUNT"] = "4"
+
+        result = subprocess.run(
+            ["ruby", str(SCRIPT_PATH)],
+            input="finished Traject::Indexer#process: 99 records in 0.1 seconds\n",
+            text=True,
+            capture_output=True,
+            check=False,
+            env=env,
+        )
+
+        self.assertEqual(result.returncode, 0)
+        self.assertEqual(result.stdout.strip(), "{ 'published': '4' }")