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:
- Collect the DOI work URIs (
info:doi:10.…) that appear in the metrics-apieventtable up to a given date (or within one calendar month, inmonthlymode). - Match them against the identifiers-api DB — URIs known there as
monograph/book/edited-bookcount as books, and their DOI prefixes (e.g.info:doi:10.5334) identify publishers. - 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.
- Write
books_…/publishers_…/not_found_uris_….jsontoscripts/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_dateThe 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 aNameError, 20 hours in. Now computed unconditionally. - Crossref was effectively never used: the prefix was passed to
api.crossref.org/prefixes/info:doi:10.5334/workswithout strippinginfo:doi:, which 404s, so every lookup fell through to DataCite. The prefix is now cleaned for both APIs. - Crossref paging:
offsetis capped at 10,000 by Crossref, so big publishers would silently truncate. Switched tocursor=*deep paging. - DataCite paging: only the first page (~25 records) was ever fetched. Now follows
links.nextwithpage[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_dateat 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": …}) raisedKeyError; 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 theMETRICS_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 secondunittest discoverpass. - flake8: the script was excluded in
.flake8; the rewrite is lint-clean so the exclusion is removed. (Note: CI only lintssrc/, so this affects local runs.)
Testing
- All 13 new unit tests pass (run in a clean venv matching the CI env, where
sqlalchemyisn't installed — the tests stub it). flake8 scripts/kpis_update_script.pyclean.⚠️ 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 indexedSELECT DISTINCT) instead of ~20 hours, with the API-lookup phase now likely dominating.