chore: reduce complexity of three C901-flagged functions

posthog/models/person/point_in_time_properties.py

@@ -100,35 +100,7 @@ def get_distinct_ids_for_person_identifier(
     return distinct_ids
 
 
-def build_person_properties_at_time(
-    team_id: int,
-    timestamp: datetime,
-    distinct_ids: list[str],
-    include_set_once: bool = False,
-    timeout: Optional[int] = 30,
-    row_limit: int = DEFAULT_PROPERTY_ROW_LIMIT,
-    lower_bound: Optional[datetime] = None,
-) -> dict[str, Any]:
-    """
-    Build person properties at a specific point in time from ClickHouse events.
-
-    Args:
-        team_id: The team ID to filter events by
-        timestamp: The point in time to build properties at (events after this are ignored)
-        distinct_ids: List of distinct_ids to query for person properties
-        include_set_once: If True, also handles $set_once operations (default: False)
-        timeout: Query timeout in seconds (default: 30)
-        row_limit: Maximum property update rows to ship back from ClickHouse (default 100_000).
-        lower_bound: Optional lower bound for the time range scan. If not provided, defaults to timestamp - 2 years.
-
-    Returns:
-        Dict containing person properties as they existed at the specified timestamp.
-
-    Raises:
-        ValueError: If parameters are invalid
-        Exception: If ClickHouse query fails
-    """
-    # Validation
+def _validate_build_inputs(team_id: int, timestamp: datetime, distinct_ids: list[str], row_limit: int) -> None:
     if not isinstance(team_id, int) or team_id <= 0:
         raise ValueError("team_id must be a positive integer")
 
@@ -144,6 +116,8 @@ def build_person_properties_at_time(
     if not isinstance(row_limit, int) or row_limit <= 0:
         raise ValueError("row_limit must be a positive integer")
 
+
+def _build_property_query(include_set_once: bool, row_limit: int) -> str:
     if include_set_once:
         event_filter = "event IN ('$set', '$set_once') OR JSONHas(properties, '$set')"
     else:
@@ -155,7 +129,7 @@ def build_person_properties_at_time(
     # the property row count answers directly. We extract $set / $set_once raw
     # JSON instead of shipping the full properties blob, and the timestamp
     # window + LIMIT keeps ClickHouse from walking dead partitions.
-    query = f"""
+    return f"""
     SELECT
         JSONExtractRaw(properties, '$set') AS set_json,
         JSONExtractRaw(properties, '$set_once') AS set_once_json,
@@ -170,45 +144,82 @@ def build_person_properties_at_time(
     LIMIT {int(row_limit)}
     """
 
-    # Use provided lower_bound or default to timestamp - 2 years
-    effective_lower_bound = lower_bound if lower_bound is not None else timestamp - _HISTORY_SCAN_FLOOR
-
-    params = {
-        "team_id": team_id,
-        "distinct_ids": distinct_ids,
-        "lower_bound": effective_lower_bound.astimezone(UTC).strftime("%Y-%m-%d %H:%M:%S"),
-        "upper_bound": timestamp.astimezone(UTC).strftime("%Y-%m-%d %H:%M:%S"),
-    }
 
+def _parse_property_json(raw: Any) -> Optional[dict]:
     try:
-        rows = sync_execute(query, params, settings={"max_execution_time": timeout})
-    except Exception as e:
-        raise Exception(f"Failed to query ClickHouse events: {str(e)}") from e
+        parsed = json.loads(raw)
+    except (json.JSONDecodeError, TypeError):
+        return None
+    return parsed if isinstance(parsed, dict) else None
+
 
+def _reconstruct_properties(rows: list, include_set_once: bool) -> dict[str, Any]:
     person_properties: dict[str, Any] = {}
 
     for row in rows:
         set_json, set_once_json, event_name = row
 
         if set_json:
-            try:
-                set_properties = json.loads(set_json)
-            except (json.JSONDecodeError, TypeError):
-                set_properties = None
-
-            if isinstance(set_properties, dict):
+            set_properties = _parse_property_json(set_json)
+            if set_properties is not None:
                 person_properties.update(set_properties)
 
         # $set_once semantics only apply to dedicated $set_once events.
         if include_set_once and event_name == "$set_once" and set_once_json:
-            try:
-                set_once_properties = json.loads(set_once_json)
-            except (json.JSONDecodeError, TypeError):
-                set_once_properties = None
-
-            if isinstance(set_once_properties, dict):
+            set_once_properties = _parse_property_json(set_once_json)
+            if set_once_properties is not None:
                 for key, value in set_once_properties.items():
                     if key not in person_properties:
                         person_properties[key] = value
 
     return person_properties
+
+
+def build_person_properties_at_time(
+    team_id: int,
+    timestamp: datetime,
+    distinct_ids: list[str],
+    include_set_once: bool = False,
+    timeout: Optional[int] = 30,
+    row_limit: int = DEFAULT_PROPERTY_ROW_LIMIT,
+    lower_bound: Optional[datetime] = None,
+) -> dict[str, Any]:
+    """
+    Build person properties at a specific point in time from ClickHouse events.
+
+    Args:
+        team_id: The team ID to filter events by
+        timestamp: The point in time to build properties at (events after this are ignored)
+        distinct_ids: List of distinct_ids to query for person properties
+        include_set_once: If True, also handles $set_once operations (default: False)
+        timeout: Query timeout in seconds (default: 30)
+        row_limit: Maximum property update rows to ship back from ClickHouse (default 100_000).
+        lower_bound: Optional lower bound for the time range scan. If not provided, defaults to timestamp - 2 years.
+
+    Returns:
+        Dict containing person properties as they existed at the specified timestamp.
+
+    Raises:
+        ValueError: If parameters are invalid
+        Exception: If ClickHouse query fails
+    """
+    _validate_build_inputs(team_id, timestamp, distinct_ids, row_limit)
+
+    query = _build_property_query(include_set_once, row_limit)
+
+    # Use provided lower_bound or default to timestamp - 2 years
+    effective_lower_bound = lower_bound if lower_bound is not None else timestamp - _HISTORY_SCAN_FLOOR
+
+    params = {
+        "team_id": team_id,
+        "distinct_ids": distinct_ids,
+        "lower_bound": effective_lower_bound.astimezone(UTC).strftime("%Y-%m-%d %H:%M:%S"),
+        "upper_bound": timestamp.astimezone(UTC).strftime("%Y-%m-%d %H:%M:%S"),
+    }
+
+    try:
+        rows = sync_execute(query, params, settings={"max_execution_time": timeout})
+    except Exception as e:
+        raise Exception(f"Failed to query ClickHouse events: {str(e)}") from e
+
+    return _reconstruct_properties(rows, include_set_once)

products/review_hog/backend/reviewer/progress.py

@@ -233,19 +233,15 @@ class ResolutionRunState:
     needs_attention: int = 0
 
 
-def resolution_states(team_id: int, reports: list[ReviewReport]) -> dict[str, ResolutionRunState]:
-    """Each report's latest resolution run — resolving or died-partway — derived from artefacts.
+def _latest_created_at(queryset: QuerySet) -> dict[str, datetime]:
+    rows = queryset.values_list("report_id").annotate(latest=Max("created_at")).values_list("report_id", "latest")
+    return {str(report_id): latest for report_id, latest in rows}
 
-    The run's `resolution_run` artefact (written at prepare) anchors it: the run's progress is its
-    queued threads' `thread_verdict` rows written since, its completion is a closing run `note`
-    (author `review_hog_resolution`) written since, and its liveness is the report's overall
-    artefact activity against the staleness window — the same signal `_in_progress_report_ids`
-    uses for review turns, so the two can't disagree about "visibly moving".
 
-    A report is absent from the result when it has no resolution run, its latest run completed
-    (closing note present), or a newer review turn superseded it (a `pr_snapshot` written after the
-    run anchor — that turn's own progress takes over the row).
-    """
+def _load_resolution_runs(
+    team_id: int, reports: list[ReviewReport]
+) -> dict[str, tuple[ResolutionRunArtefact, datetime]]:
+    """The latest `resolution_run` anchor per report, dropping unparseable and empty runs."""
     runs: dict[str, tuple[ResolutionRunArtefact, datetime]] = {}
     run_rows = (
         ReviewReportArtefact.objects.for_team(team_id)
@@ -266,35 +262,45 @@ def resolution_states(team_id: int, reports: list[ReviewReport]) -> dict[str, Re
             continue
         if run.total > 0:
             runs[report_id] = (run, row["created_at"])
-    if not runs:
-        return {}
+    return runs
+
 
-    def _latest_after(queryset: QuerySet) -> dict[str, datetime]:
-        rows = queryset.values_list("report_id").annotate(latest=Max("created_at")).values_list("report_id", "latest")
-        return {str(report_id): latest for report_id, latest in rows}
+def _live_resolution_runs(
+    team_id: int, runs: dict[str, tuple[ResolutionRunArtefact, datetime]]
+) -> tuple[dict[str, tuple[ResolutionRunArtefact, datetime]], dict[str, datetime]]:
+    """Keep only runs still in flight, and return each report's latest artefact activity.
 
+    A run drops out when a newer review turn superseded it (a `pr_snapshot` after the run anchor) or
+    it completed (a closing run `note` after the anchor). Activity liveness reuses the same signal
+    `_in_progress_report_ids` uses for review turns, so the two can't disagree about "visibly moving".
+    """
     scoped = ReviewReportArtefact.objects.for_team(team_id).filter(report_id__in=list(runs))
-    snapshot_latest = _latest_after(scoped.filter(type=ReviewReportArtefact.ArtefactType.PR_SNAPSHOT))
-    note_latest = _latest_after(
+    snapshot_latest = _latest_created_at(scoped.filter(type=ReviewReportArtefact.ArtefactType.PR_SNAPSHOT))
+    note_latest = _latest_created_at(
         scoped.filter(type=ReviewReportArtefact.ArtefactType.NOTE)
         .annotate(note_author=KeyTextTransform("author", _content_json()))
         .filter(note_author=RESOLUTION_RUN_NOTE_AUTHOR)
     )
-    activity_latest = _latest_after(scoped.exclude(type=ReviewReportArtefact.ArtefactType.FINDING_OUTCOME))
+    activity_latest = _latest_created_at(scoped.exclude(type=ReviewReportArtefact.ArtefactType.FINDING_OUTCOME))
 
     live: dict[str, tuple[ResolutionRunArtefact, datetime]] = {}
     for report_id, (run, started_at) in runs.items():
         superseded = snapshot_latest.get(report_id) is not None and snapshot_latest[report_id] > started_at
         completed = note_latest.get(report_id) is not None and note_latest[report_id] >= started_at
         if not superseded and not completed:
             live[report_id] = (run, started_at)
-    if not live:
-        return {}
+    return live, activity_latest
 
-    # Latest verdict per thread within each live run (rows come oldest-first, so later rows win) —
-    # scoped to the run's own queued threads, because redelivering a prior run's verdict also
-    # appends rows during this run. Only delivered verdicts (`reply_posted`) count: a judged thread
-    # whose GitHub writes failed has no reply yet, so it must not read as settled.
+
+def _thread_verdicts(
+    team_id: int, live: dict[str, tuple[ResolutionRunArtefact, datetime]]
+) -> dict[str, dict[str, tuple[str, bool]]]:
+    """Latest verdict per thread within each live run (rows come oldest-first, so later rows win).
+
+    Scoped to the run's own queued threads, because redelivering a prior run's verdict also appends
+    rows during this run. Only delivered verdicts (`reply_posted`) count: a judged thread whose
+    GitHub writes failed has no reply yet, so it must not read as settled.
+    """
     verdicts: dict[str, dict[str, tuple[str, bool]]] = {report_id: {} for report_id in live}
     verdict_rows = (
         ReviewReportArtefact.objects.for_team(team_id)
@@ -320,6 +326,30 @@ def _latest_after(queryset: QuerySet) -> dict[str, datetime]:
             verdict_row["outcome"],
             verdict_row["reply_posted"] == "true",
         )
+    return verdicts
+
+
+def resolution_states(team_id: int, reports: list[ReviewReport]) -> dict[str, ResolutionRunState]:
+    """Each report's latest resolution run — resolving or died-partway — derived from artefacts.
+
+    The run's `resolution_run` artefact (written at prepare) anchors it: the run's progress is its
+    queued threads' `thread_verdict` rows written since, its completion is a closing run `note`
+    (author `review_hog_resolution`) written since, and its liveness is the report's overall
+    artefact activity against the staleness window.
+
+    A report is absent from the result when it has no resolution run, its latest run completed
+    (closing note present), or a newer review turn superseded it (a `pr_snapshot` written after the
+    run anchor — that turn's own progress takes over the row).
+    """
+    runs = _load_resolution_runs(team_id, reports)
+    if not runs:
+        return {}
+
+    live, activity_latest = _live_resolution_runs(team_id, runs)
+    if not live:
+        return {}
+
+    verdicts = _thread_verdicts(team_id, live)
 
     cutoff = timezone.now() - IN_PROGRESS_STALE_AFTER
     reports_by_id = {str(report.id): report for report in reports}

products/warehouse_sources/backend/temporal/data_imports/sources/elasticsearch/elasticsearch.py

@@ -1,5 +1,5 @@
 import dataclasses
-from collections.abc import Iterator
+from collections.abc import Callable, Iterator
 from typing import Any, Optional
 from urllib.parse import quote, urlparse
 
@@ -169,16 +169,9 @@ def coerce_float_fields(doc: dict[str, Any], float_paths: set[str]) -> None:
         _coerce_path(doc, path.split("."))
 
 
-def get_rows(
-    host: str,
-    auth: ElasticsearchAuth,
-    index: str,
-    logger: FilteringBoundLogger,
-) -> Iterator[list[dict[str, Any]]]:
-    session = _get_session(auth)
-    base_url = normalize_host(host)
-    float_paths = get_float_field_paths(session, base_url, index)
-
+def _build_post_fn(
+    session: requests.Session, logger: FilteringBoundLogger
+) -> Callable[[str, dict[str, Any]], dict[str, Any]]:
     @retry(
         retry=retry_if_exception_type((ElasticsearchRetryableError, requests.ReadTimeout, requests.ConnectionError)),
         stop=stop_after_attempt(MAX_RETRY_ATTEMPTS),
@@ -199,6 +192,38 @@ def post(url: str, body: dict[str, Any]) -> dict[str, Any]:
 
         return response.json()
 
+    return post
+
+
+def _page_items(data: dict[str, Any], float_paths: set[str]) -> tuple[list, list[dict[str, Any]]]:
+    """Split a scroll response into its raw hits and the coerced `_source` rows."""
+    hits = ((data.get("hits") or {}).get("hits")) or []
+    items = [{**(hit.get("_source") or {}), "_id": hit["_id"]} for hit in hits]
+    if float_paths:
+        for item in items:
+            coerce_float_fields(item, float_paths)
+    return hits, items
+
+
+def _clear_scroll(session: requests.Session, base_url: str, scroll_id: str) -> None:
+    # Best-effort: free the server-side scroll context early.
+    try:
+        session.delete(f"{base_url}/_search/scroll", json={"scroll_id": [scroll_id]}, timeout=10)
+    except Exception:
+        pass
+
+
+def get_rows(
+    host: str,
+    auth: ElasticsearchAuth,
+    index: str,
+    logger: FilteringBoundLogger,
+) -> Iterator[list[dict[str, Any]]]:
+    session = _get_session(auth)
+    base_url = normalize_host(host)
+    float_paths = get_float_field_paths(session, base_url, index)
+    post = _build_post_fn(session, logger)
+
     # The scroll API gives a stable snapshot of the index for the duration of
     # the walk; scroll ids expire after SCROLL_KEEPALIVE of inactivity, so the
     # walk restarts from scratch on retry rather than persisting state.
@@ -210,12 +235,7 @@ def post(url: str, body: dict[str, Any]) -> dict[str, Any]:
 
     try:
         while True:
-            hits = ((data.get("hits") or {}).get("hits")) or []
-            items = [{**(hit.get("_source") or {}), "_id": hit["_id"]} for hit in hits]
-
-            if float_paths:
-                for item in items:
-                    coerce_float_fields(item, float_paths)
+            hits, items = _page_items(data, float_paths)
 
             if items:
                 yield items
@@ -227,15 +247,7 @@ def post(url: str, body: dict[str, Any]) -> dict[str, Any]:
             scroll_id = data.get("_scroll_id", scroll_id)
     finally:
         if scroll_id:
-            # Best-effort: free the server-side scroll context early.
-            try:
-                session.delete(
-                    f"{base_url}/_search/scroll",
-                    json={"scroll_id": [scroll_id]},
-                    timeout=10,
-                )
-            except Exception:
-                pass
+            _clear_scroll(session, base_url, scroll_id)
 
 
 def elasticsearch_source(