perf: optimize revenue forecast effective document links

app/document_reconciliation_service.py

@@ -181,28 +181,78 @@ def prepare_effective_document_links(conn: Any) -> str:
     result sets.  The temporary table is session-local and is populated solely
     through the group-aware resolver.
     """
-    opportunity_ids = {str(row[0]) for row in conn.execute(text("""
-        SELECT DISTINCT opportunity_id::text FROM commercial_documents
-        WHERE opportunity_id IS NOT NULL
-    """)).all()}
-    if document_reconciliation_v2_available(conn):
-        opportunity_ids.update(str(row[0]) for row in conn.execute(text("""
-            SELECT DISTINCT opportunity_id::text FROM opportunity_document_links
-            WHERE ended_at IS NULL
-        """)).all())
     conn.execute(text("""CREATE TEMP TABLE IF NOT EXISTS _effective_document_links (
         opportunity_id UUID NOT NULL, document_id UUID NOT NULL, document_kind TEXT NOT NULL,
         relationship TEXT NOT NULL, ended_at TIMESTAMPTZ) ON COMMIT DROP"""))
     conn.execute(text("TRUNCATE _effective_document_links"))
-    for opportunity_id in sorted(opportunity_ids):
-        for row in resolve_document_links(opportunity_id, conn=conn):
-            conn.execute(text("""INSERT INTO _effective_document_links
-                (opportunity_id,document_id,document_kind,relationship,ended_at)
-                VALUES(CAST(:oid AS UUID),CAST(:did AS UUID),:kind,:relationship,:ended_at)"""),
-                {"oid": opportunity_id, "did": row["document_id"],
-                 "kind": row.get("document_kind") or "unknown",
-                 "relationship": row.get("relationship") or "SECONDARY",
-                 "ended_at": row.get("ended_at")})
+    legacy_relationship = """CASE
+        WHEN d.is_active IS FALSE OR lower(COALESCE(NULLIF(d.role, ''), 'current')) = 'detached' THEN 'REMOVED'
+        WHEN d.is_primary IS TRUE AND lower(COALESCE(NULLIF(d.role, ''), 'current')) IN ('current','accepted') THEN 'PRIMARY'
+        WHEN lower(COALESCE(NULLIF(d.role, ''), 'current')) IN ('historical','history','superseded') THEN 'HISTORICAL'
+        ELSE 'SECONDARY'
+    END"""
+    if document_reconciliation_v2_available(conn):
+        # This is the relational form of resolve_document_links(): a legacy
+        # group switches to v2 only when every legacy document in that exact
+        # opportunity/kind group has a current v2 link. V2-only groups remain
+        # absent, matching the rollout resolver's current behaviour.
+        conn.execute(text(f"""
+            INSERT INTO _effective_document_links
+                (opportunity_id, document_id, document_kind, relationship, ended_at)
+            WITH legacy_groups AS (
+                SELECT DISTINCT d.opportunity_id, COALESCE(d.document_kind, '') AS document_kind
+                FROM commercial_documents d
+                WHERE d.opportunity_id IS NOT NULL
+            ), complete_groups AS (
+                SELECT g.opportunity_id, g.document_kind
+                FROM legacy_groups g
+                WHERE NOT EXISTS (
+                    SELECT 1
+                    FROM commercial_documents d
+                    WHERE d.opportunity_id = g.opportunity_id
+                      AND COALESCE(d.document_kind, '') = g.document_kind
+                      AND NOT EXISTS (
+                          SELECT 1
+                          FROM opportunity_document_links l
+                          WHERE l.opportunity_id = g.opportunity_id
+                            AND COALESCE(l.document_kind, '') = g.document_kind
+                            AND l.document_id = d.id
+                            AND l.ended_at IS NULL
+                      )
+                )
+            ), effective AS (
+                SELECT l.opportunity_id, l.document_id, l.document_kind,
+                       l.relationship, l.ended_at
+                FROM opportunity_document_links l
+                JOIN complete_groups g
+                  ON g.opportunity_id = l.opportunity_id
+                 AND g.document_kind = COALESCE(l.document_kind, '')
+                WHERE l.ended_at IS NULL
+                UNION ALL
+                SELECT d.opportunity_id, d.id, d.document_kind,
+                       {legacy_relationship} AS relationship,
+                       NULL::timestamptz AS ended_at
+                FROM commercial_documents d
+                LEFT JOIN complete_groups g
+                  ON g.opportunity_id = d.opportunity_id
+                 AND g.document_kind = COALESCE(d.document_kind, '')
+                WHERE d.opportunity_id IS NOT NULL
+                  AND g.opportunity_id IS NULL
+            )
+            SELECT opportunity_id, document_id, COALESCE(NULLIF(document_kind, ''), 'unknown'),
+                   relationship, ended_at
+            FROM effective
+        """))
+    else:
+        conn.execute(text(f"""
+            INSERT INTO _effective_document_links
+                (opportunity_id, document_id, document_kind, relationship, ended_at)
+            SELECT d.opportunity_id, d.id, COALESCE(NULLIF(d.document_kind, ''), 'unknown'),
+                   {legacy_relationship} AS relationship,
+                   NULL::timestamptz
+            FROM commercial_documents d
+            WHERE d.opportunity_id IS NOT NULL
+        """))
     return "_effective_document_links"
 
 

app/revenue_forecast_service.py

@@ -320,9 +320,13 @@ def _historical_stage_rates(conn: Any) -> dict[str, dict[str, Any]]:
     }
 
 
-def _realised_for_period(conn: Any, *, metric: str, period_start: date, period_end: date) -> dict[str, Any]:
+def _realised_for_period(
+    conn: Any, *, metric: str, period_start: date, period_end: date,
+    effective_links_prepared: bool = False,
+) -> dict[str, Any]:
     from app.document_reconciliation_service import prepare_effective_document_links
-    prepare_effective_document_links(conn)
+    if not effective_links_prepared:
+        prepare_effective_document_links(conn)
     if metric == "cash_received":
         rows = conn.execute(text("""
             WITH latest_doc AS (
@@ -385,9 +389,12 @@ def _realised_for_period(conn: Any, *, metric: str, period_start: date, period_e
     }
 
 
-def _already_realised_ids(conn: Any, *, metric: str) -> set[str]:
+def _already_realised_ids(
+    conn: Any, *, metric: str, effective_links_prepared: bool = False,
+) -> set[str]:
     from app.document_reconciliation_service import prepare_effective_document_links
-    prepare_effective_document_links(conn)
+    if not effective_links_prepared:
+        prepare_effective_document_links(conn)
     if metric == "cash_received":
         sql = """
             SELECT DISTINCT opportunity_id::text
@@ -441,8 +448,13 @@ def get_revenue_forecast(*, limit: int = 1000, month: str | None = None, metric:
         from app.document_reconciliation_service import prepare_effective_document_links
         prepare_effective_document_links(conn)
         historical = _historical_stage_rates(conn)
-        realised = _realised_for_period(conn, metric=metric, period_start=period_start, period_end=period_end)
-        already_realised_ids = _already_realised_ids(conn, metric=metric)
+        realised = _realised_for_period(
+            conn, metric=metric, period_start=period_start, period_end=period_end,
+            effective_links_prepared=True,
+        )
+        already_realised_ids = _already_realised_ids(
+            conn, metric=metric, effective_links_prepared=True,
+        )
         rows = conn.execute(text("""
             WITH latest_doc AS (
                 SELECT DISTINCT ON (l.opportunity_id)

tests/test_revenue_forecast_effective_links_performance.py

@@ -0,0 +1,102 @@
+from contextlib import nullcontext
+from datetime import date
+
+import app.document_reconciliation_service as document_service
+import app.revenue_forecast_service as forecast_service
+
+
+class _Result:
+    def __init__(self, rows=(), scalar_value=False):
+        self.rows = list(rows)
+        self.scalar_value = scalar_value
+
+    def mappings(self):
+        return self
+
+    def all(self):
+        return self.rows
+
+    def scalar(self):
+        return self.scalar_value
+
+
+class _Connection:
+    def __init__(self, *, v2=False):
+        self.v2 = v2
+        self.statements = []
+
+    def execute(self, statement, params=None):
+        sql = str(statement)
+        self.statements.append(sql)
+        if "to_regclass" in sql:
+            return _Result(scalar_value=self.v2)
+        return _Result()
+
+
+class _Engine:
+    def __init__(self, connection):
+        self.connection = connection
+
+    def begin(self):
+        return nullcontext(self.connection)
+
+
+def test_forecast_prepares_effective_links_once(monkeypatch):
+    connection = _Connection()
+    calls = []
+    flags = []
+    monkeypatch.setattr(forecast_service, "engine", _Engine(connection))
+    monkeypatch.setattr(forecast_service, "ensure_revenue_forecast_schema", lambda: None)
+    monkeypatch.setattr(forecast_service, "get_sales_target", lambda **kwargs: {"target_amount": 0})
+    monkeypatch.setattr(document_service, "prepare_effective_document_links", lambda conn: calls.append(conn) or "_effective_document_links")
+    monkeypatch.setattr(forecast_service, "_historical_stage_rates", lambda conn: {})
+
+    def realised(conn, **kwargs):
+        flags.append(kwargs["effective_links_prepared"])
+        return {"amount": 0.0, "count": 0, "opportunity_ids": set(), "items": []}
+
+    def realised_ids(conn, **kwargs):
+        flags.append(kwargs["effective_links_prepared"])
+        return set()
+
+    monkeypatch.setattr(forecast_service, "_realised_for_period", realised)
+    monkeypatch.setattr(forecast_service, "_already_realised_ids", realised_ids)
+
+    forecast_service.get_revenue_forecast()
+
+    assert calls == [connection]
+    assert flags == [True, True]
+
+
+def test_helpers_still_prepare_when_called_in_isolation(monkeypatch):
+    connection = _Connection()
+    calls = []
+    monkeypatch.setattr(document_service, "prepare_effective_document_links", lambda conn: calls.append(conn))
+
+    forecast_service._realised_for_period(
+        connection, metric="invoiced", period_start=date(2026, 8, 1), period_end=date(2026, 8, 31),
+    )
+    forecast_service._already_realised_ids(connection, metric="invoiced")
+
+    assert calls == [connection, connection]
+
+
+def test_set_based_preparation_has_constant_statement_count_and_no_resolver(monkeypatch):
+    monkeypatch.setattr(
+        document_service, "resolve_document_links",
+        lambda *args, **kwargs: (_ for _ in ()).throw(AssertionError("per-opportunity resolver called")),
+    )
+    for v2 in (False, True):
+        connection = _Connection(v2=v2)
+        document_service.prepare_effective_document_links(connection)
+        assert len(connection.statements) == 4
+        assert sum("INSERT INTO _effective_document_links" in sql for sql in connection.statements) == 1
+
+
+def test_set_based_legacy_relationship_case_matches_resolver_vocabulary():
+    source = document_service.prepare_effective_document_links.__code__.co_consts
+    sql_fragments = " ".join(value for value in source if isinstance(value, str))
+    for relationship in ("REMOVED", "PRIMARY", "HISTORICAL", "SECONDARY"):
+        assert relationship in sql_fragments
+    assert "NOT EXISTS" in sql_fragments
+    assert "l.ended_at IS NULL" in sql_fragments