Speed up KPIs update script: dedupe work URIs server-side

What this script is for (context first)

scripts/kpis_update_script.py computes the OPERAS KPI figures: how many distinct books, and how many distinct publishers, have usage events recorded in metrics-api. The intent of the original code, preserved in this rewrite:

  1. Collect the DOI work URIs (info:doi:10.…) that appear in the metrics-api event table up to a given date (or within one calendar month, in monthly mode).
  2. Match them against the identifiers-api DB — URIs known there as monograph/book/edited-book count as books, and their DOI prefixes (e.g. info:doi:10.5334) identify publishers.
  3. For each publisher prefix, query Crossref (falling back to DataCite) to resolve the publisher's name, and to catch books that have events but are missing from the identifiers DB.
  4. Write books_…/publishers_…/not_found_uris_….json to scripts/kpis_output_files/<YYYY-MM>/ for the portal.

Why it took 20 hours and then crashed

The old step 1 streamed every individual event row (135M+) through cloud_sql_proxy to the laptop, wrote each one to a temp JSON file (including a timestamp field that was never used again), then re-parsed that file line-by-line — only to reduce everything to sets of distinct URIs. The per-event granularity was discarded entirely.

That also explains the crash: a server-side cursor holds one connection/transaction open for the full 20 hours, and eventually Cloud SQL or the proxy drops it (server closed the connection unexpectedly at row ~134.9M), losing the whole run since there's no checkpointing.

The fix

Deduplicate inside Postgres:

SELECT DISTINCT work_uri, SUBSTRING(work_uri FROM '^([^/]+)')
FROM event
WHERE work_uri LIKE 'info:doi:10%' AND timestamp < :end_date

The event table already has idx_work_uri, and the distinct-URI result set is tiny compared to 135M event rows — so this is one server-side pass with a small transfer, instead of shipping every event over the proxy. The temp-JSON write/re-parse stage is deleted outright. Connections also now set TCP keepalives + pool_pre_ping as a belt-and-braces measure.

Bugs fixed along the way

  • Leaked loop variable decided the output (old line 236): if publisher.get('prefix'): ran outside the loop, so whether "not found URIs" got computed depended on whichever publisher happened to be iterated last — and when it didn't run, the code silently fell back to an unrelated module-level global (the DataCite 404 set). With zero publishers it was a NameError, 20 hours in. Now computed unconditionally.
  • Crossref was effectively never used: the prefix was passed to api.crossref.org/prefixes/info:doi:10.5334/works without stripping info:doi:, which 404s, so every lookup fell through to DataCite. The prefix is now cleaned for both APIs.
  • Crossref paging: offset is capped at 10,000 by Crossref, so big publishers would silently truncate. Switched to cursor=* deep paging.
  • DataCite paging: only the first page (~25 records) was ever fetched. Now follows links.next with page[size]=1000.
  • No HTTP timeouts/retries: a stalled response could hang a worker thread forever. All calls now use a shared session with a 30s timeout and retry/backoff on 429/5xx.
  • Monthly off-by-one: timestamp < end_date at midnight excluded the last day of the month. Monthly mode now covers the full calendar month of the given date.
  • Crash-prone DataCite parsing: error-shaped payloads ({"errors": …}) raised KeyError; responses are now guarded.

Other changes

  • CLI instead of editing source: python kpis_update_script.py 2025-04-01 [--mode monthly] with --help, replacing the hardcoded date in __main__. A module docstring documents the whole pipeline and the METRICS_API_DB_* / IDENTIFIERS_DB_* env vars (defaults suit local cloud_sql_proxy on 54322/54323).
  • Unit tests (scripts/tests/, 13 tests) covering the pure logic: URI classification, Crossref/DataCite response matching incl. error payloads, publisher dedup, unresolved-prefix reporting, and the monthly/cumulative windows. Wired into tox as a second unittest discover pass.
  • flake8: the script was excluded in .flake8; the rewrite is lint-clean so the exclusion is removed. (Note: CI only lints src/, so this affects local runs.)

Testing

  • All 13 new unit tests pass (run in a clean venv matching the CI env, where sqlalchemy isn't installed — the tests stub it).
  • flake8 scripts/kpis_update_script.py clean.
  • ⚠️ Not yet run against the real databases — whoever runs the next KPI update should sanity-check the BOOKS/PUBLISHERS counts against a previous month's output. Expected runtime is minutes (one indexed SELECT DISTINCT) instead of ~20 hours, with the API-lookup phase now likely dominating.

Merge request reports

Loading