From e75c634f7d3281064dd09a48ec77791f5682b596 Mon Sep 17 00:00:00 2001 From: debpalash <4178343+debpalash@users.noreply.github.com> Date: Tue, 13 Oct 2026 23:45:28 +0530 Subject: [PATCH] perf(workbook): select stale refresh cells with bounded SQL Refresh no longer loads every workbook row and enrichment overlay into Python to decide which cells are stale. _stale_row_columns keyset-paginates workbook_rows.id and, per page, cross-joins the rows with the enrichment column ids (each with its own TTL cutoff) and LEFT JOINs the overlays on (workbook_id, COALESCE(lead_id, id), column_id), returning only stale pairs: missing overlay, non-complete or NULL status, NULL updated_at, or updated_at older than the cutoff. Cutoffs bind as naive UTC to match the column on SQLite and PostgreSQL. Every page is scanned, so results are identical to the previous loop; WORKBOOK_REFRESH_SCAN_BATCH_SIZE (default 1000) only bounds per-query work. Part of #33. Signed-off-by: debpalash Co-Authored-By: Claude Opus 5.5 --- .github/workflows/ci.yml | 1 + apps/api/core/config.py | 6 + apps/api/services/workbook/refresh.py | 114 ++++++++++++----- ...012-refresh-dependent-and-new-row-cells.md | 19 +++ tests/refresh_stale_support.py | 115 ++++++++++++++++++ tests/test_workbook_refresh_cell_scope.py | 13 ++ tests/test_workbook_refresh_stale_sql.py | 97 +++++++++++++++ tests/test_workbook_refresh_stale_sql_pg.py | 58 +++++++++ 8 files changed, 394 insertions(+), 29 deletions(-) create mode 100644 tests/refresh_stale_support.py create mode 100644 tests/test_workbook_refresh_stale_sql.py create mode 100644 tests/test_workbook_refresh_stale_sql_pg.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2957ec35..9d2cb4ac 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -62,6 +62,7 @@ jobs: tests/test_entity_identity_pg.py tests/test_pg_tenancy_rls.py tests/test_lead_paging_pg.py + tests/test_workbook_refresh_stale_sql_pg.py frontend: name: Frontend lint, tests and build diff --git a/apps/api/core/config.py b/apps/api/core/config.py index 45859708..edd3b306 100644 --- a/apps/api/core/config.py +++ b/apps/api/core/config.py @@ -163,6 +163,12 @@ class Settings(BaseSettings): # Cloud deployments may populate this; comma-separated when set via env. AUTOMATIONS_WEBHOOK_DOMAIN_ALLOWLIST: list[str] = [] + # ── Living workbook refresh ── + # Rows per keyset page when a refresh selects stale enrichment cells in SQL. + # Bounds memory/query size per round trip; every page is still scanned, so + # the selected cells do not depend on this value. + WORKBOOK_REFRESH_SCAN_BATCH_SIZE: int = 1000 + # ── Signal scanner (legacy hiring/JobSpy scan → unified ORM signals) ── # The scanner now iterates EVERY workspace and fans out provider I/O per # lead, so per-workspace and global caps keep a multi-tenant scan polite. diff --git a/apps/api/services/workbook/refresh.py b/apps/api/services/workbook/refresh.py index 759a7a51..659b7829 100644 --- a/apps/api/services/workbook/refresh.py +++ b/apps/api/services/workbook/refresh.py @@ -52,37 +52,93 @@ def _interval_minutes(policy: dict) -> Optional[int]: return INTERVAL_MINUTES.get(str(interval).lower()) -def _stale_row_columns(db, workbook_id: str, enrichment_cols: List[dict], ttl_map: dict) -> dict[str, list[str]]: - """Exact enrichment cells missing, failed or older than their field's TTL.""" - rows = db.query(WorkbookRow).filter(WorkbookRow.workbook_id == workbook_id).all() - if not rows: - return {} - # Legacy overlay keys follow the execution key (linked lead id or row id). - overlays = db.query(WorkbookEnrichment).filter( - WorkbookEnrichment.workbook_id == workbook_id - ).all() - seen = {(o.lead_id, o.column_id): o for o in overlays} +def _stale_row_columns(db, workbook_id: str, enrichment_cols: List[dict], ttl_map: dict, + batch_size: Optional[int] = None) -> dict[str, list[str]]: + """Exact enrichment cells missing, failed or older than their field's TTL. + + Selected in SQL, keyset-paginated over ``workbook_rows.id`` so a refresh + never loads the whole workbook (or its overlays) into Python. Each page + cross-joins the page's rows with the enrichment columns and LEFT JOINs the + overlay on the execution key (linked lead id, else row id); only stale + (row, column) pairs come back. A cell is stale when its overlay is missing, + its status is anything but ``complete`` (including NULL), its + ``updated_at`` is NULL, or ``updated_at`` is older than + ``now - staleness_ttl_days[target_field or id]`` (default 30 days). + + Every page is scanned, so the result is independent of ``batch_size``. + Rows are returned in id order with columns in config order. Tenancy: the + explicit ``workbook_id`` filter plus RLS on both tables, as before. + """ + from sqlalchemy import DateTime, Integer, String, and_, func, literal, or_, select, true, union_all + + from apps.api.core.config import settings - stale = {} + # Column ids are NOT NULL on overlays; de-dup preserves config order (the + # caller de-dups the per-row lists anyway). + cutoffs: dict[str, datetime] = {} now = _now() - for row in rows: - row_id = row.id - lead_id = row.lead_id or row_id - for col in enrichment_cols: - cid = col.get("id") - ttl_days = ttl_map.get(col.get("target_field") or cid, DEFAULT_STALENESS_DAYS) - o = seen.get((lead_id, cid)) - if o is None or o.status != "complete": - stale.setdefault(str(row_id), []).append(cid) - continue - upd = o.updated_at - if upd is None: - stale.setdefault(str(row_id), []).append(cid) - continue - if upd.tzinfo is None: - upd = upd.replace(tzinfo=timezone.utc) - if upd < now - timedelta(days=ttl_days): - stale.setdefault(str(row_id), []).append(cid) + for col in enrichment_cols: + cid = col.get("id") + if not cid or cid in cutoffs: + continue + ttl_days = ttl_map.get(col.get("target_field") or cid, DEFAULT_STALENESS_DAYS) + # Overlay timestamps are naive UTC (DateTime without tz on both + # dialects); compare naive UTC to avoid session-timezone coercion on + # PostgreSQL and offset-suffixed strings on SQLite. + cutoffs[cid] = (now - timedelta(days=ttl_days)).astimezone(timezone.utc).replace(tzinfo=None) + if not cutoffs: + return {} + + columns = union_all(*( + select(literal(cid, String).label("cid"), + literal(ordinal, Integer).label("ord"), + literal(cutoff, DateTime()).label("cutoff")) + for ordinal, (cid, cutoff) in enumerate(cutoffs.items()) + )).subquery("refresh_cols") + + size = max(1, int(batch_size or settings.WORKBOOK_REFRESH_SCAN_BATCH_SIZE)) + R, E = WorkbookRow, WorkbookEnrichment + stale: dict[str, list[str]] = {} + after = None + while True: + # Keyset page bounds: the next `size` row ids after the previous page. + page_q = select(R.id).where(R.workbook_id == workbook_id) + if after is not None: + page_q = page_q.where(R.id > after) + page = page_q.order_by(R.id).limit(size).subquery("refresh_page") + hi, count = db.execute(select(func.max(page.c.id), func.count())).one() + if not count: + break + + in_page = [R.workbook_id == workbook_id, R.id <= hi] + if after is not None: + in_page.append(R.id > after) + overlay_on = and_( + E.workbook_id == workbook_id, + E.lead_id == func.coalesce(R.lead_id, R.id), + E.column_id == columns.c.cid, + ) + rows = db.execute( + select(R.id, columns.c.cid) + .select_from(R) + .join(columns, true()) + .outerjoin(E, overlay_on) + .where(*in_page) + .where(or_( + E.id.is_(None), + E.status.is_(None), + E.status != "complete", + E.updated_at.is_(None), + E.updated_at < columns.c.cutoff, + )) + .order_by(R.id, columns.c.ord) + ).all() + for row_id, cid in rows: + stale.setdefault(str(row_id), []).append(cid) + + if count < size: + break + after = hi return stale diff --git a/docs/adr/20261012-refresh-dependent-and-new-row-cells.md b/docs/adr/20261012-refresh-dependent-and-new-row-cells.md index 20160ced..8210e3b2 100644 --- a/docs/adr/20261012-refresh-dependent-and-new-row-cells.md +++ b/docs/adr/20261012-refresh-dependent-and-new-row-cells.md @@ -38,6 +38,25 @@ writer during sourcing are left alone (covered by `test_concurrently_inserted_row_is_not_treated_as_sourced`). No extra row-ID queries are issued. +## Addendum (2026-10-13): stale selection in bounded SQL (part of #33) + +Stale-cell selection no longer loads every workbook row and overlay into +Python. `_stale_row_columns` keyset-paginates `workbook_rows.id` (page size +`WORKBOOK_REFRESH_SCAN_BATCH_SIZE`, default 1000) and, per page, cross-joins +the rows with the enrichment column ids (each carrying its own cutoff, +`now - staleness_ttl_days[target_field or id]`, default 30 days) and LEFT JOINs +`workbook_enrichments` on `(workbook_id, COALESCE(lead_id, id), column_id)`. +Only stale pairs are returned: overlay missing, status not `complete` (NULL +included), `updated_at` NULL, or `updated_at` older than the cutoff. Cutoffs +are bound as naive UTC to match the naive-UTC `updated_at` column on both +SQLite and PostgreSQL. Every page is scanned, so the selected cells are +identical to the previous Python loop regardless of page size; it is a bound +on per-query work, not a cap on refresh scope. Column entries without an id +(malformed config) are skipped rather than reported as a `None` column. +Parity against the previous implementation is covered by +`tests/test_workbook_refresh_stale_sql.py` and, under FORCE RLS on +PostgreSQL, `tests/test_workbook_refresh_stale_sql_pg.py`. + ## Validation Python 3.12.3 and the repository's frozen uv dependencies: both new regressions diff --git a/tests/refresh_stale_support.py b/tests/refresh_stale_support.py new file mode 100644 index 00000000..2ad07035 --- /dev/null +++ b/tests/refresh_stale_support.py @@ -0,0 +1,115 @@ +"""Shared fixtures for the SQL stale-cell selection parity tests (SQLite + PG). + +``legacy_stale_row_columns`` is the pre-#33 Python implementation, kept +verbatim as the oracle: the SQL selection must return exactly what it did. +""" + +import random +from datetime import timedelta, timezone + +from apps.api.services.workbook.models import Workbook, WorkbookEnrichment, WorkbookRow +from apps.api.services.workbook.refresh import DEFAULT_STALENESS_DAYS + +COLUMNS = [ + {"id": "email", "type": "input"}, + {"id": "summary", "type": "ai_formula", "prompt": "s {email}"}, # ttl 7 + {"id": "pitch", "type": "ai_formula", "prompt": "p {email}"}, # ttl 30 + {"id": "domain_col", "type": "enrichment", "target_field": "domain"}, # ttl via target_field: 2 + {"id": "no_ttl", "type": "waterfall"}, # default 30 + {"id": "quick", "type": "ai_formula", "prompt": "q {email}"}, # ttl 0 +] +# "domain_col" also has its own entry: target_field must take precedence. +TTL_MAP = {"summary": 7, "pitch": 30, "domain": 2, "quick": 0, "domain_col": 999} +ENRICHMENT_COLS = [c for c in COLUMNS if c["type"] in ("enrichment", "waterfall", "ai_formula")] + + +def legacy_stale_row_columns(db, workbook_id, enrichment_cols, ttl_map, now): + rows = db.query(WorkbookRow).filter(WorkbookRow.workbook_id == workbook_id).all() + if not rows: + return {} + overlays = db.query(WorkbookEnrichment).filter( + WorkbookEnrichment.workbook_id == workbook_id + ).all() + seen = {(o.lead_id, o.column_id): o for o in overlays} + stale = {} + for row in rows: + row_id = row.id + lead_id = row.lead_id or row_id + for col in enrichment_cols: + cid = col.get("id") + ttl_days = ttl_map.get(col.get("target_field") or cid, DEFAULT_STALENESS_DAYS) + o = seen.get((lead_id, cid)) + if o is None or o.status != "complete": + stale.setdefault(str(row_id), []).append(cid) + continue + upd = o.updated_at + if upd is None: + stale.setdefault(str(row_id), []).append(cid) + continue + if upd.tzinfo is None: + upd = upd.replace(tzinfo=timezone.utc) + if upd < now - timedelta(days=ttl_days): + stale.setdefault(str(row_id), []).append(cid) + return stale + + +def _age(now, days, hours=0): + return (now - timedelta(days=days, hours=hours)).astimezone(timezone.utc).replace(tzinfo=None) + + +STATES = ["fresh", "stale", "missing", "error", "pending", "null_status", + "null_updated", "edge_fresh", "edge_stale"] + + +def populate_random(db, ws, now, *, n_rows=40, seed=33, name="Stale SQL parity"): + """A workbook with a randomized mix of fresh/stale/missing/failed cells. + + Also writes a sibling workbook whose complete, fresh overlays share the + same execution keys; they must never leak into the first workbook's + selection. Returns (workbook_id, states_seen). + """ + rng = random.Random(seed) + wb = Workbook(name=name, workspace_id=ws, columns_config=COLUMNS, + refresh_policy={"staleness_ttl_days": TTL_MAP}) + other = Workbook(name=f"{name} sibling", workspace_id=ws, columns_config=COLUMNS) + db.add_all([wb, other]) + db.flush() + seen_states = set() + for i in range(n_rows): + row = WorkbookRow(workbook_id=wb.id, workspace_id=ws, position=i, + lead_id=(100000 + i) if rng.random() < 0.5 else None, + data={"email": f"r{i}@example.test"}, enrichments={}) + db.add(row) + db.flush() + key = row.lead_id or row.id + for col in ENRICHMENT_COLS: + db.add(WorkbookEnrichment(workbook_id=other.id, workspace_id=ws, lead_id=key, + column_id=col["id"], status="complete", + updated_at=_age(now, 0))) + ttl = TTL_MAP.get(col.get("target_field") or col["id"], DEFAULT_STALENESS_DAYS) + state = rng.choice(STATES) + seen_states.add(state) + if state == "missing": + continue + status, upd = "complete", _age(now, 0) + if state == "stale": + upd = _age(now, ttl + 5) + elif state == "error": + status = "error" + elif state == "pending": + status = "pending" + elif state == "edge_fresh": + upd = _age(now, ttl, hours=-1) + elif state == "edge_stale": + upd = _age(now, ttl, hours=1) + cell = WorkbookEnrichment(workbook_id=wb.id, workspace_id=ws, lead_id=key, + column_id=col["id"], status=status, updated_at=upd) + db.add(cell) + db.flush() + if state in ("null_status", "null_updated"): + # Bypass ORM/server defaults to persist genuine NULLs. + db.query(WorkbookEnrichment).filter(WorkbookEnrichment.id == cell.id).update( + {"status": None} if state == "null_status" else {"updated_at": None}, + synchronize_session=False) + db.commit() + return wb.id, seen_states diff --git a/tests/test_workbook_refresh_cell_scope.py b/tests/test_workbook_refresh_cell_scope.py index 08321342..3d1e6c29 100644 --- a/tests/test_workbook_refresh_cell_scope.py +++ b/tests/test_workbook_refresh_cell_scope.py @@ -103,6 +103,19 @@ def test_all_fresh_cells_make_refresh_a_noop(refresh_runner, linked): assert calls == sent == [] +@pytest.mark.parametrize("linked", [False, True]) +def test_refresh_scope_is_independent_of_scan_page_size(refresh_runner, monkeypatch, linked): + from apps.api.core.config import settings + + sessions, now, calls, sent = refresh_runner + monkeypatch.setattr(settings, "WORKBOOK_REFRESH_SCAN_BATCH_SIZE", 1) + wid = populate(sessions, now, linked) + result = asyncio.run(refresh.refresh_workbook(wid, workspace_id="refresh-scope")) + assert sorted(calls) == [("pitch {email}", "row202@example.test"), + ("summary {email}", "row101@example.test")] + assert result == {"sourced": 0, "reenriched": 2, "stale_rows": 2} + + def test_refresh_recomputes_transitive_dependents_without_unrelated_spend(refresh_runner): sessions, now, calls, sent = refresh_runner columns = [ diff --git a/tests/test_workbook_refresh_stale_sql.py b/tests/test_workbook_refresh_stale_sql.py new file mode 100644 index 00000000..1bb2ff4a --- /dev/null +++ b/tests/test_workbook_refresh_stale_sql.py @@ -0,0 +1,97 @@ +"""Refresh selects stale cells in bounded SQL with the legacy Python semantics.""" +from datetime import datetime, timedelta, timezone + +import pytest +from sqlalchemy import create_engine +from sqlalchemy.orm import sessionmaker + +from apps.api.core.config import settings +from apps.api.database import Base +from apps.api.services.workbook import refresh +from apps.api.services.workbook.models import Workbook, WorkbookEnrichment, WorkbookRow +from tests.refresh_stale_support import ( + ENRICHMENT_COLS, STATES, TTL_MAP, legacy_stale_row_columns, populate_random, +) + +WS = "stale-sql" + + +@pytest.fixture +def db_now(tmp_path, monkeypatch): + engine = create_engine(f"sqlite:///{tmp_path / 'stale-sql.db'}") + Base.metadata.create_all(engine) + sessions = sessionmaker(bind=engine, autoflush=False) + now = datetime.now(timezone.utc) + monkeypatch.setattr(refresh, "_now", lambda: now) + with sessions() as db: + yield db, now + engine.dispose() + + +@pytest.mark.parametrize("batch_size", [1, 3, 7, 40, 1000]) +@pytest.mark.parametrize("seed", [33, 34, 35]) +def test_sql_selection_matches_legacy_python(db_now, batch_size, seed): + db, now = db_now + wid, seen_states = populate_random(db, WS, now, seed=seed) + assert seen_states == set(STATES) # every cell state is exercised + expected = legacy_stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, now) + got = refresh._stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, batch_size=batch_size) + assert got == expected + assert list(got) == list(expected) # row order (id) and per-row column order + assert 0 < sum(map(len, got.values())) < 40 * len(ENRICHMENT_COLS) # mixed fresh/stale + + +def test_setting_bounds_each_page(db_now, monkeypatch): + db, now = db_now + wid, _ = populate_random(db, WS, now) + expected = legacy_stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, now) + statements = [] + from sqlalchemy import event + + @event.listens_for(db.get_bind(), "before_cursor_execute") + def _count(conn, cursor, statement, params, context, executemany): # noqa: ANN001 + statements.append(statement) + + monkeypatch.setattr(settings, "WORKBOOK_REFRESH_SCAN_BATCH_SIZE", 5) + assert refresh._stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP) == expected + # 40 rows / 5 per page: 8 LIMIT-bounded page probes with one id-range + # selection each, plus a final empty probe. + assert sum("LIMIT" in s for s in statements) == 9 + assert len(statements) == 9 + 8 + + +def test_fresh_stale_missing_and_default_ttl_cells(db_now): + db, now = db_now + wb = Workbook(name="cases", workspace_id=WS, refresh_policy={}) + db.add(wb) + db.flush() + cols = [{"id": "a", "type": "ai_formula"}, {"id": "b", "type": "enrichment", "target_field": "f"}, + {"id": "c", "type": "waterfall"}] + ttl = {"a": 1, "f": 10} # c has no entry -> 30 days + cells = { # row_id -> {col: days old or None for missing} + 1: {"a": 0, "b": 0, "c": 0}, # all fresh + 2: {"a": 2, "b": 9, "c": 29}, # a stale only + 3: {"a": 0, "b": 11, "c": 31}, # b and c stale + 4: {}, # never enriched + } + for row_id, ages in cells.items(): + db.add(WorkbookRow(id=row_id, workbook_id=wb.id, workspace_id=WS, data={}, enrichments={})) + for cid, days in ages.items(): + db.add(WorkbookEnrichment(workbook_id=wb.id, workspace_id=WS, lead_id=row_id, column_id=cid, + status="complete", + updated_at=(now - timedelta(days=days)).replace(tzinfo=None))) + db.commit() + expected = {"2": ["a"], "3": ["b", "c"], "4": ["a", "b", "c"]} + assert legacy_stale_row_columns(db, wb.id, cols, ttl, now) == expected + for size in (1, 2, 100): + assert refresh._stale_row_columns(db, wb.id, cols, ttl, batch_size=size) == expected + + +def test_empty_workbook_and_no_columns(db_now): + db, now = db_now + wb = Workbook(name="empty", workspace_id=WS) + db.add(wb) + db.commit() + assert refresh._stale_row_columns(db, wb.id, ENRICHMENT_COLS, TTL_MAP, batch_size=1) == {} + wid, _ = populate_random(db, WS, now, n_rows=3) + assert refresh._stale_row_columns(db, wid, [], TTL_MAP) == {} diff --git a/tests/test_workbook_refresh_stale_sql_pg.py b/tests/test_workbook_refresh_stale_sql_pg.py new file mode 100644 index 00000000..170b5899 --- /dev/null +++ b/tests/test_workbook_refresh_stale_sql_pg.py @@ -0,0 +1,58 @@ +"""Stale refresh-cell selection on PostgreSQL with FORCE RLS. + +The SQL selection compares ``timestamp without time zone`` overlays against +bound cutoffs, cross-joins a literal column set and LEFT JOINs overlays on +``COALESCE(lead_id, id)``. Prove it matches the legacy Python oracle on PG +under the app role, and that another tenant's rows/overlays are invisible. +Set TEST_DATABASE_URL to a disposable PostgreSQL database (see pg_rls_support). +""" + +import os +from datetime import datetime, timezone +from uuid import uuid4 + +import pytest + +TEST_DATABASE_URL = os.getenv("TEST_DATABASE_URL") +pytestmark = [ + pytest.mark.postgres, + pytest.mark.skipif(not TEST_DATABASE_URL, reason="TEST_DATABASE_URL not set"), +] + + +@pytest.fixture(scope="module") +def app_session(): + from tests.pg_rls_support import rls_app_session + + factory, dispose = rls_app_session(TEST_DATABASE_URL, pool_size=5) + try: + yield factory + finally: + dispose() + + +@pytest.mark.parametrize("batch_size", [1, 4, 1000]) +def test_pg_sql_selection_matches_legacy_python(app_session, monkeypatch, batch_size): + from apps.api.core.tenancy import workspace_scope + from apps.api.services.workbook import refresh + from tests.refresh_stale_support import ( + ENRICHMENT_COLS, STATES, TTL_MAP, legacy_stale_row_columns, populate_random, + ) + + now = datetime.now(timezone.utc) + monkeypatch.setattr(refresh, "_now", lambda: now) + ws = f"stale_{uuid4().hex}" + with workspace_scope(ws): + with app_session() as db: + wid, seen_states = populate_random(db, ws, now, n_rows=30, seed=batch_size) + expected = legacy_stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, now) + got = refresh._stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, batch_size=batch_size) + assert seen_states == set(STATES) + assert got == expected and list(got) == list(expected) + assert 0 < sum(map(len, got.values())) < 30 * len(ENRICHMENT_COLS) + + # Another tenant sees neither the rows nor the overlays (RLS), even when + # it names the workbook id directly. + with workspace_scope(f"other_{uuid4().hex}"): + with app_session() as db: + assert refresh._stale_row_columns(db, wid, ENRICHMENT_COLS, TTL_MAP, batch_size=batch_size) == {}