Merge pull request #2001 from shashank03-dev/fix/issue-1966-ja-price-bulk-rpc

fix(etl): bulk-update Jan Aushadhi prices via atomic RPC

Author: dipexplorer

apps/etl/src/loaders/supabase_loader.py

@@ -44,6 +44,12 @@
 
 BATCH_SIZE = 100
 DELAY_SEC = 0.5
+
+# Postgres RPC (supabase/migrations) for atomic Jan Aushadhi price back-fill.
+# Updates medicines.jan_aushadhi_price in place by id without an INSERT, so the
+# table's NOT NULL columns (e.g. generic_name) are never validated against the
+# {id, jan_aushadhi_price}-only payload.
+JA_PRICE_BULK_RPC = "bulk_update_jan_aushadhi_price"
 BATCH_UPSERT_MAX_ATTEMPTS = 4
 BATCH_UPSERT_INITIAL_BACKOFF_SEC = 2.0
 BATCH_UPSERT_MAX_BACKOFF_SEC = 8.0
@@ -790,6 +796,12 @@ def _upsert_ja_price_update_batches(
         updates: list[dict],
         table: str,
     ) -> tuple[int, int]:
+        # The bulk RPC updates public.medicines.jan_aushadhi_price in place by id.
+        # Any other table (none today) must not reach it — fall back to per-row
+        # updates so behaviour stays correct if a caller passes a different table.
+        if table != "medicines":
+            return self._update_ja_price_rows_one_by_one(updates, table)
+
         updated = failed = 0
         total = len(updates)
         total_batches = (total + BATCH_SIZE - 1) // BATCH_SIZE
@@ -800,21 +812,44 @@ def _upsert_ja_price_update_batches(
             batch_end = batch_start + len(batch)
 
             try:
-                self._run_upsert_with_transient_retries(
-                    lambda: self.client.table(table).upsert(batch).execute(),
-                    f"merge_jan_aushadhi_price batch {batch_number}/{total_batches} upsert",
-                )
-                updated += len(batch)
+                response = self._bulk_update_ja_price(batch, batch_number, total_batches)
+                rpc_count = self._coerce_rpc_updated_count(response)
+                if rpc_count is None:
+                    # RPC committed (no exception → no INSERT, NOT NULL never hit)
+                    # but the row count came back in an unexpected shape. Assume the
+                    # whole batch landed and log loudly so a response-shape change is
+                    # visible rather than silently miscounted.
+                    logger.error(
+                        f"[Loader] merge_jan_aushadhi_price: batch "
+                        f"{batch_number}/{total_batches} bulk RPC returned an "
+                        f"unrecognized count shape {getattr(response, 'data', None)!r}; "
+                        f"assuming all {len(batch)} rows updated"
+                    )
+                    rpc_count = len(batch)
+                updated += rpc_count
+                # The selection guarantees every id was a real, still-NULL row, so a
+                # short count means rows vanished between the page scan and the UPDATE
+                # (e.g. deleted concurrently). Account for them as failed so the
+                # caller's checked == updated + skipped + failed invariant holds.
+                shortfall = len(batch) - rpc_count
+                if shortfall > 0:
+                    failed += shortfall
+                    logger.warning(
+                        f"[Loader] merge_jan_aushadhi_price: batch "
+                        f"{batch_number}/{total_batches} updated {rpc_count}/{len(batch)} "
+                        f"rows; {shortfall} unaccounted (rows missing at UPDATE time) "
+                        f"— counted as failed"
+                    )
                 logger.info(
                     f"[Loader] merge_jan_aushadhi_price: batch "
-                    f"{batch_number}/{total_batches} upserted {len(batch)} rows "
+                    f"{batch_number}/{total_batches} updated {rpc_count} rows "
                     f"({updated}/{total} page matches)"
                 )
             except Exception as e:
-                logger.warning(
+                logger.error(
                     f"[Loader] merge_jan_aushadhi_price: batch "
                     f"{batch_number}/{total_batches} rows {batch_start}-{batch_end} "
-                    f"failed: {e} - retrying row-by-row"
+                    f"bulk RPC failed: {e} - retrying row-by-row"
                 )
                 batch_updated, batch_failed = self._update_ja_price_rows_one_by_one(
                     batch,
@@ -825,6 +860,44 @@ def _upsert_ja_price_update_batches(
 
         return updated, failed
 
+    def _bulk_update_ja_price(
+        self,
+        batch: list[dict],
+        batch_number: int,
+        total_batches: int,
+    ) -> object:
+        """Atomically update one batch via the bulk RPC, with transient retries."""
+        captured: dict = {}
+
+        def _call() -> None:
+            captured["response"] = self.client.rpc(
+                JA_PRICE_BULK_RPC,
+                {"p_updates": batch},
+            ).execute()
+
+        self._run_upsert_with_transient_retries(
+            _call,
+            f"merge_jan_aushadhi_price batch {batch_number}/{total_batches} bulk RPC",
+        )
+        return captured.get("response")
+
+    @staticmethod
+    def _coerce_rpc_updated_count(response: object) -> "int | None":
+        """Read the integer row count returned by the bulk RPC.
+
+        PostgREST returns a scalar function result as the raw value, but tolerate
+        a single-element list too. Returns ``None`` when the shape is unrecognized
+        so the caller can decide how to account for it rather than guessing here.
+        """
+        data = getattr(response, "data", None)
+        if isinstance(data, bool):  # bool is an int subclass — exclude it
+            return None
+        if isinstance(data, int):
+            return data
+        if isinstance(data, list) and len(data) == 1 and isinstance(data[0], int):
+            return data[0]
+        return None
+
     def _update_ja_price_rows_one_by_one(
         self,
         updates: list[dict],

apps/etl/tests/test_loader.py

@@ -575,8 +575,47 @@ def __init__(self, medicines=None, transient_batch_failures=0):
         self.medicines = medicines or []
         self.update_calls = []
         self.upsert_calls = []
+        self.rpc_calls = []
         self.transient_batch_failures = transient_batch_failures
         self.transient_batch_attempts = 0
+        # When set, the RPC returns this as response.data instead of the real
+        # changed-row count — used to exercise short-count / unrecognized shapes.
+        self.rpc_data_override = None
+        self.rpc_override_set = False
+
+    def rpc(self, name, params):
+        """Fake the bulk_update_jan_aushadhi_price RPC: atomic UPDATE by id."""
+        client = self
+
+        class _FakeRpc:
+            def execute(self_inner):
+                client.rpc_calls.append((name, params))
+                updates = params.get("p_updates") or []
+
+                # Mirror the real loader's batch retry surface: a multi-row batch
+                # can hit a transient error before succeeding on a later attempt.
+                if len(updates) > 1:
+                    client.transient_batch_attempts += 1
+                    if client.transient_batch_attempts <= client.transient_batch_failures:
+                        raise TimeoutError(
+                            "connection timed out during Jan Aushadhi price bulk RPC"
+                        )
+
+                changed = 0
+                for update in updates:
+                    row_id = update.get("id")
+                    new_price = update.get("jan_aushadhi_price")
+                    if row_id is None or new_price is None:
+                        continue
+                    for med in client.medicines:
+                        if med.get("id") == row_id:
+                            med["jan_aushadhi_price"] = new_price
+                            changed += 1
+                if client.rpc_override_set:
+                    return FakeExecuteResponse(client.rpc_data_override)
+                return FakeExecuteResponse(changed)
+
+        return _FakeRpc()
 
     def table(self, name):
         t = MergeFakeTable(name, self)
@@ -644,7 +683,78 @@ def test_ja_backfill_updates_null_jan_aushadhi_price_rows(tmp_path):
     assert medicines[1]["jan_aushadhi_price"] == 25.00
 
 
-def test_ja_backfill_retries_transient_batch_upsert_before_fallback(tmp_path, monkeypatch):
+def test_ja_backfill_uses_bulk_rpc_with_id_and_price_only(tmp_path):
+    """Regression for #1966: back-fill goes through the bulk_update RPC with a
+    {id, jan_aushadhi_price}-only payload, never a PostgREST upsert (which would
+    fail the medicines.generic_name NOT NULL constraint and fall back to slow
+    row-by-row PATCHes)."""
+    medicines = [
+        {"id": "m1", "generic_name": "Paracetamol", "strength": "500mg",
+         "source": "commercial", "jan_aushadhi_price": None},
+    ]
+    nppa_csv = _write_nppa_csv(tmp_path, [
+        {"generic_name": "paracetamol", "strength": "500mg", "mrp": "18.50"},
+    ])
+    client = MergeFakeSupabaseClient(medicines=medicines)
+    loader = make_merge_loader(client, tmp_path)
+
+    stats = loader.merge_jan_aushadhi_price(nppa_csv=nppa_csv)
+
+    assert stats["updated"] == 1
+    assert stats["failed"] == 0
+    # No upsert and no row-by-row fallback were used.
+    assert client.upsert_calls == []
+    assert client.update_calls == []
+    # Exactly one bulk RPC call, carrying only id + jan_aushadhi_price.
+    assert len(client.rpc_calls) == 1
+    name, params = client.rpc_calls[0]
+    assert name == "bulk_update_jan_aushadhi_price"
+    assert params["p_updates"] == [{"id": "m1", "jan_aushadhi_price": 18.50}]
+
+
+def test_ja_backfill_counts_short_rpc_result_as_failed(tmp_path):
+    """If the bulk RPC updates fewer rows than the batch (a row vanished between
+    the page scan and the UPDATE), the shortfall is counted as failed so the
+    checked == updated + skipped + failed invariant holds — not silently dropped."""
+    medicines = [
+        {"id": "m1", "generic_name": "Paracetamol", "jan_aushadhi_price": None},
+    ]
+    client = MergeFakeSupabaseClient(medicines=medicines)
+    loader = make_merge_loader(client, tmp_path)
+
+    # m2 has no matching medicine row, so the RPC reports 1 updated, not 2.
+    batch = [
+        {"id": "m1", "jan_aushadhi_price": 18.50},
+        {"id": "m2", "jan_aushadhi_price": 25.00},
+    ]
+    updated, failed = loader._upsert_ja_price_update_batches(batch, "medicines")
+
+    assert updated == 1
+    assert failed == 1
+    assert medicines[0]["jan_aushadhi_price"] == 18.50
+    assert client.update_calls == []  # no row-by-row fallback was triggered
+
+
+def test_ja_backfill_assumes_full_batch_on_unrecognized_rpc_shape(tmp_path):
+    """An unrecognized RPC count shape is assumed to be a full success (the RPC
+    committed without raising) rather than miscounted as failed."""
+    medicines = [
+        {"id": "m1", "generic_name": "Paracetamol", "jan_aushadhi_price": None},
+    ]
+    client = MergeFakeSupabaseClient(medicines=medicines)
+    client.rpc_override_set = True
+    client.rpc_data_override = {"unexpected": "shape"}
+    loader = make_merge_loader(client, tmp_path)
+
+    batch = [{"id": "m1", "jan_aushadhi_price": 18.50}]
+    updated, failed = loader._upsert_ja_price_update_batches(batch, "medicines")
+
+    assert updated == 1
+    assert failed == 0
+    assert client.update_calls == []
+
+
+def test_ja_backfill_retries_transient_batch_rpc_before_fallback(tmp_path, monkeypatch):
     medicines = [
         {"id": "m1", "generic_name": "Paracetamol", "strength": "500mg",
          "source": "commercial", "jan_aushadhi_price": None},
@@ -671,7 +781,7 @@ def test_ja_backfill_retries_transient_batch_upsert_before_fallback(tmp_path, mo
     assert stats["updated"] == 2
     assert stats["failed"] == 0
     assert client.transient_batch_attempts == 3
-    assert len(client.upsert_calls) == 3
+    assert len(client.rpc_calls) == 3
     assert client.update_calls == []
     assert len(sleep_calls) == 2
 

supabase/migrations/20260616174742_add_bulk_update_jan_aushadhi_price_rpc.sql

@@ -0,0 +1,67 @@
+-- Bulk Jan Aushadhi price back-fill RPC
+--
+-- WHY THIS EXISTS
+-- ---------------
+-- The ETL loader (apps/etl/src/loaders/supabase_loader.py :: merge_jan_aushadhi_price)
+-- back-fills medicines.jan_aushadhi_price in batches. It previously used PostgREST
+-- .upsert(), which compiles to INSERT ... ON CONFLICT. Even for rows that already
+-- exist, PostgREST sends a full INSERT, so every payload was validated against the
+-- table's NOT NULL columns (notably medicines.generic_name). The batch payload only
+-- carries {id, jan_aushadhi_price}, so the INSERT path raised
+--   "null value in column generic_name violates not-null constraint"
+-- which forced a row-by-row PATCH fallback (~100x slower).
+--
+-- This function performs an atomic, set-based UPDATE of jan_aushadhi_price only,
+-- keyed by id, in a single statement. It never inserts, so NOT NULL columns are
+-- never touched and no other field can be corrupted.
+--
+-- Payload shape (single jsonb array argument):
+--   [{"id": "<uuid>", "jan_aushadhi_price": 10.5}, ...]
+--
+-- Returns the number of rows actually updated.
+
+CREATE OR REPLACE FUNCTION public.bulk_update_jan_aushadhi_price(p_updates jsonb)
+RETURNS integer
+LANGUAGE plpgsql
+-- SECURITY INVOKER (the default): this function WRITES to medicines, so it must
+-- run with the caller's privileges and stay subject to RLS. The ETL connects as
+-- service_role (which the medicines_service_write policy in
+-- 20260529000000_add_rls_policies.sql allows), so the back-fill works; any other
+-- role's UPDATE is filtered out by RLS. Unlike the read-only RPCs in this repo,
+-- exposing a writer with definer rights would let anon/authenticated callers edit
+-- prices via PostgREST.
+SET search_path = public, pg_temp
+AS $$
+DECLARE
+  v_updated integer;
+BEGIN
+  IF p_updates IS NULL OR jsonb_typeof(p_updates) <> 'array' THEN
+    RAISE EXCEPTION 'p_updates must be a JSON array of {id, jan_aushadhi_price} objects, got %',
+      COALESCE(jsonb_typeof(p_updates), 'null');
+  END IF;
+
+  WITH payload AS (
+    SELECT id, jan_aushadhi_price
+    FROM jsonb_to_recordset(p_updates)
+      AS x(id uuid, jan_aushadhi_price numeric)
+    WHERE id IS NOT NULL
+      AND jan_aushadhi_price IS NOT NULL
+  ),
+  changed AS (
+    UPDATE public.medicines m
+    SET jan_aushadhi_price = p.jan_aushadhi_price
+    FROM payload p
+    WHERE m.id = p.id
+    RETURNING m.id
+  )
+  SELECT count(*) INTO v_updated FROM changed;
+
+  RETURN v_updated;
+END;
+$$;
+
+-- Postgres grants EXECUTE on new functions to PUBLIC by default, which would make
+-- this writer callable by anon/authenticated through PostgREST. Lock it down to
+-- the ETL's service_role so prices can only be back-filled by the loader.
+REVOKE EXECUTE ON FUNCTION public.bulk_update_jan_aushadhi_price(jsonb) FROM PUBLIC;
+GRANT EXECUTE ON FUNCTION public.bulk_update_jan_aushadhi_price(jsonb) TO service_role;