fix: harden schema probe cache against concurrency and stale code versions

Two latent issues, both surfaced while populating new fields:

- The schema-file and probe-cache writers used a shared "<path>.tmp" and the
  schema-file writer did not catch errors. Concurrent extraction workers that
  all cache-miss and re-probe raced on that temp and one worker's rename failed
  with FileNotFoundError, aborting the run. Each writer now uses a unique temp
  file (tempfile.mkstemp) and atomically renames; failures are logged, not
  raised. As the proper fix, convert_relationships now builds the canonical
  write schemas once in the parent before forking, so workers inherit the cache
  and never re-probe concurrently.

- A probe cache is keyed by the source-file set, which is unchanged across a
  code change, so a stale cache silently shadowed updated extraction logic
  (e.g. the new affiliation `years` field). The cache now records and validates
  the schema version on load and is rejected on mismatch; bumped to v2.

Tests in test_probe_cache.py cover version rejection and concurrency-safe writes
(blocking the legacy shared temp name to deterministically distinguish the fix).

Author: Mearman

sync/extract.py

@@ -1318,6 +1318,12 @@ def convert_relationships(
         " + ".join(all_pending_types),
     )
 
+    # Build the canonical write schemas once, here in the parent, before forking
+    # the pool. Workers inherit the populated lru cache across fork, so they
+    # never re-probe (which would have every worker sample the source and race
+    # to rewrite the shared schema-probe cache).
+    _entity_arrow_schemas(entity_type)
+
     # Distribute files across workers — each worker processes all types
     actual_workers = min(n_workers, len(files_to_process))
     actual_workers = max(1, actual_workers)

sync/schema.py

@@ -16,6 +16,7 @@
 import json
 import logging
 import os
+import tempfile
 from dataclasses import dataclass, field
 from functools import lru_cache
 from pathlib import Path
@@ -1767,7 +1768,14 @@ def probe_schema_multi(
 
 # ── Committed schema file ──────────────────────────────────────────────
 
-_SCHEMA_FILE_VERSION = 1
+# Bump this whenever the probe/extraction logic changes what columns a schema
+# produces. Both the committed schema file and the on-disk probe cache record
+# the version they were written with and are rejected on load if it no longer
+# matches, so a logic change automatically invalidates stale caches instead of
+# silently shadowing the new behaviour (the source-file hash alone can't catch
+# a code change). v2: locations capture all scalar leaves; relationship probes
+# carry lists-of-scalars as native list columns.
+_SCHEMA_FILE_VERSION = 2
 _SCHEMA_FILE = Path(__file__).resolve().parent.parent / "openalex.schema.json"
 _PROBE_SAMPLE_SIZE = 100  # records sampled per file
 _PROBE_SAMPLE_FILES = 25  # files sampled, evenly spaced across the date range
@@ -1805,9 +1813,27 @@ def _write_schema_file(cache: dict[str, dict[str, Any]]) -> None:
         "version": _SCHEMA_FILE_VERSION,
         "entities": cache,
     }
-    tmp_path = path.with_suffix(".tmp")
-    tmp_path.write_text(json.dumps(payload, indent=2, sort_keys=True), encoding="utf-8")
-    tmp_path.replace(path)
+    # Unique temp file per writer (not a shared "<path>.tmp") so concurrent
+    # workers that all cache-miss and re-probe at once can't clobber each
+    # other's temp and fail the rename. Failures are logged, not raised, so a
+    # best-effort cache update never aborts the run.
+    tmp_path: str | None = None
+    try:
+        fd, tmp_path = tempfile.mkstemp(
+            dir=str(path.parent), prefix=".openalex.schema.", suffix=".tmp",
+        )
+        with os.fdopen(fd, "w", encoding="utf-8") as fh:
+            json.dump(payload, fh, indent=2, sort_keys=True)
+        os.replace(tmp_path, path)
+        tmp_path = None
+    except OSError as exc:
+        log.warning("Cannot write schema file %s: %s", path, exc)
+    finally:
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
     _load_schema_file.cache_clear()
 
 
@@ -1905,6 +1931,11 @@ def _load_probe_cache(entity: str, key: str) -> EntitySchema | None:
         return None
     if payload.get("cache_key") != key or payload.get("entity") != entity:
         return None
+    # Reject caches written by a different probe-logic version — the source-file
+    # hash is unchanged across a code change, so without this a stale cache would
+    # shadow the updated extraction behaviour.
+    if payload.get("version") != _SCHEMA_FILE_VERSION:
+        return None
     schema_dict = payload.get("schema")
     if not isinstance(schema_dict, dict):
         return None
@@ -1929,15 +1960,28 @@ def _store_probe_cache(entity: str, key: str, schema: EntitySchema) -> None:
         "cache_key": key,
         "schema": schema.to_dict(),
     }
-    tmp_path = path.with_suffix(path.suffix + ".tmp")
+    # Write to a unique temp file, then atomically rename onto the final path.
+    # A shared "<path>.tmp" would let concurrent writers (e.g. a forked
+    # extraction pool that all cache-miss at once) clobber each other's temp and
+    # fail the rename with FileNotFoundError. Every writer produces the same
+    # schema, so a per-writer temp plus last-writer-wins rename is safe.
+    tmp_path: str | None = None
     try:
-        tmp_path.write_text(
-            json.dumps(payload, indent=2, sort_keys=True),
-            encoding="utf-8",
+        fd, tmp_path = tempfile.mkstemp(
+            dir=str(_PROBE_CACHE_DIR), prefix=f"schema_{entity}_", suffix=".tmp",
         )
-        tmp_path.replace(path)
+        with os.fdopen(fd, "w", encoding="utf-8") as fh:
+            json.dump(payload, fh, indent=2, sort_keys=True)
+        os.replace(tmp_path, path)
+        tmp_path = None
     except OSError as exc:
         log.warning("Cannot write schema probe cache %s: %s", path, exc)
+    finally:
+        if tmp_path is not None:
+            try:
+                os.unlink(tmp_path)
+            except OSError:
+                pass
 
 
 # ── Schema probing from JSONL ───────────────────────────────────────────

tests/test_probe_cache.py

@@ -0,0 +1,89 @@
+"""Regression tests for the schema probe cache: version invalidation and
+concurrency-safe writes."""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import pytest
+
+import sync.schema as schema
+from sync.schema import EntitySchema, FieldSchema
+
+
+def _tiny_schema(entity: str = "widgets") -> EntitySchema:
+    return EntitySchema(
+        entity=entity,
+        id_col="widget_id",
+        id_path="id",
+        id_type="int",
+        fields=[
+            FieldSchema(
+                json_key="cited_by_count", pattern="scalar", rel_name="widget_main",
+                scalar_cols=[{"col": "cited_by_count", "path": "cited_by_count", "type": "int"}],
+            )
+        ],
+    )
+
+
+@pytest.fixture
+def cache_dir(tmp_path, monkeypatch):
+    d = tmp_path / "probe-cache"
+    monkeypatch.setattr(schema, "_PROBE_CACHE_DIR", d)
+    return d
+
+
+class TestVersionInvalidation:
+    def test_roundtrip_same_version_loads(self, cache_dir):
+        schema._store_probe_cache("widgets", "key0", _tiny_schema())
+        loaded = schema._load_probe_cache("widgets", "key0")
+        assert loaded is not None
+        assert loaded.entity == "widgets"
+
+    def test_stale_version_is_rejected(self, cache_dir, monkeypatch):
+        # Written under the current version...
+        schema._store_probe_cache("widgets", "key0", _tiny_schema())
+        path = schema._probe_cache_path("widgets", "key0")
+        assert path.is_file()
+        # ...then the probe logic version moves on: the cache must be ignored,
+        # not silently shadow the new behaviour (the cache_key is unchanged).
+        monkeypatch.setattr(schema, "_SCHEMA_FILE_VERSION", schema._SCHEMA_FILE_VERSION + 1)
+        assert schema._load_probe_cache("widgets", "key0") is None
+
+    def test_cache_payload_records_version(self, cache_dir):
+        schema._store_probe_cache("widgets", "key0", _tiny_schema())
+        payload = json.loads(schema._probe_cache_path("widgets", "key0").read_text())
+        assert payload["version"] == schema._SCHEMA_FILE_VERSION
+
+
+class TestConcurrentWriteSafe:
+    """A unique temp file per writer means a contended/blocked legacy shared
+    ``<path>.tmp`` name can't break a write. Blocking that exact name (with a
+    directory, so the old shared-tmp write would fail) deterministically
+    distinguishes the fix from the racy original — no timing luck involved.
+    """
+
+    def test_write_schema_file_survives_blocked_shared_tmp(self, tmp_path, monkeypatch):
+        sf = tmp_path / "openalex.schema.json"
+        monkeypatch.setattr(schema, "_SCHEMA_FILE", sf)
+        # Occupy the legacy shared temp name; the old code wrote here and would
+        # raise (and it didn't catch the error, which is what aborted the job).
+        (tmp_path / "openalex.schema.tmp").mkdir()
+        schema._write_schema_file({"widgets": _tiny_schema().to_dict()})
+        assert sf.is_file()
+        payload = json.loads(sf.read_text())
+        assert payload["version"] == schema._SCHEMA_FILE_VERSION
+        assert "widgets" in payload["entities"]
+        assert not list(tmp_path.glob(".openalex.schema.*.tmp")), "temp not cleaned up"
+
+    def test_store_probe_cache_survives_blocked_shared_tmp(self, cache_dir):
+        cache_dir.mkdir(parents=True, exist_ok=True)
+        path = schema._probe_cache_path("widgets", "key0")
+        # Block the legacy shared temp name (old: ``<path>.tmp``).
+        Path(str(path) + ".tmp").mkdir()
+        schema._store_probe_cache("widgets", "key0", _tiny_schema())
+        # New code uses a unique temp, so the cache is written and loadable;
+        # the old code would have failed to write it (and logged a warning).
+        loaded = schema._load_probe_cache("widgets", "key0")
+        assert loaded is not None and loaded.entity == "widgets"
+        assert not list(cache_dir.glob("schema_widgets_key0.*.tmp")), "temp not cleaned up"