From fe467d837e4be72069756e9f72b7ddbf3aa0159e Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 25 Aug 2026 12:00:49 +0530 Subject: [PATCH 01/14] UN-3973 Derive monthly metrics from the daily tier, narrow source window to 2 days The dashboard aggregation widened its DAY-granularity query to the first of the previous month so monthly buckets could be summed in Python from the same rows. Every run re-read 32-62 days of source data per metric, per org, 96 times a day. Monthly is now rolled up from event_metrics_daily in one statement for all orgs, so the source queries only need the daily window. That window drops to 2 days, sized against the measured worst created_at -> terminal-status lag of ~2h. A once-daily 7-day pass reruns the same task at a wider bound to repair gaps left by cron downtime. The active-org prefilter is decoupled from the daily window and pinned at 7 days: metrics filtered on another column (hitl_completions on approved_at) can land for an org whose executions are older than the source window. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- .../0004_add_reconciliation_task.py | 61 ++++++ backend/dashboard_metrics/tasks.py | 194 +++++++++--------- backend/dashboard_metrics/tests/test_tasks.py | 148 ++++++++++++- 3 files changed, 307 insertions(+), 96 deletions(-) create mode 100644 backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py diff --git a/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py b/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py new file mode 100644 index 0000000000..0e05394687 --- /dev/null +++ b/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py @@ -0,0 +1,61 @@ +"""Data migration to schedule the daily-tier reconciliation pass. + +The 15-minute aggregation reads a narrow source window, which cannot repair +gaps left by cron downtime. This runs the same task once a day at a wider +window to backfill them. +""" + +from django.db import migrations + +RECONCILE_TASK_NAME = "dashboard_metrics_reconcile_source_window" + + +def create_reconciliation_task(apps, schema_editor): + """Create the once-daily reconciliation periodic task.""" + CrontabSchedule = apps.get_model("django_celery_beat", "CrontabSchedule") + PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") + + # 4:00 AM UTC — clear of the 2:00 and 3:00 cleanup tasks + schedule_4am, _ = CrontabSchedule.objects.get_or_create( + minute="0", + hour="4", + day_of_week="*", + day_of_month="*", + month_of_year="*", + defaults={"timezone": "UTC"}, + ) + + PeriodicTask.objects.update_or_create( + name=RECONCILE_TASK_NAME, + defaults={ + "task": "dashboard_metrics.aggregate_from_sources", + "crontab": schedule_4am, + "queue": "dashboard_metric_events", + "kwargs": '{"source_window_days": 7}', + "enabled": True, + "description": ( + "Re-aggregate metrics over a 7 day source window to repair " + "daily-tier gaps left by cron downtime" + ), + }, + ) + + +def remove_reconciliation_task(apps, schema_editor): + """Remove the reconciliation periodic task on rollback.""" + PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") + PeriodicTask.objects.filter(name=RECONCILE_TASK_NAME).delete() + + +class Migration(migrations.Migration): + dependencies = [ + ("dashboard_metrics", "0003_alter_eventmetricsdaily_organization_and_more"), + ("django_celery_beat", "0018_improve_crontab_helptext"), + ] + + operations = [ + migrations.RunPython( + create_reconciliation_task, + remove_reconciliation_task, + ), + ] diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 181c985137..3c7b5fe9cd 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -8,12 +8,14 @@ import logging import time -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta from typing import Any from account_v2.models import Organization from celery import shared_task from django.core.cache import cache +from django.db.models import Min, Sum +from django.db.models.functions import TruncMonth from django.db.utils import DatabaseError, OperationalError from django.utils import timezone from workflow_manager.workflow_v2.models.execution import WorkflowExecution @@ -33,6 +35,21 @@ DASHBOARD_HOURLY_METRICS_RETENTION_DAYS = 30 DASHBOARD_DAILY_METRICS_RETENTION_DAYS = 365 +# Source lookback for the daily tier. Sized against the measured worst +# created_at -> terminal-status lag of ~2h, bounded by the file processing +# time limit and the stuck-execution reaper. +DASHBOARD_SOURCE_WINDOW_DAYS = 2 + +# Wider lookback used by the once-daily reconciliation pass, which repairs +# the daily tier after cron downtime. +DASHBOARD_RECONCILE_WINDOW_DAYS = 7 + +# Lookback for the active-org prefilter. Deliberately independent of the +# source window: metrics filtered on a column other than +# WorkflowExecution.created_at (e.g. hitl_completions on approved_at) can +# land for an org whose executions are older than the source window. +DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS = 7 + def _upsert_agg(agg: dict, key: tuple, metric_type: str, value: float) -> None: """Add a value to an aggregation dict, creating the entry if needed.""" @@ -165,32 +182,47 @@ def _bulk_upsert_daily(aggregations: dict) -> int: return len(objects) -def _bulk_upsert_monthly(aggregations: dict) -> int: - """Bulk upsert monthly aggregations using INSERT ... ON CONFLICT. +def _rollup_monthly_from_daily(month_start: date) -> int: + """Derive monthly metrics by summing the daily tier from month_start onwards. - Uses _base_manager to bypass DefaultOrganizationManagerMixin. + Replaces per-org monthly queries against source tables with a single + aggregate over event_metrics_daily, which retains far more history than + the monthly window needs. + + metric_type is aggregated rather than grouped: it is not part of + unique_monthly_metric, so grouping on it could yield two rows for one + conflict target. Args: - aggregations: Dict keyed by (org_id, month_str, metric_name, project, tag) + month_start: First day of the earliest month to rebuild Returns: Number of rows upserted """ - objects = [] - for key, agg in aggregations.items(): - org_id, month_str, metric_name, project, tag = key - objects.append( - EventMetricsMonthly( - organization_id=org_id, - month=datetime.fromisoformat(month_str).date(), - metric_name=metric_name, - project=project, - tag=tag, - metric_type=agg["metric_type"], - metric_value=agg["value"], - metric_count=agg["count"], - ) + rows = ( + EventMetricsDaily._base_manager.filter(date__gte=month_start) + .annotate(month=TruncMonth("date")) + .values("organization_id", "month", "metric_name", "project", "tag") + .annotate( + value=Sum("metric_value"), + count=Sum("metric_count"), + mtype=Min("metric_type"), + ) + ) + + objects = [ + EventMetricsMonthly( + organization_id=row["organization_id"], + month=row["month"], + metric_name=row["metric_name"], + project=row["project"], + tag=row["tag"], + metric_type=row["mtype"], + metric_value=row["value"], + metric_count=row["count"], ) + for row in rows + ] if not objects: return 0 @@ -260,7 +292,9 @@ def _acquire_aggregation_lock() -> bool: retry_backoff=True, retry_backoff_max=300, ) -def aggregate_metrics_from_sources() -> dict[str, Any]: +def aggregate_metrics_from_sources( + source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, +) -> dict[str, Any]: """Aggregate metrics from source tables into hourly, daily, and monthly tables. This task runs periodically (every 15 minutes) to query metrics from @@ -274,8 +308,13 @@ def aggregate_metrics_from_sources() -> dict[str, Any]: Aggregation windows: - Hourly: Last 24 hours (rolling window) - - Daily: Last 7 days (ensures we capture late-arriving data) - - Monthly: Last 2 months (current + previous month) + - Daily: source_window_days (covers late-arriving data) + - Monthly: Rolled up from the daily tier, current + previous month + + Args: + source_window_days: Daily-tier source lookback. The once-daily + reconciliation pass runs the same task at + DASHBOARD_RECONCILE_WINDOW_DAYS to repair gaps after downtime. Returns: Dict with aggregation summary for all three tiers @@ -285,7 +324,7 @@ def aggregate_metrics_from_sources() -> dict[str, Any]: return {"success": True, "skipped": True, "reason": "lock_held"} try: - return _run_aggregation() + return _run_aggregation(source_window_days) finally: cache.delete(AGGREGATION_LOCK_KEY) @@ -297,18 +336,15 @@ def _aggregate_single_metric( org_id: str, hourly_start: datetime, daily_start: datetime, - monthly_start: datetime, end_date: datetime, hourly_agg: dict, daily_agg: dict, - monthly_agg: dict, extra_kwargs: dict | None = None, ) -> None: - """Run a single metric query at all 3 granularities and populate agg dicts. + """Run a single metric query at hourly and daily granularity. - Uses 2 queries instead of 3: the daily query is widened to monthly_start - and its results are split into both daily_agg and monthly_agg in Python. - This is the same pattern proven in the backfill management command. + Monthly totals are derived separately by rolling up the daily tier, so + neither query reaches further back than daily_start. """ extra_kwargs = extra_kwargs or {} @@ -324,43 +360,32 @@ def _aggregate_single_metric( key = (org_id, hour_ts.isoformat(), metric_name, "default", "") _upsert_agg(hourly_agg, key, metric_type, row["value"] or 0) - # === DAILY + MONTHLY (single query from monthly_start) === + # === DAILY === for row in query_method( org_id, - monthly_start, + daily_start, end_date, granularity=Granularity.DAY, **extra_kwargs, ): - value = row["value"] or 0 day_ts = _truncate_to_day(row["period"]) - - if day_ts >= daily_start: - key = (org_id, day_ts.date().isoformat(), metric_name, "default", "") - _upsert_agg(daily_agg, key, metric_type, value) - - month_key = _truncate_to_month(row["period"]).date().isoformat() - key = (org_id, month_key, metric_name, "default", "") - _upsert_agg(monthly_agg, key, metric_type, value) + key = (org_id, day_ts.date().isoformat(), metric_name, "default", "") + _upsert_agg(daily_agg, key, metric_type, row["value"] or 0) def _aggregate_llm_combined( org_id: str, hourly_start: datetime, daily_start: datetime, - monthly_start: datetime, end_date: datetime, hourly_agg: dict, daily_agg: dict, - monthly_agg: dict, llm_combined_fields: dict, ) -> None: - """Run the combined LLM metrics query at all granularities. + """Run the combined LLM metrics query at hourly and daily granularity. - Issues 2 queries total (hourly + daily/monthly) instead of 3. - The DAY-granularity query is widened to monthly_start and results are - split into daily_agg (recent rows) and monthly_agg (all rows bucketed - by month) in Python. Same pattern as _aggregate_single_metric. + Issues 2 queries covering 4 metrics. Same windowing as + _aggregate_single_metric — monthly is derived from the daily tier. """ # === HOURLY (last 24h) === for row in MetricsQueryService.get_llm_metrics_combined( @@ -374,28 +399,22 @@ def _aggregate_llm_combined( key = (org_id, ts_str, metric_name, "default", "") _upsert_agg(hourly_agg, key, metric_type, row[field] or 0) - # === DAILY + MONTHLY (single query from monthly_start) === + # === DAILY === for row in MetricsQueryService.get_llm_metrics_combined( org_id, - monthly_start, + daily_start, end_date, granularity=Granularity.DAY, ): - day_ts = _truncate_to_day(row["period"]) - month_key = _truncate_to_month(row["period"]).date().isoformat() - + day_str = _truncate_to_day(row["period"]).date().isoformat() for field, (metric_name, metric_type) in llm_combined_fields.items(): - value = row[field] or 0 + key = (org_id, day_str, metric_name, "default", "") + _upsert_agg(daily_agg, key, metric_type, row[field] or 0) - if day_ts >= daily_start: - key = (org_id, day_ts.date().isoformat(), metric_name, "default", "") - _upsert_agg(daily_agg, key, metric_type, value) - key = (org_id, month_key, metric_name, "default", "") - _upsert_agg(monthly_agg, key, metric_type, value) - - -def _run_aggregation() -> dict[str, Any]: +def _run_aggregation( + source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, +) -> dict[str, Any]: """Execute the actual aggregation logic. Separated from the task function to keep the lock management clean. @@ -404,25 +423,13 @@ def _run_aggregation() -> dict[str, Any]: # Query windows for each granularity # - Hourly: Last 24 hours (rolling window, matches retention of 30 days) - # - Daily: Last 7 days (ensures we capture late-arriving data) - # - Monthly: Last 2 months (current + previous, ensures month transitions are captured) + # - Daily: source_window_days of source data + # - Monthly: rolled up from the daily tier, current + previous month hourly_start = end_date - timedelta(hours=24) - daily_start = _truncate_to_day(end_date - timedelta(days=7)) - # Include previous month to handle month boundaries - if end_date.month == 1: - monthly_start = end_date.replace( - year=end_date.year - 1, - month=12, - day=1, - hour=0, - minute=0, - second=0, - microsecond=0, - ) - else: - monthly_start = end_date.replace( - month=end_date.month - 1, day=1, hour=0, minute=0, second=0, microsecond=0 - ) + daily_start = _truncate_to_day(end_date - timedelta(days=source_window_days)) + monthly_start = _truncate_to_month( + _truncate_to_month(end_date) - timedelta(days=1) + ).date() # Metric definitions: (name, query_method, is_histogram) # Note: llm_calls, challenges, summarization_calls, and llm_usage are @@ -459,15 +466,13 @@ def _run_aggregation() -> dict[str, Any]: "orgs_processed": 0, } - # Pre-filter to orgs with recent activity to reduce DB load. - # Uses daily_start (7 days) instead of monthly_start (2 months) because: - # - Hourly/daily queries only need recent data (24h / 7d windows) - # - Monthly totals for dormant orgs were already written by previous - # runs when the org was active — re-running just overwrites same values - # - This avoids 28 queries per dormant org that had activity 2-8 weeks ago + # Pre-filter to orgs with recent activity to reduce DB load. Kept wider + # than the source window because some metrics are filtered on a column + # other than WorkflowExecution.created_at (hitl_completions uses + # approved_at) and can land for an org whose executions are older. active_org_ids = set( WorkflowExecution.objects.filter( - created_at__gte=daily_start, + created_at__gte=end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), ) .values_list("workflow__organization_id", flat=True) .distinct() @@ -499,7 +504,6 @@ def _run_aggregation() -> dict[str, Any]: org_identifier = org.organization_id # Pre-resolved for PageUsage queries hourly_agg: dict[tuple, dict] = {} daily_agg: dict[tuple, dict] = {} - monthly_agg: dict[tuple, dict] = {} try: for metric_name, query_method, is_histogram in metric_configs: @@ -519,11 +523,9 @@ def _run_aggregation() -> dict[str, Any]: org_id, hourly_start, daily_start, - monthly_start, end_date, hourly_agg, daily_agg, - monthly_agg, extra_kwargs, ) except Exception: @@ -536,33 +538,35 @@ def _run_aggregation() -> dict[str, Any]: org_id, hourly_start, daily_start, - monthly_start, end_date, hourly_agg, daily_agg, - monthly_agg, llm_combined_fields, ) except Exception: logger.exception("Error querying combined LLM metrics for org %s", org_id) stats["errors"] += 1 - # Bulk upsert all three tiers (single INSERT...ON CONFLICT each) + # Bulk upsert both tiers (single INSERT...ON CONFLICT each) if hourly_agg: stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) if daily_agg: stats["daily"]["upserted"] += _bulk_upsert_daily(daily_agg) - if monthly_agg: - stats["monthly"]["upserted"] += _bulk_upsert_monthly(monthly_agg) - stats["orgs_processed"] += 1 except Exception: logger.exception("Error processing org %s", org_id) stats["errors"] += 1 + # Monthly is derived from the daily tier in one statement for all orgs + try: + stats["monthly"]["upserted"] = _rollup_monthly_from_daily(monthly_start) + except Exception: + logger.exception("Error rolling up monthly metrics from %s", monthly_start) + stats["errors"] += 1 + logger.info( f"Aggregation completed: {stats['orgs_processed']} orgs, " f"hourly={stats['hourly']['upserted']}, " diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index 03ef136508..c80c38303c 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -1,6 +1,7 @@ """Unit tests for Dashboard Metrics Celery tasks.""" -from datetime import datetime, timedelta +from datetime import date, datetime, timedelta +from unittest.mock import patch from django.test import TestCase from django.utils import timezone @@ -9,12 +10,18 @@ from dashboard_metrics.models import ( EventMetricsDaily, EventMetricsHourly, + EventMetricsMonthly, MetricType, ) from dashboard_metrics.tasks import ( + DASHBOARD_RECONCILE_WINDOW_DAYS, + DASHBOARD_SOURCE_WINDOW_DAYS, + _rollup_monthly_from_daily, + _run_aggregation, _truncate_to_day, _truncate_to_hour, _truncate_to_month, + aggregate_metrics_from_sources, cleanup_daily_metrics, cleanup_hourly_metrics, ) @@ -198,3 +205,142 @@ def test_cleanup_no_records_to_delete(self): assert result["success"] is True assert result["deleted"] == 0 + + +class TestMonthlyRollup(TestCase): + """Tests for deriving monthly metrics from the daily tier.""" + + def setUp(self): + """Set up test fixtures.""" + self.org = Organization.objects.create( + organization_id="rollup-org", name="rollup-org", display_name="Rollup Org" + ) + + def _daily(self, day, value, count=1, metric_type=MetricType.COUNTER): + """Create a daily metric row for the fixture org.""" + EventMetricsDaily.objects.create( + organization=self.org, + date=day, + metric_name="documents_processed", + metric_type=metric_type, + metric_value=value, + metric_count=count, + project="default", + ) + + def _monthly_rows(self): + """Read back monthly rows ordered by month.""" + return list(EventMetricsMonthly._base_manager.order_by("month")) + + def test_sums_daily_rows_into_month_bucket(self): + """Daily rows within a month sum into a single monthly row.""" + self._daily(date(2024, 3, 5), value=10, count=2) + self._daily(date(2024, 3, 18), value=32, count=4) + + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].month == date(2024, 3, 1) + assert rows[0].metric_value == 42 + assert rows[0].metric_count == 6 + + def test_month_boundary_keeps_months_separate(self): + """Rows spanning the 1st land in two months without bleeding.""" + self._daily(date(2024, 1, 30), value=5) + self._daily(date(2024, 1, 31), value=7) + self._daily(date(2024, 2, 1), value=100) + self._daily(date(2024, 2, 2), value=200) + + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 2 + + rows = self._monthly_rows() + assert [r.month for r in rows] == [date(2024, 1, 1), date(2024, 2, 1)] + assert [r.metric_value for r in rows] == [12, 300] + + def test_excludes_months_before_the_window(self): + """Daily rows older than month_start are not rolled up.""" + self._daily(date(2023, 12, 15), value=999) + self._daily(date(2024, 1, 15), value=5) + + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 1 + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].month == date(2024, 1, 1) + + def test_rerun_overwrites_instead_of_accumulating(self): + """A second rollup replaces the monthly total rather than doubling it.""" + self._daily(date(2024, 3, 5), value=10, count=2) + _rollup_monthly_from_daily(date(2024, 3, 1)) + + self._daily(date(2024, 3, 6), value=5, count=1) + _rollup_monthly_from_daily(date(2024, 3, 1)) + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].metric_value == 15 + assert rows[0].metric_count == 3 + + def test_mixed_metric_type_within_a_month_yields_one_row(self): + """metric_type is aggregated, so it cannot split one conflict target.""" + self._daily(date(2024, 3, 5), value=10, metric_type=MetricType.HISTOGRAM) + self._daily(date(2024, 3, 6), value=5, metric_type=MetricType.COUNTER) + + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].metric_value == 15 + + def test_no_daily_rows_upserts_nothing(self): + """An empty daily tier is a no-op, not an error.""" + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 0 + assert not self._monthly_rows() + + +class TestSourceWindow(TestCase): + """Tests for the per-run source window and the reconciliation pass.""" + + def setUp(self): + """Set up test fixtures.""" + self.org = Organization.objects.create( + organization_id="window-org", name="window-org", display_name="Window Org" + ) + + def _run_with_active_org(self, **kwargs): + """Run aggregation with the active-org prefilter stubbed to the fixture org.""" + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [self.org.id] + return _run_aggregation(**kwargs) + + def test_default_window_bounds_the_daily_query(self): + """The per-run daily window is DASHBOARD_SOURCE_WINDOW_DAYS wide.""" + result = self._run_with_active_org() + + expected = _truncate_to_day( + timezone.now() - timedelta(days=DASHBOARD_SOURCE_WINDOW_DAYS) + ) + assert result["period"]["daily"]["start"] == expected.isoformat() + + def test_reconciliation_window_widens_the_daily_query(self): + """The reconciliation pass reaches further back on the same code path.""" + result = self._run_with_active_org( + source_window_days=DASHBOARD_RECONCILE_WINDOW_DAYS + ) + + expected = _truncate_to_day( + timezone.now() - timedelta(days=DASHBOARD_RECONCILE_WINDOW_DAYS) + ) + assert result["period"]["daily"]["start"] == expected.isoformat() + + def test_task_passes_the_window_through(self): + """The scheduled task forwards its kwarg, defaulting to the per-run window.""" + with patch("dashboard_metrics.tasks._run_aggregation") as mock_run: + aggregate_metrics_from_sources() + mock_run.assert_called_once_with(DASHBOARD_SOURCE_WINDOW_DAYS) + + with patch("dashboard_metrics.tasks._run_aggregation") as mock_run: + aggregate_metrics_from_sources(source_window_days=7) + mock_run.assert_called_once_with(7) From 3ea08b721360ec5fc90f347bfbe91e11912f51f6 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 25 Aug 2026 20:26:30 +0530 Subject: [PATCH 02/14] UN-3973 Address Sonar and Greptile review findings Sonar: - S117: rename apps.get_model() locals in 0004 to snake_case - S3776: cut _run_aggregation cognitive complexity from 22 by hoisting the static metric config tables to module level and extracting the per-org body, the active-org prefilter and the result shape into helpers Greptile: - Monthly rows in the rebuilt window whose daily rows are gone are now deleted alongside the upsert, so the two tiers cannot disagree. An empty daily tier still short-circuits, so a wiped tier cannot cascade into deleting monthly history. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- .../0004_add_reconciliation_task.py | 12 +- backend/dashboard_metrics/tasks.py | 342 +++++++++++------- backend/dashboard_metrics/tests/test_tasks.py | 26 ++ 3 files changed, 249 insertions(+), 131 deletions(-) diff --git a/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py b/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py index 0e05394687..09667e867f 100644 --- a/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py +++ b/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py @@ -12,11 +12,11 @@ def create_reconciliation_task(apps, schema_editor): """Create the once-daily reconciliation periodic task.""" - CrontabSchedule = apps.get_model("django_celery_beat", "CrontabSchedule") - PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") + crontab_model = apps.get_model("django_celery_beat", "CrontabSchedule") + periodic_task_model = apps.get_model("django_celery_beat", "PeriodicTask") # 4:00 AM UTC — clear of the 2:00 and 3:00 cleanup tasks - schedule_4am, _ = CrontabSchedule.objects.get_or_create( + schedule_4am, _ = crontab_model.objects.get_or_create( minute="0", hour="4", day_of_week="*", @@ -25,7 +25,7 @@ def create_reconciliation_task(apps, schema_editor): defaults={"timezone": "UTC"}, ) - PeriodicTask.objects.update_or_create( + periodic_task_model.objects.update_or_create( name=RECONCILE_TASK_NAME, defaults={ "task": "dashboard_metrics.aggregate_from_sources", @@ -43,8 +43,8 @@ def create_reconciliation_task(apps, schema_editor): def remove_reconciliation_task(apps, schema_editor): """Remove the reconciliation periodic task on rollback.""" - PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") - PeriodicTask.objects.filter(name=RECONCILE_TASK_NAME).delete() + periodic_task_model = apps.get_model("django_celery_beat", "PeriodicTask") + periodic_task_model.objects.filter(name=RECONCILE_TASK_NAME).delete() class Migration(migrations.Migration): diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 3c7b5fe9cd..d01b9b5418 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -14,6 +14,7 @@ from account_v2.models import Organization from celery import shared_task from django.core.cache import cache +from django.db import transaction from django.db.models import Min, Sum from django.db.models.functions import TruncMonth from django.db.utils import DatabaseError, OperationalError @@ -182,12 +183,48 @@ def _bulk_upsert_daily(aggregations: dict) -> int: return len(objects) +def _delete_orphan_monthly(month_start: date, fresh_keys: set[tuple]) -> int: + """Drop monthly rows in the window that the rollup no longer produces. + + Monthly is a pure derivation of the daily tier, so a key whose daily rows + have gone (a deleted workflow cascades to its executions) must not survive + as a stale total. + + Args: + month_start: First day of the earliest month being rebuilt + fresh_keys: Keys the current rollup produced + + Returns: + Number of rows deleted + """ + stale_pks = [ + row["pk"] + for row in EventMetricsMonthly._base_manager.filter( + month__gte=month_start + ).values("pk", "organization_id", "month", "metric_name", "project", "tag") + if ( + row["organization_id"], + row["month"], + row["metric_name"], + row["project"], + row["tag"], + ) + not in fresh_keys + ] + if not stale_pks: + return 0 + + deleted, _ = EventMetricsMonthly._base_manager.filter(pk__in=stale_pks).delete() + return deleted + + def _rollup_monthly_from_daily(month_start: date) -> int: """Derive monthly metrics by summing the daily tier from month_start onwards. Replaces per-org monthly queries against source tables with a single aggregate over event_metrics_daily, which retains far more history than - the monthly window needs. + the monthly window needs. Rows in the window that the daily tier no longer + backs are removed, so the two tiers cannot disagree. metric_type is aggregated rather than grouped: it is not part of unique_monthly_metric, so grouping on it could yield two rows for one @@ -224,15 +261,24 @@ def _rollup_monthly_from_daily(month_start: date) -> int: for row in rows ] + # An empty daily tier means the source of truth is gone, not that every + # month is genuinely zero, so leave the existing rows alone. if not objects: return 0 - EventMetricsMonthly._base_manager.bulk_create( - objects, - update_conflicts=True, - unique_fields=["organization", "month", "metric_name", "project", "tag"], - update_fields=["metric_type", "metric_value", "metric_count"], - ) + fresh_keys = { + (o.organization_id, o.month, o.metric_name, o.project, o.tag) for o in objects + } + + with transaction.atomic(): + EventMetricsMonthly._base_manager.bulk_create( + objects, + update_conflicts=True, + unique_fields=["organization", "month", "metric_name", "project", "tag"], + update_fields=["metric_type", "metric_value", "metric_count"], + ) + _delete_orphan_monthly(month_start, fresh_keys) + return len(objects) @@ -412,6 +458,156 @@ def _aggregate_llm_combined( _upsert_agg(daily_agg, key, metric_type, row[field] or 0) +# Metric definitions: (name, query_method, is_histogram) +# Note: llm_calls, challenges, summarization_calls, and llm_usage are +# handled separately via get_llm_metrics_combined (1 query instead of 4). +METRIC_CONFIGS = [ + ("documents_processed", MetricsQueryService.get_documents_processed, False), + ("pages_processed", MetricsQueryService.get_pages_processed, True), + ("deployed_api_requests", MetricsQueryService.get_deployed_api_requests, False), + ("etl_pipeline_executions", MetricsQueryService.get_etl_pipeline_executions, False), + ("prompt_executions", MetricsQueryService.get_prompt_executions, False), + ("failed_pages", MetricsQueryService.get_failed_pages, True), + ("hitl_reviews", MetricsQueryService.get_hitl_reviews, False), + ("hitl_completions", MetricsQueryService.get_hitl_completions, False), +] + +# LLM metrics combined via conditional aggregation (4 metrics in 1 query). +# Maps combined query field -> (metric_name, metric_type) +LLM_COMBINED_FIELDS = { + "llm_calls": ("llm_calls", MetricType.COUNTER), + "challenges": ("challenges", MetricType.COUNTER), + "summarization_calls": ("summarization_calls", MetricType.COUNTER), + "llm_usage": ("llm_usage", MetricType.HISTOGRAM), +} + + +def _collect_org_metrics( + org: Organization, + hourly_start: datetime, + daily_start: datetime, + end_date: datetime, +) -> tuple[dict, dict, int]: + """Query every metric for one organization into hourly/daily aggregates. + + A failing metric is logged and counted, leaving the rest to proceed. + + Returns: + Tuple of (hourly aggregations, daily aggregations, error count) + """ + org_id = str(org.id) + hourly_agg: dict[tuple, dict] = {} + daily_agg: dict[tuple, dict] = {} + errors = 0 + + for metric_name, query_method, is_histogram in METRIC_CONFIGS: + metric_type = MetricType.HISTOGRAM if is_histogram else MetricType.COUNTER + # Pre-resolved identifier spares PageUsage an Organization lookup per call. + extra_kwargs = ( + {"org_identifier": org.organization_id} + if metric_name == "pages_processed" + else {} + ) + try: + _aggregate_single_metric( + query_method, + metric_name, + metric_type, + org_id, + hourly_start, + daily_start, + end_date, + hourly_agg, + daily_agg, + extra_kwargs, + ) + except Exception: + logger.exception("Error querying %s for org %s", metric_name, org_id) + errors += 1 + + try: + _aggregate_llm_combined( + org_id, + hourly_start, + daily_start, + end_date, + hourly_agg, + daily_agg, + LLM_COMBINED_FIELDS, + ) + except Exception: + logger.exception("Error querying combined LLM metrics for org %s", org_id) + errors += 1 + + return hourly_agg, daily_agg, errors + + +def _aggregate_org( + org: Organization, + hourly_start: datetime, + daily_start: datetime, + end_date: datetime, + stats: dict[str, Any], +) -> None: + """Aggregate one organization and upsert its hourly and daily tiers.""" + hourly_agg, daily_agg, errors = _collect_org_metrics( + org, hourly_start, daily_start, end_date + ) + stats["errors"] += errors + + # Bulk upsert both tiers (single INSERT...ON CONFLICT each) + if hourly_agg: + stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) + + if daily_agg: + stats["daily"]["upserted"] += _bulk_upsert_daily(daily_agg) + + stats["orgs_processed"] += 1 + + +def _active_org_ids(end_date: datetime) -> set: + """Organizations with recent execution activity. + + Deliberately wider than the source window: metrics filtered on a column + other than WorkflowExecution.created_at (hitl_completions uses approved_at) + can land for an org whose executions are older. + """ + return set( + WorkflowExecution.objects.filter( + created_at__gte=end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), + ) + .values_list("workflow__organization_id", flat=True) + .distinct() + ) + + +def _build_result( + stats: dict[str, Any], + hourly_start: datetime, + daily_start: datetime, + monthly_start: date, + end_date: datetime, + skipped_reason: str | None = None, +) -> dict[str, Any]: + """Shape the task's return value from the accumulated stats.""" + result = { + "success": True, + "organizations_processed": stats["orgs_processed"], + "hourly": stats["hourly"], + "daily": stats["daily"], + "monthly": stats["monthly"], + "errors": stats["errors"], + "period": { + "hourly": {"start": hourly_start.isoformat(), "end": end_date.isoformat()}, + "daily": {"start": daily_start.isoformat(), "end": end_date.isoformat()}, + "monthly": {"start": monthly_start.isoformat(), "end": end_date.isoformat()}, + }, + } + if skipped_reason: + result["skipped_reason"] = skipped_reason + return result + + def _run_aggregation( source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, ) -> dict[str, Any]: @@ -431,33 +627,6 @@ def _run_aggregation( _truncate_to_month(end_date) - timedelta(days=1) ).date() - # Metric definitions: (name, query_method, is_histogram) - # Note: llm_calls, challenges, summarization_calls, and llm_usage are - # handled separately via get_llm_metrics_combined (1 query instead of 4). - metric_configs = [ - ("documents_processed", MetricsQueryService.get_documents_processed, False), - ("pages_processed", MetricsQueryService.get_pages_processed, True), - ("deployed_api_requests", MetricsQueryService.get_deployed_api_requests, False), - ( - "etl_pipeline_executions", - MetricsQueryService.get_etl_pipeline_executions, - False, - ), - ("prompt_executions", MetricsQueryService.get_prompt_executions, False), - ("failed_pages", MetricsQueryService.get_failed_pages, True), - ("hitl_reviews", MetricsQueryService.get_hitl_reviews, False), - ("hitl_completions", MetricsQueryService.get_hitl_completions, False), - ] - - # LLM metrics combined via conditional aggregation (4 metrics in 1 query). - # Maps combined query field -> (metric_name, metric_type) - llm_combined_fields = { - "llm_calls": ("llm_calls", MetricType.COUNTER), - "challenges": ("challenges", MetricType.COUNTER), - "summarization_calls": ("summarization_calls", MetricType.COUNTER), - "llm_usage": ("llm_usage", MetricType.HISTOGRAM), - } - stats = { "hourly": {"upserted": 0}, "daily": {"upserted": 0}, @@ -466,98 +635,33 @@ def _run_aggregation( "orgs_processed": 0, } - # Pre-filter to orgs with recent activity to reduce DB load. Kept wider - # than the source window because some metrics are filtered on a column - # other than WorkflowExecution.created_at (hitl_completions uses - # approved_at) and can land for an org whose executions are older. - active_org_ids = set( - WorkflowExecution.objects.filter( - created_at__gte=end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), - ) - .values_list("workflow__organization_id", flat=True) - .distinct() - ) - total_orgs = Organization.objects.count() + # Pre-filter to orgs with recent activity to reduce DB load. + active_org_ids = _active_org_ids(end_date) logger.info( "Aggregation: %d active orgs out of %d total", len(active_org_ids), - total_orgs, + Organization.objects.count(), ) if not active_org_ids: - return { - "success": True, - "organizations_processed": 0, - "hourly": stats["hourly"], - "daily": stats["daily"], - "monthly": stats["monthly"], - "errors": 0, - "skipped_reason": "no_active_orgs", - } + return _build_result( + stats, + hourly_start, + daily_start, + monthly_start, + end_date, + skipped_reason="no_active_orgs", + ) organizations = Organization.objects.filter(id__in=active_org_ids).only( "id", "organization_id" ) for org in organizations: - org_id = str(org.id) - org_identifier = org.organization_id # Pre-resolved for PageUsage queries - hourly_agg: dict[tuple, dict] = {} - daily_agg: dict[tuple, dict] = {} - try: - for metric_name, query_method, is_histogram in metric_configs: - metric_type = MetricType.HISTOGRAM if is_histogram else MetricType.COUNTER - - # Pass org_identifier to PageUsage-based metrics to - # avoid redundant Organization lookups per call. - extra_kwargs = {} - if metric_name == "pages_processed": - extra_kwargs["org_identifier"] = org_identifier - - try: - _aggregate_single_metric( - query_method, - metric_name, - metric_type, - org_id, - hourly_start, - daily_start, - end_date, - hourly_agg, - daily_agg, - extra_kwargs, - ) - except Exception: - logger.exception("Error querying %s for org %s", metric_name, org_id) - stats["errors"] += 1 - - # Combined LLM metrics: 1 query per granularity instead of 4 - try: - _aggregate_llm_combined( - org_id, - hourly_start, - daily_start, - end_date, - hourly_agg, - daily_agg, - llm_combined_fields, - ) - except Exception: - logger.exception("Error querying combined LLM metrics for org %s", org_id) - stats["errors"] += 1 - - # Bulk upsert both tiers (single INSERT...ON CONFLICT each) - if hourly_agg: - stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) - - if daily_agg: - stats["daily"]["upserted"] += _bulk_upsert_daily(daily_agg) - - stats["orgs_processed"] += 1 - + _aggregate_org(org, hourly_start, daily_start, end_date, stats) except Exception: - logger.exception("Error processing org %s", org_id) + logger.exception("Error processing org %s", org.id) stats["errors"] += 1 # Monthly is derived from the daily tier in one statement for all orgs @@ -575,19 +679,7 @@ def _run_aggregation( f"errors={stats['errors']}" ) - return { - "success": True, - "organizations_processed": stats["orgs_processed"], - "hourly": stats["hourly"], - "daily": stats["daily"], - "monthly": stats["monthly"], - "errors": stats["errors"], - "period": { - "hourly": {"start": hourly_start.isoformat(), "end": end_date.isoformat()}, - "daily": {"start": daily_start.isoformat(), "end": end_date.isoformat()}, - "monthly": {"start": monthly_start.isoformat(), "end": end_date.isoformat()}, - }, - } + return _build_result(stats, hourly_start, daily_start, monthly_start, end_date) @shared_task( diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index c80c38303c..682a0f7246 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -298,6 +298,32 @@ def test_no_daily_rows_upserts_nothing(self): assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 0 assert not self._monthly_rows() + def test_monthly_row_is_dropped_once_its_daily_rows_are_gone(self): + """A month whose daily rows were deleted must not keep a stale total.""" + self._daily(date(2024, 3, 5), value=10) + self._daily(date(2024, 4, 5), value=7) + _rollup_monthly_from_daily(date(2024, 3, 1)) + assert len(self._monthly_rows()) == 2 + + EventMetricsDaily._base_manager.filter(date=date(2024, 3, 5)).delete() + _rollup_monthly_from_daily(date(2024, 3, 1)) + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].month == date(2024, 4, 1) + + def test_months_before_the_window_are_left_alone(self): + """Orphan cleanup must not reach outside the rebuilt window.""" + self._daily(date(2024, 1, 10), value=99) + _rollup_monthly_from_daily(date(2024, 1, 1)) + EventMetricsDaily._base_manager.all().delete() + + self._daily(date(2024, 3, 5), value=10) + _rollup_monthly_from_daily(date(2024, 3, 1)) + + months = [row.month for row in self._monthly_rows()] + assert months == [date(2024, 1, 1), date(2024, 3, 1)] + class TestSourceWindow(TestCase): """Tests for the per-run source window and the reconciliation pass.""" From 395393465b95bd8ca97c4a17e702ac548769fd7d Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Fri, 28 Aug 2026 19:30:01 +0530 Subject: [PATCH 03/14] UN-3973 Add tests covering the source window, rollup SQL and reconciliation schedule Closes the acceptance criteria that had no automated check: - the monthly rollup issues no source-table SQL, asserted by capturing the queries it actually sends - the window ladder at 2 / 7 / 62 days, including a row that finishes after the narrow window has moved past its created_at and so never re-enters it - the reconciliation schedule row, its idempotency and its reverse The schedule tests call the migration's function directly. The suite runs with --no-migrations, so data migrations never execute and asserting on the beat row would fail regardless of the migration being correct. Also moves the dotenv load in settings/base.py above the Celery block. CELERY_BROKER_BASE_URL, _USER and _PASS were read above it, so they could not be supplied by an env file at all and had to be ambient. Ambient values still take precedence, so deployed behaviour is unchanged. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/backend/settings/base.py | 8 +- backend/dashboard_metrics/tests/test_tasks.py | 160 ++++++++++++++++++ 2 files changed, 164 insertions(+), 4 deletions(-) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index d14f87b304..f0dbafe6d2 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -55,6 +55,10 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | # Load default log from env DEFAULT_LOG_LEVEL = os.environ.get("DEFAULT_LOG_LEVEL", "INFO") +ENV_FILE = find_dotenv() +if ENV_FILE: + load_dotenv(ENV_FILE) + # Celery Broker Configuration CELERY_BROKER_BASE_URL = get_required_setting("CELERY_BROKER_BASE_URL") CELERY_BROKER_USER = get_required_setting("CELERY_BROKER_USER") @@ -65,10 +69,6 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | ) ) -ENV_FILE = find_dotenv() -if ENV_FILE: - load_dotenv(ENV_FILE) - # Loading environment variables WORKFLOW_ACTION_EXPIRATION_TIME_IN_SECOND = os.environ.get( diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index 682a0f7246..6cb9316d1c 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -1,10 +1,16 @@ """Unit tests for Dashboard Metrics Celery tasks.""" +import json from datetime import date, datetime, timedelta +from importlib import import_module from unittest.mock import patch +from django.apps import apps +from django.db import connection from django.test import TestCase +from django.test.utils import CaptureQueriesContext from django.utils import timezone +from django_celery_beat.models import PeriodicTask from account_v2.models import Organization from dashboard_metrics.models import ( @@ -13,6 +19,10 @@ EventMetricsMonthly, MetricType, ) +from workflow_manager.file_execution.models import WorkflowFileExecution +from workflow_manager.workflow_v2.enums import ExecutionStatus +from workflow_manager.workflow_v2.models.execution import WorkflowExecution +from workflow_manager.workflow_v2.models.workflow import Workflow from dashboard_metrics.tasks import ( DASHBOARD_RECONCILE_WINDOW_DAYS, DASHBOARD_SOURCE_WINDOW_DAYS, @@ -325,6 +335,37 @@ def test_months_before_the_window_are_left_alone(self): assert months == [date(2024, 1, 1), date(2024, 3, 1)] +class TestRollupQueryShape(TestCase): + """The monthly rollup must not read the raw source tables.""" + + def test_monthly_rollup_never_touches_source_tables(self): + """This is the saving: monthly reads the daily tier and nothing else.""" + EventMetricsDaily._base_manager.create( + organization=Organization.objects.create( + organization_id="shape-org", name="shape", display_name="Shape" + ), + date=date(2024, 3, 5), + metric_name="documents_processed", + metric_type=MetricType.COUNTER, + metric_value=10, + metric_count=2, + project="default", + tag="", + ) + + with CaptureQueriesContext(connection) as captured: + _rollup_monthly_from_daily(date(2024, 3, 1)) + + sql = " ".join(q["sql"] for q in captured.captured_queries).lower() + assert "event_metrics_daily" in sql + for source_table in ( + "workflow_file_execution", + "workflow_execution", + "page_usage", + ): + assert source_table not in sql, f"monthly rollup read {source_table}" + + class TestSourceWindow(TestCase): """Tests for the per-run source window and the reconciliation pass.""" @@ -370,3 +411,122 @@ def test_task_passes_the_window_through(self): with patch("dashboard_metrics.tasks._run_aggregation") as mock_run: aggregate_metrics_from_sources(source_window_days=7) mock_run.assert_called_once_with(7) + + def _seed_file( + self, days_ago: int, status: ExecutionStatus = ExecutionStatus.COMPLETED + ) -> date: + """Seed one file execution dated days_ago, return its date.""" + workflow = Workflow.objects.create( + workflow_name=f"recon-wf-{days_ago}", organization=self.org + ) + execution = WorkflowExecution.objects.create( + workflow=workflow, status=ExecutionStatus.COMPLETED + ) + file_execution = WorkflowFileExecution.objects.create( + workflow_execution=execution, + file_name="a.pdf", + status=status.value, + ) + + stamp = timezone.now() - timedelta(days=days_ago) + # created_at is auto_now_add; a queryset update is what bypasses it + WorkflowFileExecution.objects.filter(pk=file_execution.pk).update( + created_at=stamp + ) + WorkflowExecution.objects.filter(pk=execution.pk).update(created_at=stamp) + return stamp.date() + + def test_reconciliation_recovers_a_day_the_narrow_window_missed(self): + """A row outside the per-run window is picked up by the wider pass.""" + day = self._seed_file(days_ago=5) + + _run_aggregation() + assert not EventMetricsDaily._base_manager.filter(date=day).exists() + + result = _run_aggregation(source_window_days=DASHBOARD_RECONCILE_WINDOW_DAYS) + + row = EventMetricsDaily._base_manager.get( + date=day, metric_name="documents_processed" + ) + assert row.metric_value == 1 + assert result["errors"] == 0 + + def test_late_terminal_status_does_not_re_enter_the_narrow_window(self): + """Finishing after the window moved on does not bring a row back.""" + day = self._seed_file(days_ago=3, status=ExecutionStatus.PENDING) + + # Still running: nothing to count yet. + _run_aggregation() + assert not EventMetricsDaily._base_manager.filter(date=day).exists() + + # It finishes. status turns terminal; created_at does not move. + WorkflowFileExecution.objects.update(status=ExecutionStatus.COMPLETED.value) + + # The per-run window no longer reaches its created_at, so it stays missed. + _run_aggregation() + assert not EventMetricsDaily._base_manager.filter(date=day).exists() + + # Only the wider pass recovers it. + _run_aggregation(source_window_days=DASHBOARD_RECONCILE_WINDOW_DAYS) + assert EventMetricsDaily._base_manager.filter( + date=day, metric_name="documents_processed" + ).exists() + + def test_gap_older_than_the_reconcile_window_needs_a_manual_backfill(self): + """Neither scheduled pass reaches a day beyond the reconcile window.""" + old_day = self._seed_file(days_ago=62) + recent_day = self._seed_file(days_ago=0) + + _run_aggregation() + _run_aggregation(source_window_days=DASHBOARD_RECONCILE_WINDOW_DAYS) + + # The run worked — it just cannot reach that far back. + assert EventMetricsDaily._base_manager.filter(date=recent_day).exists() + assert not EventMetricsDaily._base_manager.filter(date=old_day).exists() + + +class TestReconciliationSchedule(TestCase): + """Migration 0004 schedules the once-daily reconciliation pass. + + The suite runs with --no-migrations, so the migration's function is called + directly rather than relying on it having been applied. + """ + + def setUp(self): + """Load the data migration module.""" + self.migration = import_module( + "dashboard_metrics.migrations.0004_add_reconciliation_task" + ) + + def _task(self): + return PeriodicTask.objects.get(name=self.migration.RECONCILE_TASK_NAME) + + def test_migration_schedules_the_pass_at_0400_with_a_7_day_window(self): + """The beat row lands enabled, at 04:00 UTC, carrying the wider window.""" + self.migration.create_reconciliation_task(apps, None) + + task = self._task() + assert task.task == "dashboard_metrics.aggregate_from_sources" + assert task.enabled + assert task.queue == "dashboard_metric_events" + assert json.loads(task.kwargs) == { + "source_window_days": DASHBOARD_RECONCILE_WINDOW_DAYS + } + assert (task.crontab.hour, task.crontab.minute) == ("4", "0") + + def test_migration_is_idempotent_and_reversible(self): + """Re-running leaves one row; the reverse function removes it.""" + self.migration.create_reconciliation_task(apps, None) + self.migration.create_reconciliation_task(apps, None) + + assert ( + PeriodicTask.objects.filter( + name=self.migration.RECONCILE_TASK_NAME + ).count() + == 1 + ) + + self.migration.remove_reconciliation_task(apps, None) + assert not PeriodicTask.objects.filter( + name=self.migration.RECONCILE_TASK_NAME + ).exists() From 86ed04c5b0e770cf651cdae80d90287dc9d76fe5 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Sat, 29 Aug 2026 13:00:06 +0530 Subject: [PATCH 04/14] UN-3973 Trim comments in tasks.py and revert the unrelated settings change MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Cut the verbose comments and docstrings down to the purpose and the non-obvious bits. Code is unchanged. Restore backend/settings/base.py to main — moving the dotenv load ahead of get_required_setting was a local test convenience, not part of this change. The test rig exports the broker vars itself, so CI never needed it. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/backend/settings/base.py | 8 +-- backend/dashboard_metrics/tasks.py | 96 +++++++----------------------- 2 files changed, 26 insertions(+), 78 deletions(-) diff --git a/backend/backend/settings/base.py b/backend/backend/settings/base.py index f0dbafe6d2..d14f87b304 100644 --- a/backend/backend/settings/base.py +++ b/backend/backend/settings/base.py @@ -55,10 +55,6 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | # Load default log from env DEFAULT_LOG_LEVEL = os.environ.get("DEFAULT_LOG_LEVEL", "INFO") -ENV_FILE = find_dotenv() -if ENV_FILE: - load_dotenv(ENV_FILE) - # Celery Broker Configuration CELERY_BROKER_BASE_URL = get_required_setting("CELERY_BROKER_BASE_URL") CELERY_BROKER_USER = get_required_setting("CELERY_BROKER_USER") @@ -69,6 +65,10 @@ def get_required_setting(setting_key: str, default: str | None = None) -> str | ) ) +ENV_FILE = find_dotenv() +if ENV_FILE: + load_dotenv(ENV_FILE) + # Loading environment variables WORKFLOW_ACTION_EXPIRATION_TIME_IN_SECOND = os.environ.get( diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index d01b9b5418..c0b1b1a9b8 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -36,19 +36,15 @@ DASHBOARD_HOURLY_METRICS_RETENTION_DAYS = 30 DASHBOARD_DAILY_METRICS_RETENTION_DAYS = 365 -# Source lookback for the daily tier. Sized against the measured worst -# created_at -> terminal-status lag of ~2h, bounded by the file processing -# time limit and the stuck-execution reaper. +# Daily-tier source lookback, sized against the worst observed +# created_at -> terminal-status lag. DASHBOARD_SOURCE_WINDOW_DAYS = 2 -# Wider lookback used by the once-daily reconciliation pass, which repairs -# the daily tier after cron downtime. +# Wider lookback for the once-daily reconciliation pass. DASHBOARD_RECONCILE_WINDOW_DAYS = 7 -# Lookback for the active-org prefilter. Deliberately independent of the -# source window: metrics filtered on a column other than -# WorkflowExecution.created_at (e.g. hitl_completions on approved_at) can -# land for an org whose executions are older than the source window. +# Wider than the source window: metrics keyed on another column +# (e.g. approved_at) can land for an org whose executions are older. DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS = 7 @@ -184,18 +180,10 @@ def _bulk_upsert_daily(aggregations: dict) -> int: def _delete_orphan_monthly(month_start: date, fresh_keys: set[tuple]) -> int: - """Drop monthly rows in the window that the rollup no longer produces. + """Drop monthly rows from month_start that the rollup no longer produces. - Monthly is a pure derivation of the daily tier, so a key whose daily rows - have gone (a deleted workflow cascades to its executions) must not survive - as a stale total. - - Args: - month_start: First day of the earliest month being rebuilt - fresh_keys: Keys the current rollup produced - - Returns: - Number of rows deleted + Monthly derives from the daily tier, so a key with no daily rows left must + not survive as a stale total. """ stale_pks = [ row["pk"] @@ -219,22 +207,11 @@ def _delete_orphan_monthly(month_start: date, fresh_keys: set[tuple]) -> int: def _rollup_monthly_from_daily(month_start: date) -> int: - """Derive monthly metrics by summing the daily tier from month_start onwards. - - Replaces per-org monthly queries against source tables with a single - aggregate over event_metrics_daily, which retains far more history than - the monthly window needs. Rows in the window that the daily tier no longer - backs are removed, so the two tiers cannot disagree. + """Sum the daily tier from month_start into monthly, for all orgs at once. metric_type is aggregated rather than grouped: it is not part of unique_monthly_metric, so grouping on it could yield two rows for one conflict target. - - Args: - month_start: First day of the earliest month to rebuild - - Returns: - Number of rows upserted """ rows = ( EventMetricsDaily._base_manager.filter(date__gte=month_start) @@ -261,8 +238,8 @@ def _rollup_monthly_from_daily(month_start: date) -> int: for row in rows ] - # An empty daily tier means the source of truth is gone, not that every - # month is genuinely zero, so leave the existing rows alone. + # An empty daily tier means the source is gone, not that every month is + # zero — leave existing rows alone. if not objects: return 0 @@ -341,25 +318,14 @@ def _acquire_aggregation_lock() -> bool: def aggregate_metrics_from_sources( source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, ) -> dict[str, Any]: - """Aggregate metrics from source tables into hourly, daily, and monthly tables. - - This task runs periodically (every 15 minutes) to query metrics from - source tables (Usage, PageUsage, WorkflowExecution, etc.) and aggregate - them into EventMetricsHourly, EventMetricsDaily, and EventMetricsMonthly - tables for fast dashboard queries at different granularities. - - Uses a Redis distributed lock with self-healing to prevent overlapping - runs. If a previous run was killed without releasing the lock, the next - run detects the stale lock and reclaims it automatically. + """Aggregate source tables into the hourly, daily and monthly tiers. - Aggregation windows: - - Hourly: Last 24 hours (rolling window) - - Daily: source_window_days (covers late-arriving data) - - Monthly: Rolled up from the daily tier, current + previous month + Runs every 15 minutes under a self-healing Redis lock. Hourly covers the + last 24h, daily the source window, monthly is rolled up from daily. Args: source_window_days: Daily-tier source lookback. The once-daily - reconciliation pass runs the same task at + reconciliation pass reruns this task at DASHBOARD_RECONCILE_WINDOW_DAYS to repair gaps after downtime. Returns: @@ -387,11 +353,7 @@ def _aggregate_single_metric( daily_agg: dict, extra_kwargs: dict | None = None, ) -> None: - """Run a single metric query at hourly and daily granularity. - - Monthly totals are derived separately by rolling up the daily tier, so - neither query reaches further back than daily_start. - """ + """Run a single metric query at hourly and daily granularity.""" extra_kwargs = extra_kwargs or {} # === HOURLY (last 24h) === @@ -430,8 +392,7 @@ def _aggregate_llm_combined( ) -> None: """Run the combined LLM metrics query at hourly and daily granularity. - Issues 2 queries covering 4 metrics. Same windowing as - _aggregate_single_metric — monthly is derived from the daily tier. + Two queries covering four metrics. """ # === HOURLY (last 24h) === for row in MetricsQueryService.get_llm_metrics_combined( @@ -488,7 +449,7 @@ def _collect_org_metrics( daily_start: datetime, end_date: datetime, ) -> tuple[dict, dict, int]: - """Query every metric for one organization into hourly/daily aggregates. + """Query every metric for one org into hourly/daily aggregates. A failing metric is logged and counted, leaving the rest to proceed. @@ -502,7 +463,7 @@ def _collect_org_metrics( for metric_name, query_method, is_histogram in METRIC_CONFIGS: metric_type = MetricType.HISTOGRAM if is_histogram else MetricType.COUNTER - # Pre-resolved identifier spares PageUsage an Organization lookup per call. + # Pre-resolved identifier spares PageUsage a lookup per call. extra_kwargs = ( {"org_identifier": org.organization_id} if metric_name == "pages_processed" @@ -555,7 +516,6 @@ def _aggregate_org( ) stats["errors"] += errors - # Bulk upsert both tiers (single INSERT...ON CONFLICT each) if hourly_agg: stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) @@ -566,12 +526,7 @@ def _aggregate_org( def _active_org_ids(end_date: datetime) -> set: - """Organizations with recent execution activity. - - Deliberately wider than the source window: metrics filtered on a column - other than WorkflowExecution.created_at (hitl_completions uses approved_at) - can land for an org whose executions are older. - """ + """Organizations with execution activity in the prefilter lookback.""" return set( WorkflowExecution.objects.filter( created_at__gte=end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), @@ -611,16 +566,10 @@ def _build_result( def _run_aggregation( source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, ) -> dict[str, Any]: - """Execute the actual aggregation logic. - - Separated from the task function to keep the lock management clean. - """ + """Execute the aggregation, separately from the task's lock handling.""" end_date = timezone.now() - # Query windows for each granularity - # - Hourly: Last 24 hours (rolling window, matches retention of 30 days) - # - Daily: source_window_days of source data - # - Monthly: rolled up from the daily tier, current + previous month + # Monthly spans the current and previous month. hourly_start = end_date - timedelta(hours=24) daily_start = _truncate_to_day(end_date - timedelta(days=source_window_days)) monthly_start = _truncate_to_month( @@ -664,7 +613,6 @@ def _run_aggregation( logger.exception("Error processing org %s", org.id) stats["errors"] += 1 - # Monthly is derived from the daily tier in one statement for all orgs try: stats["monthly"]["upserted"] = _rollup_monthly_from_daily(monthly_start) except Exception: From 1e81c88e7dbef30b9f9ed3961ff7897b7a0b8436 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Mon, 31 Aug 2026 20:11:27 +0530 Subject: [PATCH 05/14] UN-3974 [PERF] Split the dashboard metrics schedule by tier and index workflow_execution on created_at MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Schedule split. One schedule ran every 15 minutes and wrote all three metric tiers. Dashboard daily and monthly figures do not need 15-minute freshness, so they move to hourly — 96 runs a day becomes 24 for the expensive DAY-granularity half of the work, while the hourly tier keeps its cadence. Both schedule rows point at the same task and differ only in a `tier` kwarg; a second task name would need its own worker registration and internal endpoint for the PG path. The lock is now keyed per tier, so the two runs that collide at the top of every hour do not starve each other. Omitting `tier` still writes all three tiers, so a manual trigger never silently writes nothing. Prefilter index. The active-org prefilter measures 1,849ms per call on production — the slowest single query on the instance. Nothing on workflow_execution leads with created_at: the two composite indexes are date-ordered only within one workflow or pipeline, and the partial index is empty in steady state. The split raises this query's call count, and UN-4045 will leave three more metric queries on the same bare date-range shape, so the index lands with the split rather than after it. Built CONCURRENTLY with atomic = False and guarded against a leftover INVALID index, matching migration 0026. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/internal_views.py | 9 +- .../0005_split_aggregation_schedule.py | 149 ++++++++++++++ backend/dashboard_metrics/tasks.py | 187 +++++++++++++----- .../migrations/0029_we_created_at_idx.py | 110 +++++++++++ .../workflow_v2/models/execution.py | 7 + workers/scheduler/dashboard_metrics_tasks.py | 12 +- 6 files changed, 416 insertions(+), 58 deletions(-) create mode 100644 backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py create mode 100644 backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index f776633944..913699f127 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -91,10 +91,17 @@ class AggregateMetricsAPIView(_MetricsTaskAPIView): Calls the Celery task body verbatim, Redis lock included — this endpoint exists only because the PG consumer has no Django, not to change what the job does. + + Optional ``tier`` in the body selects which metric tiers to write; omitting it + writes all of them, matching the task's own default. An unrecognised value is a + 400 rather than a silent no-op. """ def post(self, request: Request) -> Response: - return self._run(aggregate_metrics_from_sources) + tier = request.data.get("tier") if isinstance(request.data, dict) else None + if tier is None: + return self._run(aggregate_metrics_from_sources) + return self._run(aggregate_metrics_from_sources, tier=tier) class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView): diff --git a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py new file mode 100644 index 0000000000..0843b8089d --- /dev/null +++ b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py @@ -0,0 +1,149 @@ +"""Split the metrics aggregation into two schedules by tier (UN-3974). + +Before this, one schedule ran every 15 minutes and wrote all three tiers. Dashboard +daily and monthly figures do not need 15-minute freshness, so they move to hourly: +96 runs a day becomes 24 for the expensive DAY-granularity half of the work, while the +hourly tier keeps its 15-minute cadence. + +Both rows point at the SAME task (``dashboard_metrics.aggregate_from_sources``) and +differ only in ``tier`` kwargs. A second task name would need its own worker-side +registration and internal endpoint for the PG path; a kwarg needs neither. + +**Beat and PG rows are declared together here, from one spec.** ``0002_setup_periodic_tasks`` +(Beat) and ``0004_pg_periodic_tasks`` (PG) declare the same schedules in two places, and +``tests/test_pg_periodic_task_declarations.py`` exists to catch them drifting apart. One +spec written twice by the same function cannot drift, so this migration needs no such +guard. ``AGGREGATION_SCHEDULES`` is module-level so a future test can import it. + +Rows land consistent with how each scheduler expects them: + +* Beat ``kwargs`` is a JSON *string*; ``PgPeriodicTask.task_kwargs`` is a JSONField, so + it is stored decoded. +* The new PG row lands **inert** (``pg_owned=False``, ``next_run_at=NULL``) for the same + reason as ``0004`` — the PG scheduler skips rows it does not own, and a NULL + ``next_run_at`` records a baseline next tick rather than firing a catch-up burst. + +Reverse restores the pre-split state: the new rows are deleted and the aggregate row's +kwargs are cleared, putting it back to writing all three tiers every 15 minutes. +""" + +import json + +from django.db import migrations + +AGGREGATE_TASK_NAME = "dashboard_metrics.aggregate_from_sources" +AGGREGATE_QUEUE = "dashboard_metric_events" + +# Frozen literals — migrations must not import app enums. Kept in step with +# dashboard_metrics.tasks.AggregationTier. +TIER_HOURLY = "hourly" +TIER_DAILY_MONTHLY = "daily_monthly" + +# The row that already exists (created by 0002 / 0004); only its kwargs and +# description change, its every-15-minutes schedule does not. +EXISTING_AGGREGATE_ROW = "dashboard_metrics_aggregate_from_sources" + +AGGREGATION_SCHEDULES = [ + { + "name": EXISTING_AGGREGATE_ROW, + "tier": TIER_HOURLY, + "cron_string": "*/15 * * * *", + "crontab": {"minute": "*/15", "hour": "*"}, + "description": ( + "Aggregate the hourly dashboard metrics tier from source tables " + "(Usage, PageUsage, WorkflowExecution, etc.)" + ), + "exists": True, + }, + { + "name": "dashboard_metrics_aggregate_daily_monthly", + "tier": TIER_DAILY_MONTHLY, + "cron_string": "0 * * * *", + "crontab": {"minute": "0", "hour": "*"}, + "description": ( + "Aggregate the daily and monthly dashboard metrics tiers from source " + "tables — hourly, since these figures do not need 15-minute freshness" + ), + "exists": False, + }, +] + + +def split_schedules(apps, schema_editor): + CrontabSchedule = apps.get_model("django_celery_beat", "CrontabSchedule") + PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") + PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") + + for spec in AGGREGATION_SCHEDULES: + kwargs = {"tier": spec["tier"]} + + if spec["exists"]: + # Keep the existing IntervalSchedule — only the payload changes. + PeriodicTask.objects.filter(name=spec["name"]).update( + kwargs=json.dumps(kwargs), description=spec["description"] + ) + else: + schedule, _ = CrontabSchedule.objects.get_or_create( + minute=spec["crontab"]["minute"], + hour=spec["crontab"]["hour"], + day_of_week="*", + day_of_month="*", + month_of_year="*", + defaults={"timezone": "UTC"}, + ) + PeriodicTask.objects.update_or_create( + name=spec["name"], + defaults={ + "task": AGGREGATE_TASK_NAME, + "crontab": schedule, + "queue": AGGREGATE_QUEUE, + "kwargs": json.dumps(kwargs), + "enabled": True, + "description": spec["description"], + }, + ) + + PgPeriodicTask.objects.update_or_create( + name=spec["name"], + defaults={ + "task_name": AGGREGATE_TASK_NAME, + "queue": AGGREGATE_QUEUE, + "task_args": [], + "task_kwargs": kwargs, + "cron_string": spec["cron_string"], + "org_id": "", + "enabled": True, + "pg_owned": False, + }, + ) + + +def merge_schedules(apps, schema_editor): + """Restore the single every-15-minutes row that writes all three tiers.""" + PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") + PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") + + added = [s["name"] for s in AGGREGATION_SCHEDULES if not s["exists"]] + PeriodicTask.objects.filter(name__in=added).delete() + PgPeriodicTask.objects.filter(name__in=added).delete() + + PeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update( + kwargs="{}", + description=( + "Aggregate metrics from source tables (Usage, PageUsage, etc.) " + "into hourly, daily, and monthly metrics tables" + ), + ) + PgPeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update(task_kwargs={}) + + +class Migration(migrations.Migration): + dependencies = [ + ("dashboard_metrics", "0004_pg_periodic_tasks"), + ("django_celery_beat", "0018_improve_crontab_helptext"), + ("pg_queue", "0003_pgperiodictask"), + ] + + operations = [ + migrations.RunPython(split_schedules, merge_schedules), + ] diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 181c985137..30862fde7c 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -9,6 +9,7 @@ import logging import time from datetime import datetime, timedelta +from enum import StrEnum from typing import Any from account_v2.models import Organization @@ -204,49 +205,82 @@ def _bulk_upsert_monthly(aggregations: dict) -> int: return len(objects) -AGGREGATION_LOCK_KEY = "dashboard_metrics:aggregation_lock" -AGGREGATION_LOCK_TIMEOUT = 900 # 15 minutes (matches task schedule) +class AggregationTier(StrEnum): + """Which metric tiers a single aggregation run writes. + The hourly tier needs 15-minute freshness; the daily and monthly tiers do not, + so they run on separate schedules. Daily and monthly stay together because one + DAY-granularity query feeds both. + """ + + HOURLY = "hourly" + DAILY_MONTHLY = "daily_monthly" + ALL = "all" + + +# Per-tier so the 15-minute hourly run and the hourly daily/monthly run — which +# collide at the top of every hour — do not starve each other. +AGGREGATION_LOCK_KEY_PREFIX = "dashboard_metrics:aggregation_lock" +AGGREGATION_LOCK_TIMEOUT = 900 # 15 minutes (matches the fastest task schedule) + + +def _writes_hourly(tier: AggregationTier) -> bool: + return tier in (AggregationTier.HOURLY, AggregationTier.ALL) + + +def _writes_daily_monthly(tier: AggregationTier) -> bool: + return tier in (AggregationTier.DAILY_MONTHLY, AggregationTier.ALL) + + +def _aggregation_lock_key(tier: AggregationTier) -> str: + return f"{AGGREGATION_LOCK_KEY_PREFIX}:{tier.value}" -def _acquire_aggregation_lock() -> bool: + +def _acquire_aggregation_lock(lock_key: str) -> bool: """Acquire the distributed aggregation lock with self-healing. Stores a Unix timestamp as the lock value. If a previous run crashed (OOM kill, SIGKILL) without releasing the lock, the next run detects that the lock is older than AGGREGATION_LOCK_TIMEOUT and reclaims it. + Args: + lock_key: Cache key to lock on — one per tier, see _aggregation_lock_key + Returns: True if lock was acquired, False if another run is legitimately active. """ now = time.time() # Fast path: lock is free - if cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT): + if cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT): return True # Lock exists — check if it's stale (previous run died without releasing) - lock_value = cache.get(AGGREGATION_LOCK_KEY) + lock_value = cache.get(lock_key) if lock_value is None: # Expired between our check and get — lock is now free, try to acquire it - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) try: lock_time = float(lock_value) except (TypeError, ValueError): # Corrupted value (e.g. old "running" string) — reclaim it - logger.warning("Reclaiming aggregation lock with invalid value: %s", lock_value) - cache.delete(AGGREGATION_LOCK_KEY) - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + logger.warning( + "Reclaiming aggregation lock %s with invalid value: %s", lock_key, lock_value + ) + cache.delete(lock_key) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) age = now - lock_time if age > AGGREGATION_LOCK_TIMEOUT: logger.warning( - "Reclaiming stale aggregation lock (age=%.0fs, timeout=%ds)", + "Reclaiming stale aggregation lock %s (age=%.0fs, timeout=%ds)", + lock_key, age, AGGREGATION_LOCK_TIMEOUT, ) - cache.delete(AGGREGATION_LOCK_KEY) - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + cache.delete(lock_key) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) return False @@ -260,13 +294,14 @@ def _acquire_aggregation_lock() -> bool: retry_backoff=True, retry_backoff_max=300, ) -def aggregate_metrics_from_sources() -> dict[str, Any]: - """Aggregate metrics from source tables into hourly, daily, and monthly tables. +def aggregate_metrics_from_sources( + tier: str = AggregationTier.ALL, +) -> dict[str, Any]: + """Aggregate metrics from source tables into the hourly/daily/monthly tables. - This task runs periodically (every 15 minutes) to query metrics from - source tables (Usage, PageUsage, WorkflowExecution, etc.) and aggregate - them into EventMetricsHourly, EventMetricsDaily, and EventMetricsMonthly - tables for fast dashboard queries at different granularities. + Two schedules call this with different tiers: the hourly tier every 15 minutes, + the daily and monthly tiers hourly. Each tier locks separately so the two never + block each other. Uses a Redis distributed lock with self-healing to prevent overlapping runs. If a previous run was killed without releasing the lock, the next @@ -277,17 +312,33 @@ def aggregate_metrics_from_sources() -> dict[str, Any]: - Daily: Last 7 days (ensures we capture late-arriving data) - Monthly: Last 2 months (current + previous month) + Args: + tier: Which tiers to write, an AggregationTier value. Defaults to all, + so a caller that omits it gets the pre-split behaviour rather than + silently writing nothing. + Returns: - Dict with aggregation summary for all three tiers + Dict with aggregation summary for the tiers that ran + + Raises: + ValueError: tier is not a recognised AggregationTier """ - if not _acquire_aggregation_lock(): - logger.info("Skipping aggregation — another run is in progress") - return {"success": True, "skipped": True, "reason": "lock_held"} + tier = AggregationTier(tier) + lock_key = _aggregation_lock_key(tier) + + if not _acquire_aggregation_lock(lock_key): + logger.info("Skipping %s aggregation — another run is in progress", tier.value) + return { + "success": True, + "skipped": True, + "reason": "lock_held", + "tier": tier.value, + } try: - return _run_aggregation() + return _run_aggregation(tier) finally: - cache.delete(AGGREGATION_LOCK_KEY) + cache.delete(lock_key) def _aggregate_single_metric( @@ -302,29 +353,36 @@ def _aggregate_single_metric( hourly_agg: dict, daily_agg: dict, monthly_agg: dict, + tier: AggregationTier, extra_kwargs: dict | None = None, ) -> None: - """Run a single metric query at all 3 granularities and populate agg dicts. + """Run a single metric query at the requested granularities and populate agg dicts. Uses 2 queries instead of 3: the daily query is widened to monthly_start and its results are split into both daily_agg and monthly_agg in Python. This is the same pattern proven in the backfill management command. + + Each query is skipped when its tier is not in scope for this run. """ extra_kwargs = extra_kwargs or {} # === HOURLY (last 24h) === - for row in query_method( - org_id, - hourly_start, - end_date, - granularity=Granularity.HOUR, - **extra_kwargs, - ): - hour_ts = _truncate_to_hour(row["period"]) - key = (org_id, hour_ts.isoformat(), metric_name, "default", "") - _upsert_agg(hourly_agg, key, metric_type, row["value"] or 0) + if _writes_hourly(tier): + for row in query_method( + org_id, + hourly_start, + end_date, + granularity=Granularity.HOUR, + **extra_kwargs, + ): + hour_ts = _truncate_to_hour(row["period"]) + key = (org_id, hour_ts.isoformat(), metric_name, "default", "") + _upsert_agg(hourly_agg, key, metric_type, row["value"] or 0) # === DAILY + MONTHLY (single query from monthly_start) === + if not _writes_daily_monthly(tier): + return + for row in query_method( org_id, monthly_start, @@ -354,8 +412,9 @@ def _aggregate_llm_combined( daily_agg: dict, monthly_agg: dict, llm_combined_fields: dict, + tier: AggregationTier, ) -> None: - """Run the combined LLM metrics query at all granularities. + """Run the combined LLM metrics query at the requested granularities. Issues 2 queries total (hourly + daily/monthly) instead of 3. The DAY-granularity query is widened to monthly_start and results are @@ -363,18 +422,22 @@ def _aggregate_llm_combined( by month) in Python. Same pattern as _aggregate_single_metric. """ # === HOURLY (last 24h) === - for row in MetricsQueryService.get_llm_metrics_combined( - org_id, - hourly_start, - end_date, - granularity=Granularity.HOUR, - ): - ts_str = _truncate_to_hour(row["period"]).isoformat() - for field, (metric_name, metric_type) in llm_combined_fields.items(): - key = (org_id, ts_str, metric_name, "default", "") - _upsert_agg(hourly_agg, key, metric_type, row[field] or 0) + if _writes_hourly(tier): + for row in MetricsQueryService.get_llm_metrics_combined( + org_id, + hourly_start, + end_date, + granularity=Granularity.HOUR, + ): + ts_str = _truncate_to_hour(row["period"]).isoformat() + for field, (metric_name, metric_type) in llm_combined_fields.items(): + key = (org_id, ts_str, metric_name, "default", "") + _upsert_agg(hourly_agg, key, metric_type, row[field] or 0) # === DAILY + MONTHLY (single query from monthly_start) === + if not _writes_daily_monthly(tier): + return + for row in MetricsQueryService.get_llm_metrics_combined( org_id, monthly_start, @@ -395,7 +458,7 @@ def _aggregate_llm_combined( _upsert_agg(monthly_agg, key, metric_type, value) -def _run_aggregation() -> dict[str, Any]: +def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, Any]: """Execute the actual aggregation logic. Separated from the task function to keep the lock management clean. @@ -474,7 +537,8 @@ def _run_aggregation() -> dict[str, Any]: ) total_orgs = Organization.objects.count() logger.info( - "Aggregation: %d active orgs out of %d total", + "Aggregation (%s): %d active orgs out of %d total", + tier.value, len(active_org_ids), total_orgs, ) @@ -482,6 +546,7 @@ def _run_aggregation() -> dict[str, Any]: if not active_org_ids: return { "success": True, + "tier": tier.value, "organizations_processed": 0, "hourly": stats["hourly"], "daily": stats["daily"], @@ -524,6 +589,7 @@ def _run_aggregation() -> dict[str, Any]: hourly_agg, daily_agg, monthly_agg, + tier, extra_kwargs, ) except Exception: @@ -542,6 +608,7 @@ def _run_aggregation() -> dict[str, Any]: daily_agg, monthly_agg, llm_combined_fields, + tier, ) except Exception: logger.exception("Error querying combined LLM metrics for org %s", org_id) @@ -564,25 +631,37 @@ def _run_aggregation() -> dict[str, Any]: stats["errors"] += 1 logger.info( - f"Aggregation completed: {stats['orgs_processed']} orgs, " + f"Aggregation completed ({tier.value}): {stats['orgs_processed']} orgs, " f"hourly={stats['hourly']['upserted']}, " f"daily={stats['daily']['upserted']}, " f"monthly={stats['monthly']['upserted']}, " f"errors={stats['errors']}" ) + # Only the windows this run actually queried — a period reported for a tier that + # was skipped reads as work that happened. + period = {} + if _writes_hourly(tier): + period["hourly"] = { + "start": hourly_start.isoformat(), + "end": end_date.isoformat(), + } + if _writes_daily_monthly(tier): + period["daily"] = {"start": daily_start.isoformat(), "end": end_date.isoformat()} + period["monthly"] = { + "start": monthly_start.isoformat(), + "end": end_date.isoformat(), + } + return { "success": True, + "tier": tier.value, "organizations_processed": stats["orgs_processed"], "hourly": stats["hourly"], "daily": stats["daily"], "monthly": stats["monthly"], "errors": stats["errors"], - "period": { - "hourly": {"start": hourly_start.isoformat(), "end": end_date.isoformat()}, - "daily": {"start": daily_start.isoformat(), "end": end_date.isoformat()}, - "monthly": {"start": monthly_start.isoformat(), "end": end_date.isoformat()}, - }, + "period": period, } diff --git a/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py b/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py new file mode 100644 index 0000000000..8c9403d4d3 --- /dev/null +++ b/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py @@ -0,0 +1,110 @@ +"""Add an unqualified created_at index to workflow_execution. + +Serves any bare "rows in this date window" question on this table. The immediate +caller is the dashboard metrics cron's active-org prefilter +(``dashboard_metrics/tasks.py``):: + + WorkflowExecution.objects.filter(created_at__gte=window_start) + .values_list("workflow__organization_id", flat=True).distinct() + +Measured on production 2026-08-31 at 1,849ms per call — the slowest single query on +the instance by average execution time. Analysis in UN-3883. + +WHY THE EXISTING INDEXES DO NOT COVER IT. ``(workflow_id, -created_at)`` and +``(pipeline_id, -created_at)`` are date-ordered only *within* one workflow or pipeline, +so a date range with no leading column value has to scan them whole. +``we_active_by_workflow_idx`` is keyed on workflow_id, and ``we_undispatched_idx`` is a +partial index that is empty in steady state. Nothing leads with ``created_at``. + +Beyond the prefilter, this is a prerequisite for the grouped-query rewrite in UN-4045. +Grouping by organization removes the per-org predicate that +``deployed_api_requests`` / ``etl_pipeline_executions`` / ``prompt_executions`` use as +their index entry point, leaving each of them on a bare ``created_at`` range over this +table — the same shape as the prefilter, at the same cost, three more times per run. + +Design +------ +* UNPARTIAL and single-column — the predicate has no other constant to key on, and the + callers differ in what they select, so a covering column would help one and not the + others. +* CONCURRENTLY + ``atomic = False`` — ``workflow_execution`` is a multi-million-row + table in production; a plain ``AddIndex`` holds a SHARE lock for the whole build and + blocks writes, i.e. blocks every execution in flight. +* INVALID-INDEX GUARD — ``IF NOT EXISTS`` silently no-ops over a leftover INVALID index + from an interrupted CONCURRENTLY build, and Django would then record this migration as + applied while the index is physically unusable (never read, write overhead only). The + second statement RAISEs in that case, so the failure is loud rather than + green-but-broken. + +Deployment +---------- +``CREATE INDEX CONCURRENTLY`` scans the table and can run for minutes at this size — +long enough to time out a deploy's ``migrate`` step. Prefer building it OUT OF BAND +*before* the deploy; the migration then no-ops via ``IF NOT EXISTS``:: + + CREATE INDEX CONCURRENTLY IF NOT EXISTS we_created_at_idx + ON workflow_execution (created_at); + +Then confirm it is valid and that the planner picks it up:: + + SELECT c.relname, i.indisvalid FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = 'we_created_at_idx'; + -- indisvalid must be 't' + +Recovery +-------- +An interrupted build leaves an INVALID index that adds write overhead but is never read. +``IF NOT EXISTS`` will NOT rebuild over it (and the guard below RAISEs on it), so drop +it first and re-run:: + + DROP INDEX CONCURRENTLY IF EXISTS we_created_at_idx; +""" + +from django.db import migrations, models + +INDEX_NAME = "we_created_at_idx" + +_ASSERT_INDEX_VALID = f""" +DO $$ +BEGIN + IF EXISTS ( + SELECT 1 FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + WHERE c.relname = '{INDEX_NAME}' AND NOT i.indisvalid + ) THEN + RAISE EXCEPTION 'Index {INDEX_NAME} exists but is INVALID (a prior CREATE INDEX CONCURRENTLY was interrupted). Drop it and re-run: DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};'; + END IF; +END +$$; +""" + + +class Migration(migrations.Migration): + # CREATE / DROP INDEX CONCURRENTLY cannot run inside a transaction block. + atomic = False + + dependencies = [("workflow_v2", "0028_undispatched_idx_dispatched_at")] + + operations = [ + migrations.SeparateDatabaseAndState( + database_operations=[ + migrations.RunSQL( + sql=( + f"CREATE INDEX CONCURRENTLY IF NOT EXISTS {INDEX_NAME} " + "ON workflow_execution (created_at);" + ), + reverse_sql=f"DROP INDEX CONCURRENTLY IF EXISTS {INDEX_NAME};", + ), + migrations.RunSQL( + sql=_ASSERT_INDEX_VALID, reverse_sql=migrations.RunSQL.noop + ), + ], + state_operations=[ + migrations.AddIndex( + model_name="workflowexecution", + index=models.Index(fields=["created_at"], name=INDEX_NAME), + ), + ], + ), + ] diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index d6082aa423..926960b026 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -272,6 +272,13 @@ class Meta: queue_message_id__isnull=True, ), ), + # Unqualified created_at range — see migration 0029. The two indexes + # above lead with workflow_id / pipeline_id, so they are date-ordered + # only *within* one workflow and cannot serve a bare date window; the + # partial index above is empty in steady state. The dashboard metrics + # cron's active-org prefilter asks exactly that bare question and + # currently full-scans the table for it. + models.Index(fields=["created_at"], name="we_created_at_idx"), ] @property diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 44bbe50440..62970de49c 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -105,9 +105,15 @@ def _log_if_skipped(name: str, result: dict[str, Any]) -> None: @worker_task(name="dashboard_metrics.aggregate_from_sources") -def dashboard_metrics_aggregate() -> dict[str, Any]: - """Aggregate source tables into the hourly/daily/monthly metrics tables.""" - result = _call_internal(_AGGREGATE_PATH) +def dashboard_metrics_aggregate(tier: str | None = None) -> dict[str, Any]: + """Aggregate source tables into the hourly/daily/monthly metrics tables. + + ``tier`` comes from the schedule row's kwargs and selects which tiers to write — + the hourly tier and the daily/monthly pair run on separate schedules. Omitted + means all tiers, matching the backend task's default. + """ + body = {"tier": tier} if tier is not None else None + result = _call_internal(_AGGREGATE_PATH, body=body) _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) return result From 7c32887d54d37d00f2fd62cc366981420b2c534e Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Mon, 31 Aug 2026 20:21:55 +0530 Subject: [PATCH 06/14] UN-3974 [PERF] Keep scheduler ownership out of the split migration and cut _run_aggregation's complexity MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Migration 0005 used update_or_create for the PG row of the schedule it was only re-keying, which reset pg_owned to False. converge_pg_scheduler disables a row's Beat twin when the PG scheduler adopts it, so on an adopted deployment the migration would have left the aggregation with no firer at all — Beat disabled, PG no longer owning it. It now updates only task_kwargs on that row, leaving enabled and pg_owned to the scheduler that owns them. Rollback is symmetric. Threading the tier through _run_aggregation took its cognitive complexity from 25 to 27 against a limit of 15. Extracted _collect_org_metrics and _aggregate_org, and hoisted the two static metric tables to module level so they are not rebuilt per call. Names match the same extraction on #2255 so the two reconcile cleanly. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- .../0005_split_aggregation_schedule.py | 56 +++-- backend/dashboard_metrics/tasks.py | 225 ++++++++++-------- 2 files changed, 162 insertions(+), 119 deletions(-) diff --git a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py index 0843b8089d..59d2cf3fa2 100644 --- a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py +++ b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py @@ -78,31 +78,36 @@ def split_schedules(apps, schema_editor): kwargs = {"tier": spec["tier"]} if spec["exists"]: - # Keep the existing IntervalSchedule — only the payload changes. + # Update ONLY the payload. `enabled` and `pg_owned` say which scheduler + # currently fires this row, and converge_pg_scheduler owns them: adopting + # a row on PG disables its Beat twin. Rewriting either here would hand the + # row back to a scheduler that is no longer running it — or, on an adopted + # row, to neither. Its cadence does not change, so nothing else needs to. PeriodicTask.objects.filter(name=spec["name"]).update( kwargs=json.dumps(kwargs), description=spec["description"] ) - else: - schedule, _ = CrontabSchedule.objects.get_or_create( - minute=spec["crontab"]["minute"], - hour=spec["crontab"]["hour"], - day_of_week="*", - day_of_month="*", - month_of_year="*", - defaults={"timezone": "UTC"}, - ) - PeriodicTask.objects.update_or_create( - name=spec["name"], - defaults={ - "task": AGGREGATE_TASK_NAME, - "crontab": schedule, - "queue": AGGREGATE_QUEUE, - "kwargs": json.dumps(kwargs), - "enabled": True, - "description": spec["description"], - }, - ) - + PgPeriodicTask.objects.filter(name=spec["name"]).update(task_kwargs=kwargs) + continue + + schedule, _ = CrontabSchedule.objects.get_or_create( + minute=spec["crontab"]["minute"], + hour=spec["crontab"]["hour"], + day_of_week="*", + day_of_month="*", + month_of_year="*", + defaults={"timezone": "UTC"}, + ) + PeriodicTask.objects.update_or_create( + name=spec["name"], + defaults={ + "task": AGGREGATE_TASK_NAME, + "crontab": schedule, + "queue": AGGREGATE_QUEUE, + "kwargs": json.dumps(kwargs), + "enabled": True, + "description": spec["description"], + }, + ) PgPeriodicTask.objects.update_or_create( name=spec["name"], defaults={ @@ -119,7 +124,12 @@ def split_schedules(apps, schema_editor): def merge_schedules(apps, schema_editor): - """Restore the single every-15-minutes row that writes all three tiers.""" + """Restore the single every-15-minutes row that writes all three tiers. + + Symmetric with the forward direction: the added rows go, and the surviving row + gets its payload back. `enabled` / `pg_owned` are left alone in both directions, + so whichever scheduler was firing the aggregation before the rollback still is. + """ PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 30862fde7c..040d6a1c15 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -458,6 +458,132 @@ def _aggregate_llm_combined( _upsert_agg(monthly_agg, key, metric_type, value) +# Metric definitions: (name, query_method, is_histogram) +# Note: llm_calls, challenges, summarization_calls, and llm_usage are +# handled separately via get_llm_metrics_combined (1 query instead of 4). +METRIC_CONFIGS = [ + ("documents_processed", MetricsQueryService.get_documents_processed, False), + ("pages_processed", MetricsQueryService.get_pages_processed, True), + ("deployed_api_requests", MetricsQueryService.get_deployed_api_requests, False), + ("etl_pipeline_executions", MetricsQueryService.get_etl_pipeline_executions, False), + ("prompt_executions", MetricsQueryService.get_prompt_executions, False), + ("failed_pages", MetricsQueryService.get_failed_pages, True), + ("hitl_reviews", MetricsQueryService.get_hitl_reviews, False), + ("hitl_completions", MetricsQueryService.get_hitl_completions, False), +] + +# LLM metrics combined via conditional aggregation (4 metrics in 1 query). +# Maps combined query field -> (metric_name, metric_type) +LLM_COMBINED_FIELDS = { + "llm_calls": ("llm_calls", MetricType.COUNTER), + "challenges": ("challenges", MetricType.COUNTER), + "summarization_calls": ("summarization_calls", MetricType.COUNTER), + "llm_usage": ("llm_usage", MetricType.HISTOGRAM), +} + + +def _collect_org_metrics( + org: Organization, + hourly_start: datetime, + daily_start: datetime, + monthly_start: datetime, + end_date: datetime, + tier: AggregationTier, +) -> tuple[dict, dict, dict, int]: + """Query every metric for one org into per-tier aggregate dicts. + + A failing metric is logged and counted, not raised: one bad query should not + cost the org its other metrics. + + Returns: + (hourly_agg, daily_agg, monthly_agg, error_count) + """ + org_id = str(org.id) + hourly_agg: dict[tuple, dict] = {} + daily_agg: dict[tuple, dict] = {} + monthly_agg: dict[tuple, dict] = {} + errors = 0 + + for metric_name, query_method, is_histogram in METRIC_CONFIGS: + metric_type = MetricType.HISTOGRAM if is_histogram else MetricType.COUNTER + + # Pass org_identifier to PageUsage-based metrics to + # avoid redundant Organization lookups per call. + extra_kwargs = {} + if metric_name == "pages_processed": + extra_kwargs["org_identifier"] = org.organization_id + + try: + _aggregate_single_metric( + query_method, + metric_name, + metric_type, + org_id, + hourly_start, + daily_start, + monthly_start, + end_date, + hourly_agg, + daily_agg, + monthly_agg, + tier, + extra_kwargs, + ) + except Exception: + logger.exception("Error querying %s for org %s", metric_name, org_id) + errors += 1 + + # Combined LLM metrics: 1 query per granularity instead of 4 + try: + _aggregate_llm_combined( + org_id, + hourly_start, + daily_start, + monthly_start, + end_date, + hourly_agg, + daily_agg, + monthly_agg, + LLM_COMBINED_FIELDS, + tier, + ) + except Exception: + logger.exception("Error querying combined LLM metrics for org %s", org_id) + errors += 1 + + return hourly_agg, daily_agg, monthly_agg, errors + + +def _aggregate_org( + org: Organization, + hourly_start: datetime, + daily_start: datetime, + monthly_start: datetime, + end_date: datetime, + tier: AggregationTier, + stats: dict[str, Any], +) -> None: + """Aggregate one organization and upsert the tiers this run writes.""" + try: + hourly_agg, daily_agg, monthly_agg, errors = _collect_org_metrics( + org, hourly_start, daily_start, monthly_start, end_date, tier + ) + stats["errors"] += errors + + # Bulk upsert each populated tier (single INSERT...ON CONFLICT each) + if hourly_agg: + stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) + if daily_agg: + stats["daily"]["upserted"] += _bulk_upsert_daily(daily_agg) + if monthly_agg: + stats["monthly"]["upserted"] += _bulk_upsert_monthly(monthly_agg) + + stats["orgs_processed"] += 1 + except Exception: + logger.exception("Error processing org %s", org.id) + stats["errors"] += 1 + + def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, Any]: """Execute the actual aggregation logic. @@ -487,33 +613,6 @@ def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, A month=end_date.month - 1, day=1, hour=0, minute=0, second=0, microsecond=0 ) - # Metric definitions: (name, query_method, is_histogram) - # Note: llm_calls, challenges, summarization_calls, and llm_usage are - # handled separately via get_llm_metrics_combined (1 query instead of 4). - metric_configs = [ - ("documents_processed", MetricsQueryService.get_documents_processed, False), - ("pages_processed", MetricsQueryService.get_pages_processed, True), - ("deployed_api_requests", MetricsQueryService.get_deployed_api_requests, False), - ( - "etl_pipeline_executions", - MetricsQueryService.get_etl_pipeline_executions, - False, - ), - ("prompt_executions", MetricsQueryService.get_prompt_executions, False), - ("failed_pages", MetricsQueryService.get_failed_pages, True), - ("hitl_reviews", MetricsQueryService.get_hitl_reviews, False), - ("hitl_completions", MetricsQueryService.get_hitl_completions, False), - ] - - # LLM metrics combined via conditional aggregation (4 metrics in 1 query). - # Maps combined query field -> (metric_name, metric_type) - llm_combined_fields = { - "llm_calls": ("llm_calls", MetricType.COUNTER), - "challenges": ("challenges", MetricType.COUNTER), - "summarization_calls": ("summarization_calls", MetricType.COUNTER), - "llm_usage": ("llm_usage", MetricType.HISTOGRAM), - } - stats = { "hourly": {"upserted": 0}, "daily": {"upserted": 0}, @@ -560,75 +659,9 @@ def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, A ) for org in organizations: - org_id = str(org.id) - org_identifier = org.organization_id # Pre-resolved for PageUsage queries - hourly_agg: dict[tuple, dict] = {} - daily_agg: dict[tuple, dict] = {} - monthly_agg: dict[tuple, dict] = {} - - try: - for metric_name, query_method, is_histogram in metric_configs: - metric_type = MetricType.HISTOGRAM if is_histogram else MetricType.COUNTER - - # Pass org_identifier to PageUsage-based metrics to - # avoid redundant Organization lookups per call. - extra_kwargs = {} - if metric_name == "pages_processed": - extra_kwargs["org_identifier"] = org_identifier - - try: - _aggregate_single_metric( - query_method, - metric_name, - metric_type, - org_id, - hourly_start, - daily_start, - monthly_start, - end_date, - hourly_agg, - daily_agg, - monthly_agg, - tier, - extra_kwargs, - ) - except Exception: - logger.exception("Error querying %s for org %s", metric_name, org_id) - stats["errors"] += 1 - - # Combined LLM metrics: 1 query per granularity instead of 4 - try: - _aggregate_llm_combined( - org_id, - hourly_start, - daily_start, - monthly_start, - end_date, - hourly_agg, - daily_agg, - monthly_agg, - llm_combined_fields, - tier, - ) - except Exception: - logger.exception("Error querying combined LLM metrics for org %s", org_id) - stats["errors"] += 1 - - # Bulk upsert all three tiers (single INSERT...ON CONFLICT each) - if hourly_agg: - stats["hourly"]["upserted"] += _bulk_upsert_hourly(hourly_agg) - - if daily_agg: - stats["daily"]["upserted"] += _bulk_upsert_daily(daily_agg) - - if monthly_agg: - stats["monthly"]["upserted"] += _bulk_upsert_monthly(monthly_agg) - - stats["orgs_processed"] += 1 - - except Exception: - logger.exception("Error processing org %s", org_id) - stats["errors"] += 1 + _aggregate_org( + org, hourly_start, daily_start, monthly_start, end_date, tier, stats + ) logger.info( f"Aggregation completed ({tier.value}): {stats['orgs_processed']} orgs, " From 9a8285a172402f87e9f70fb83a746bddbcae8459 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Mon, 31 Aug 2026 20:28:27 +0530 Subject: [PATCH 07/14] UN-3974 [PERF] Trim comments and docstrings to the project ceiling The two index migrations carried 50-60 line docstrings restating the prod plan, deployment runbook and recovery steps. That detail belongs in the PR, not in files every future agent scans. Cut to purpose and key behaviour. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/internal_views.py | 3 +- .../0005_split_aggregation_schedule.py | 48 ++++--------- backend/dashboard_metrics/tasks.py | 27 +++---- .../migrations/0029_we_created_at_idx.py | 71 ++++--------------- .../workflow_v2/models/execution.py | 8 +-- workers/scheduler/dashboard_metrics_tasks.py | 5 +- 6 files changed, 40 insertions(+), 122 deletions(-) diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index 913699f127..529e8f34af 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -93,8 +93,7 @@ class AggregateMetricsAPIView(_MetricsTaskAPIView): only because the PG consumer has no Django, not to change what the job does. Optional ``tier`` in the body selects which metric tiers to write; omitting it - writes all of them, matching the task's own default. An unrecognised value is a - 400 rather than a silent no-op. + writes all of them. An unrecognised value is a 400, not a silent no-op. """ def post(self, request: Request) -> Response: diff --git a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py index 59d2cf3fa2..9d6b40454f 100644 --- a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py +++ b/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py @@ -1,30 +1,13 @@ """Split the metrics aggregation into two schedules by tier (UN-3974). -Before this, one schedule ran every 15 minutes and wrote all three tiers. Dashboard -daily and monthly figures do not need 15-minute freshness, so they move to hourly: -96 runs a day becomes 24 for the expensive DAY-granularity half of the work, while the -hourly tier keeps its 15-minute cadence. - -Both rows point at the SAME task (``dashboard_metrics.aggregate_from_sources``) and -differ only in ``tier`` kwargs. A second task name would need its own worker-side -registration and internal endpoint for the PG path; a kwarg needs neither. - -**Beat and PG rows are declared together here, from one spec.** ``0002_setup_periodic_tasks`` -(Beat) and ``0004_pg_periodic_tasks`` (PG) declare the same schedules in two places, and -``tests/test_pg_periodic_task_declarations.py`` exists to catch them drifting apart. One -spec written twice by the same function cannot drift, so this migration needs no such -guard. ``AGGREGATION_SCHEDULES`` is module-level so a future test can import it. - -Rows land consistent with how each scheduler expects them: - -* Beat ``kwargs`` is a JSON *string*; ``PgPeriodicTask.task_kwargs`` is a JSONField, so - it is stored decoded. -* The new PG row lands **inert** (``pg_owned=False``, ``next_run_at=NULL``) for the same - reason as ``0004`` — the PG scheduler skips rows it does not own, and a NULL - ``next_run_at`` records a baseline next tick rather than firing a catch-up burst. - -Reverse restores the pre-split state: the new rows are deleted and the aggregate row's -kwargs are cleared, putting it back to writing all three tiers every 15 minutes. +The hourly tier keeps its 15-minute cadence; the daily and monthly tiers move to +hourly, taking the expensive DAY-granularity half of the work from 96 runs a day to +24. Both rows run the same task and differ only in their ``tier`` kwargs — a second +task name would need its own worker registration and internal endpoint. + +Beat and PG rows are declared here from one spec, so this pair cannot drift the way +0002 and 0004 can. Beat stores kwargs as a JSON string, PgPeriodicTask decoded. The +new PG row lands inert (``pg_owned=False``) like 0004's. """ import json @@ -39,8 +22,7 @@ TIER_HOURLY = "hourly" TIER_DAILY_MONTHLY = "daily_monthly" -# The row that already exists (created by 0002 / 0004); only its kwargs and -# description change, its every-15-minutes schedule does not. +# Created by 0002 / 0004; only its kwargs and description change here. EXISTING_AGGREGATE_ROW = "dashboard_metrics_aggregate_from_sources" AGGREGATION_SCHEDULES = [ @@ -78,11 +60,9 @@ def split_schedules(apps, schema_editor): kwargs = {"tier": spec["tier"]} if spec["exists"]: - # Update ONLY the payload. `enabled` and `pg_owned` say which scheduler - # currently fires this row, and converge_pg_scheduler owns them: adopting - # a row on PG disables its Beat twin. Rewriting either here would hand the - # row back to a scheduler that is no longer running it — or, on an adopted - # row, to neither. Its cadence does not change, so nothing else needs to. + # Payload only. `enabled` and `pg_owned` say which scheduler fires this + # row and belong to converge_pg_scheduler; rewriting them here can leave + # an adopted row with no firer. Its cadence does not change. PeriodicTask.objects.filter(name=spec["name"]).update( kwargs=json.dumps(kwargs), description=spec["description"] ) @@ -126,9 +106,7 @@ def split_schedules(apps, schema_editor): def merge_schedules(apps, schema_editor): """Restore the single every-15-minutes row that writes all three tiers. - Symmetric with the forward direction: the added rows go, and the surviving row - gets its payload back. `enabled` / `pg_owned` are left alone in both directions, - so whichever scheduler was firing the aggregation before the rollback still is. + Leaves `enabled` / `pg_owned` alone, as the forward direction does. """ PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 040d6a1c15..8af0cba530 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -206,11 +206,9 @@ def _bulk_upsert_monthly(aggregations: dict) -> int: class AggregationTier(StrEnum): - """Which metric tiers a single aggregation run writes. + """Which metric tiers one aggregation run writes. - The hourly tier needs 15-minute freshness; the daily and monthly tiers do not, - so they run on separate schedules. Daily and monthly stay together because one - DAY-granularity query feeds both. + Daily and monthly stay together because one DAY-granularity query feeds both. """ HOURLY = "hourly" @@ -218,8 +216,7 @@ class AggregationTier(StrEnum): ALL = "all" -# Per-tier so the 15-minute hourly run and the hourly daily/monthly run — which -# collide at the top of every hour — do not starve each other. +# Keyed per tier: the two schedules collide hourly and must not block each other. AGGREGATION_LOCK_KEY_PREFIX = "dashboard_metrics:aggregation_lock" AGGREGATION_LOCK_TIMEOUT = 900 # 15 minutes (matches the fastest task schedule) @@ -244,7 +241,7 @@ def _acquire_aggregation_lock(lock_key: str) -> bool: that the lock is older than AGGREGATION_LOCK_TIMEOUT and reclaims it. Args: - lock_key: Cache key to lock on — one per tier, see _aggregation_lock_key + lock_key: Cache key to lock on, one per tier Returns: True if lock was acquired, False if another run is legitimately active. @@ -299,9 +296,8 @@ def aggregate_metrics_from_sources( ) -> dict[str, Any]: """Aggregate metrics from source tables into the hourly/daily/monthly tables. - Two schedules call this with different tiers: the hourly tier every 15 minutes, - the daily and monthly tiers hourly. Each tier locks separately so the two never - block each other. + Two schedules call this with different tiers: hourly every 15 minutes, daily + and monthly hourly. Each tier locks separately. Uses a Redis distributed lock with self-healing to prevent overlapping runs. If a previous run was killed without releasing the lock, the next @@ -313,9 +309,8 @@ def aggregate_metrics_from_sources( - Monthly: Last 2 months (current + previous month) Args: - tier: Which tiers to write, an AggregationTier value. Defaults to all, - so a caller that omits it gets the pre-split behaviour rather than - silently writing nothing. + tier: An AggregationTier value. Defaults to all, so a caller that omits + it writes every tier rather than none. Returns: Dict with aggregation summary for the tiers that ran @@ -492,8 +487,7 @@ def _collect_org_metrics( ) -> tuple[dict, dict, dict, int]: """Query every metric for one org into per-tier aggregate dicts. - A failing metric is logged and counted, not raised: one bad query should not - cost the org its other metrics. + A failing metric is counted, not raised, so it does not cost the org the rest. Returns: (hourly_agg, daily_agg, monthly_agg, error_count) @@ -671,8 +665,7 @@ def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, A f"errors={stats['errors']}" ) - # Only the windows this run actually queried — a period reported for a tier that - # was skipped reads as work that happened. + # Only the windows this run queried; a skipped tier reports no period. period = {} if _writes_hourly(tier): period["hourly"] = { diff --git a/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py b/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py index 8c9403d4d3..11054900fd 100644 --- a/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py +++ b/backend/workflow_manager/workflow_v2/migrations/0029_we_created_at_idx.py @@ -1,70 +1,23 @@ -"""Add an unqualified created_at index to workflow_execution. +"""Add a created_at index to workflow_execution. -Serves any bare "rows in this date window" question on this table. The immediate -caller is the dashboard metrics cron's active-org prefilter -(``dashboard_metrics/tasks.py``):: +Serves bare "rows in this date window" queries with no leading column value — the +dashboard metrics active-org prefilter today, and the grouped metric queries in +UN-4045. The composite indexes lead with workflow_id / pipeline_id, so they are +date-ordered only within one workflow or pipeline; the partial index is empty in +steady state. Measurements in UN-3883. - WorkflowExecution.objects.filter(created_at__gte=window_start) - .values_list("workflow__organization_id", flat=True).distinct() - -Measured on production 2026-08-31 at 1,849ms per call — the slowest single query on -the instance by average execution time. Analysis in UN-3883. - -WHY THE EXISTING INDEXES DO NOT COVER IT. ``(workflow_id, -created_at)`` and -``(pipeline_id, -created_at)`` are date-ordered only *within* one workflow or pipeline, -so a date range with no leading column value has to scan them whole. -``we_active_by_workflow_idx`` is keyed on workflow_id, and ``we_undispatched_idx`` is a -partial index that is empty in steady state. Nothing leads with ``created_at``. - -Beyond the prefilter, this is a prerequisite for the grouped-query rewrite in UN-4045. -Grouping by organization removes the per-org predicate that -``deployed_api_requests`` / ``etl_pipeline_executions`` / ``prompt_executions`` use as -their index entry point, leaving each of them on a bare ``created_at`` range over this -table — the same shape as the prefilter, at the same cost, three more times per run. - -Design ------- -* UNPARTIAL and single-column — the predicate has no other constant to key on, and the - callers differ in what they select, so a covering column would help one and not the - others. -* CONCURRENTLY + ``atomic = False`` — ``workflow_execution`` is a multi-million-row - table in production; a plain ``AddIndex`` holds a SHARE lock for the whole build and - blocks writes, i.e. blocks every execution in flight. -* INVALID-INDEX GUARD — ``IF NOT EXISTS`` silently no-ops over a leftover INVALID index - from an interrupted CONCURRENTLY build, and Django would then record this migration as - applied while the index is physically unusable (never read, write overhead only). The - second statement RAISEs in that case, so the failure is loud rather than - green-but-broken. - -Deployment ----------- -``CREATE INDEX CONCURRENTLY`` scans the table and can run for minutes at this size — -long enough to time out a deploy's ``migrate`` step. Prefer building it OUT OF BAND -*before* the deploy; the migration then no-ops via ``IF NOT EXISTS``:: - - CREATE INDEX CONCURRENTLY IF NOT EXISTS we_created_at_idx - ON workflow_execution (created_at); - -Then confirm it is valid and that the planner picks it up:: - - SELECT c.relname, i.indisvalid FROM pg_class c - JOIN pg_index i ON i.indexrelid = c.oid - WHERE c.relname = 'we_created_at_idx'; - -- indisvalid must be 't' - -Recovery --------- -An interrupted build leaves an INVALID index that adds write overhead but is never read. -``IF NOT EXISTS`` will NOT rebuild over it (and the guard below RAISEs on it), so drop -it first and re-run:: - - DROP INDEX CONCURRENTLY IF EXISTS we_created_at_idx; +Built CONCURRENTLY (atomic = False): a plain AddIndex holds a SHARE lock for the +whole build and would block every execution in flight. Prefer building it out of +band before the deploy; the migration then no-ops via IF NOT EXISTS. """ from django.db import migrations, models INDEX_NAME = "we_created_at_idx" +# An interrupted CONCURRENTLY build leaves an INVALID index that costs on every +# write and is never read. IF NOT EXISTS would keep it while Django recorded the +# migration as applied, so fail loudly instead. _ASSERT_INDEX_VALID = f""" DO $$ BEGIN diff --git a/backend/workflow_manager/workflow_v2/models/execution.py b/backend/workflow_manager/workflow_v2/models/execution.py index 926960b026..a8104898d9 100644 --- a/backend/workflow_manager/workflow_v2/models/execution.py +++ b/backend/workflow_manager/workflow_v2/models/execution.py @@ -272,12 +272,8 @@ class Meta: queue_message_id__isnull=True, ), ), - # Unqualified created_at range — see migration 0029. The two indexes - # above lead with workflow_id / pipeline_id, so they are date-ordered - # only *within* one workflow and cannot serve a bare date window; the - # partial index above is empty in steady state. The dashboard metrics - # cron's active-org prefilter asks exactly that bare question and - # currently full-scans the table for it. + # Bare created_at range scans; the indexes above are date-ordered only + # within one workflow or pipeline. See migration 0029. models.Index(fields=["created_at"], name="we_created_at_idx"), ] diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 62970de49c..18e002c6e3 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -108,9 +108,8 @@ def _log_if_skipped(name: str, result: dict[str, Any]) -> None: def dashboard_metrics_aggregate(tier: str | None = None) -> dict[str, Any]: """Aggregate source tables into the hourly/daily/monthly metrics tables. - ``tier`` comes from the schedule row's kwargs and selects which tiers to write — - the hourly tier and the daily/monthly pair run on separate schedules. Omitted - means all tiers, matching the backend task's default. + ``tier`` comes from the schedule row's kwargs and selects which tiers to write; + omitted means all of them. """ body = {"tier": tier} if tier is not None else None result = _call_internal(_AGGREGATE_PATH, body=body) From 3688d749b40e2a230b87517a1f73cf19676608dc Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 1 Sep 2026 17:29:23 +0530 Subject: [PATCH 08/14] UN-3974 [PERF] Cover all three acceptance criteria with tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The suite runs with --no-migrations, so neither 0005 nor 0029 ever executes in CI, and nothing pinned the schedule split's behaviour at all. 46 tests, at least one per acceptance criterion. AC-1 — cadence, and the tier reaching the task. 0005 creates one row and rewrites one, both scheduler tables agreeing, and the rewrite touches neither pg_owned nor enabled: on an adopted deployment converge_pg_scheduler has already disabled the Beat twin, so handing ownership back would leave the aggregation with no firer. Separately the internal endpoint and the worker proxy are pinned to carry `tier` — that leg fails silently, since _call_internal builds a body only when a tier is given and the existing worker test called the task without one. AC-2 — the split changes no figure. Runs the real _run_aggregation three times and diffs the metrics tables: `hourly` reproduces the pre-split hourly figures exactly, and hourly + daily_monthly reproduce every row `all` writes. Two guards keep it from going vacuous, the second because mutation testing caught the first version passing while _aggregate_single_metric was broken — the fixture produced only LLM metrics, leaving half the split unverified. AC-3 — the index. Migration shape (non-atomic, CONCURRENTLY both directions, the INVALID guard, AddIndex confined to state_operations), plus an integration test that EXPLAINs the query the aggregation actually issues, captured rather than rewritten: a hand-copied queryset would keep passing after the prefilter changed, which is the one thing it is for. Rows are inserted in ascending created_at order so the heap matches production's append order. The Query Insights half of AC-3 is a production reading and is deliberately not faked here. Every test verified to fail when the thing it guards breaks. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- .../tests/test_active_org_prefilter.py | 113 +++++++++++ .../tests/test_aggregation_dispatch.py | 123 ++++++++++++ .../tests/test_aggregation_tier.py | 91 +++++++++ .../test_pg_periodic_task_declarations.py | 117 ++++++++++- .../tests/test_tier_split_equivalence.py | 181 ++++++++++++++++++ .../tests/test_we_created_at_idx.py | 110 +++++++++++ workers/tests/test_dashboard_metrics_tasks.py | 21 ++ 7 files changed, 754 insertions(+), 2 deletions(-) create mode 100644 backend/dashboard_metrics/tests/test_active_org_prefilter.py create mode 100644 backend/dashboard_metrics/tests/test_aggregation_dispatch.py create mode 100644 backend/dashboard_metrics/tests/test_aggregation_tier.py create mode 100644 backend/dashboard_metrics/tests/test_tier_split_equivalence.py create mode 100644 backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py diff --git a/backend/dashboard_metrics/tests/test_active_org_prefilter.py b/backend/dashboard_metrics/tests/test_active_org_prefilter.py new file mode 100644 index 0000000000..a1dbe082ad --- /dev/null +++ b/backend/dashboard_metrics/tests/test_active_org_prefilter.py @@ -0,0 +1,113 @@ +"""The active-org prefilter can actually use we_created_at_idx (UN-3974, AC-3). + +AC-3 is worded as a production observation — "no longer appears in the top 10 by total +execution time in Query Insights" — and that half can only be read off production. The +half that is answerable here is the one underneath it: the prefilter bounds nothing but +`created_at`, and the index exists to serve exactly that shape. + +What this pins is the pairing. `workflow_manager/workflow_v2/tests/test_we_created_at_idx.py` +proves the index is declared and built safely; this proves the query still looks like +something it can serve. Either half can drift without the other noticing — someone +narrowing the prefilter to lead with a different column leaves the index built, valid, +and dead. + +Rows are inserted in ascending `created_at` order so the heap matches production, where +executions are appended as they happen. With them scattered the planner reads the whole +composite (workflow_id, created_at DESC) index instead, which is an artefact of the +fixture rather than anything about the query. + +DB-bound, so conftest marks it integration. +""" + +from __future__ import annotations + + +from account_v2.models import Organization +from django.db import connection +from django.test import TestCase +from django.test.utils import CaptureQueriesContext +from workflow_manager.workflow_v2.models.workflow import Workflow + +from dashboard_metrics.tasks import AggregationTier, _run_aggregation + +INDEX_NAME = "we_created_at_idx" +_ROWS = 12000 +_SPAN_DAYS = 255 + + +class TestThePrefilterCanUseTheIndex(TestCase): + """Production ratios rather than production size: ~2.7% of rows in the 7-day window + is what decides whether the planner reaches for an index or scans. + """ + + def setUp(self) -> None: + self.org = Organization.objects.create( + organization_id="prefilter-org", name="prefilter", display_name="Prefilter" + ) + self.workflow = Workflow.objects.create( + workflow_name="prefilter-wf", organization=self.org + ) + with connection.cursor() as cur: + cur.execute( + """ + INSERT INTO workflow_execution ( + id, created_at, modified_at, workflow_id, execution_mode, + execution_method, execution_type, execution_log_id, status, + error_message, attempts, execution_time, result_acknowledged, + total_files) + SELECT gen_random_uuid(), ts, ts, %s, 'INSTANT', 'DIRECT', 'COMPLETE', + '', 'COMPLETED', '', 0, 1.0, false, 1 + FROM generate_series(1, %s) g + CROSS JOIN LATERAL ( + SELECT now() - (%s - (g::float / %s) * %s) * interval '1 day' + ) AS t(ts) + """, + [self.workflow.id, _ROWS, _SPAN_DAYS, _ROWS, _SPAN_DAYS], + ) + cur.execute("ANALYZE workflow_execution") + + def _prefilter_sql(self) -> str: + """The real query, taken from the task rather than rewritten here. + + A hand-copied queryset would keep passing after the prefilter changed, which is + the one thing this test is for. + """ + with CaptureQueriesContext(connection) as ctx: + _run_aggregation(AggregationTier.HOURLY) + candidates = [ + q["sql"] + for q in ctx.captured_queries + if "workflow_execution" in q["sql"] + and "DISTINCT" in q["sql"].upper() + and "created_at" in q["sql"] + ] + assert candidates, "the aggregation issued no active-org prefilter query" + return str(candidates[0]) + + def test_the_window_is_the_selectivity_the_index_is_for(self) -> None: + """If the prefilter ever widened to most of the table, an index on created_at + would stop being the right answer — the planner would scan regardless. + """ + with connection.cursor() as cur: + cur.execute( + "SELECT count(*) FILTER (WHERE created_at >= now() - interval '7 days')" + "::float / count(*) FROM workflow_execution" + ) + share = cur.fetchone()[0] + assert 0 < share < 0.10 + + def test_the_planner_reaches_for_the_index(self) -> None: + """The whole point of 2a. An index that exists but is never chosen costs on + every insert and buys nothing. + """ + with connection.cursor() as cur: + cur.execute("EXPLAIN " + self._prefilter_sql()) + plan = "\n".join(row[0] for row in cur.fetchall()) + assert INDEX_NAME in plan, f"expected {INDEX_NAME} in:\n{plan}" + + def test_the_prefilter_does_not_scan_the_executions_table(self) -> None: + """The regression the index is meant to remove.""" + with connection.cursor() as cur: + cur.execute("EXPLAIN " + self._prefilter_sql()) + plan = "\n".join(row[0] for row in cur.fetchall()) + assert "Seq Scan on workflow_execution" not in plan, plan diff --git a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py new file mode 100644 index 0000000000..c4892ea59d --- /dev/null +++ b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py @@ -0,0 +1,123 @@ +"""Guard: the tier a schedule row declares reaches the task (UN-3974, AC-1). + +Two schedulers fire the same task name at two different implementations. Beat reads +``PeriodicTask.kwargs`` and calls the Django ``@shared_task`` directly; the PG scheduler +reads ``PgPeriodicTask.task_kwargs`` and goes through the worker proxy and the internal +endpoint to the same function. Both legs have to carry ``tier``, and a break in either is +invisible — the job still runs, still returns success, and just writes the wrong tiers. + +The worker half of the PG leg is pinned in ``workers/tests/test_dashboard_metrics_tasks.py``; +this covers the endpoint that receives it and the Beat leg's kwargs. + +DB-free: the task is mocked, and the Beat kwargs are read from the migration spec rather +than from a migrated database. +""" + +from __future__ import annotations + +import importlib +import inspect +import json +import os +from typing import Any +from unittest import mock + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from rest_framework.test import APIRequestFactory # noqa: E402 + +from dashboard_metrics import internal_views # noqa: E402 +from dashboard_metrics.tasks import ( # noqa: E402 + AggregationTier, + aggregate_metrics_from_sources, +) + +_SPLIT_MIGRATION = "dashboard_metrics.migrations.0005_split_aggregation_schedule" +_ENDPOINT = "/internal/v1/dashboard-metrics/aggregate/" + + +def _post(body: dict[str, Any]) -> tuple[int, Any]: + """POST to the aggregate endpoint with the task mocked; return status and its kwargs.""" + view = internal_views.AggregateMetricsAPIView.as_view() + request = APIRequestFactory().post(_ENDPOINT, body, format="json") + with mock.patch.object( + internal_views, "aggregate_metrics_from_sources", return_value={"ok": True} + ) as task: + response = view(request) + return response.status_code, (task.call_args.kwargs if task.call_args else None) + + +class TestThePgLegCarriesTheTier: + """The endpoint the worker proxy POSTs to.""" + + @pytest.mark.parametrize("tier", ["hourly", "daily_monthly", "all"]) + def test_the_endpoint_forwards_the_tier_to_the_task(self, tier: str) -> None: + status, called_with = _post({"tier": tier}) + assert status == 200 + assert called_with == {"tier": tier} + + def test_an_omitted_tier_leaves_the_task_default_in_place(self) -> None: + """Not 'hourly', and not nothing: the task's own default is `all`, and passing + anything here would override it during the pre-0005 deploy window. + """ + status, called_with = _post({}) + assert status == 200 + assert called_with == {} + + def test_an_unrecognised_tier_is_rejected_rather_than_ignored(self) -> None: + """A silent no-op would look like a successful run that wrote nothing. + + Runs against the real task, not the mock: the 400 comes from the ValueError the + task raises, and a mock would accept anything and return 200. The tier is + validated on the task's first line, so nothing touches the database. + """ + view = internal_views.AggregateMetricsAPIView.as_view() + response = view(APIRequestFactory().post(_ENDPOINT, {"tier": "houry"}, format="json")) + assert response.status_code == 400 + assert "houry" in str(response.data) + + +class TestTheBeatLegCarriesTheTier: + """Beat passes the row's stored JSON kwargs straight into the task signature.""" + + @pytest.fixture(scope="class") + def declared_kwargs(self) -> dict[str, dict[str, Any]]: + mod = importlib.import_module(_SPLIT_MIGRATION) + return {s["name"]: {"tier": s["tier"]} for s in mod.AGGREGATION_SCHEDULES} + + def test_both_rows_declare_a_tier( + self, declared_kwargs: dict[str, dict[str, Any]] + ) -> None: + assert len(declared_kwargs) == 2 + assert all("tier" in kw for kw in declared_kwargs.values()) + + def test_every_declared_kwarg_set_binds_to_the_task_signature( + self, declared_kwargs: dict[str, dict[str, Any]] + ) -> None: + """A row declaring a kwarg the task does not accept fails at call time, inside + the worker, where it surfaces as a retrying task rather than a bad schedule. + """ + signature = inspect.signature(aggregate_metrics_from_sources) + for kwargs in declared_kwargs.values(): + signature.bind(**kwargs) + + def test_every_declared_tier_is_a_real_tier( + self, declared_kwargs: dict[str, dict[str, Any]] + ) -> None: + """The migration cannot import the enum, so it repeats the literals. A typo + there raises inside the task on every single run.""" + for kwargs in declared_kwargs.values(): + AggregationTier(kwargs["tier"]) + + def test_beat_stores_the_kwargs_as_json_the_task_can_receive(self) -> None: + """Beat's kwargs column is a JSON *string*; PgPeriodicTask's is a JSONField. + The Beat side has to round-trip back to the same mapping.""" + mod = importlib.import_module(_SPLIT_MIGRATION) + for spec in mod.AGGREGATION_SCHEDULES: + assert json.loads(json.dumps({"tier": spec["tier"]})) == {"tier": spec["tier"]} diff --git a/backend/dashboard_metrics/tests/test_aggregation_tier.py b/backend/dashboard_metrics/tests/test_aggregation_tier.py new file mode 100644 index 0000000000..10c9fb6183 --- /dev/null +++ b/backend/dashboard_metrics/tests/test_aggregation_tier.py @@ -0,0 +1,91 @@ +"""Guard: the tier a schedule row asks for is the tier that gets written. + +The split runs one task on two schedules that differ only in their ``tier`` kwarg, so +the gating predicates and the per-tier lock key are the whole mechanism. Each property +here is one way the split fails silently — writing nothing, writing both tiers from one +schedule, or the two schedules starving each other on the lock. + +DB-free, so this runs in the unit tier alongside test_pg_periodic_task_declarations.py. +""" + +from __future__ import annotations + +import os + +import django +import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from dashboard_metrics.tasks import ( # noqa: E402 + AggregationTier, + _aggregation_lock_key, + _writes_daily_monthly, + _writes_hourly, +) + + +class TestWhichTiersEachRunWrites: + @pytest.mark.parametrize( + "tier,hourly,daily_monthly", + [ + (AggregationTier.HOURLY, True, False), + (AggregationTier.DAILY_MONTHLY, False, True), + (AggregationTier.ALL, True, True), + ], + ) + def test_the_predicates_partition_the_work( + self, tier: AggregationTier, hourly: bool, daily_monthly: bool + ) -> None: + assert _writes_hourly(tier) is hourly + assert _writes_daily_monthly(tier) is daily_monthly + + def test_the_two_schedules_together_cover_every_tier(self) -> None: + """Neither schedule may leave a tier unwritten: hourly and daily_monthly are + the only two rows, so between them they have to do everything `all` does. + """ + scheduled = (AggregationTier.HOURLY, AggregationTier.DAILY_MONTHLY) + assert any(_writes_hourly(t) for t in scheduled) + assert any(_writes_daily_monthly(t) for t in scheduled) + + def test_no_tier_is_written_by_both_schedules(self) -> None: + """Overlap would mean duplicate work every hour on the hour. The upserts make + it harmless, not free. + """ + assert not _writes_daily_monthly(AggregationTier.HOURLY) + assert not _writes_hourly(AggregationTier.DAILY_MONTHLY) + + +class TestTheDefaultIsAll: + def test_omitting_the_tier_writes_everything_rather_than_nothing(self) -> None: + """Between the code deploying and migration 0005 running, the schedule row + still carries no tier kwarg. Defaulting to anything narrower would stop writing + tiers during that window; defaulting to none would stop writing entirely. + """ + assert AggregationTier("all") is AggregationTier.ALL + assert _writes_hourly(AggregationTier.ALL) + assert _writes_daily_monthly(AggregationTier.ALL) + + def test_an_unrecognised_tier_raises(self) -> None: + """The internal view turns this into a 400. A silent no-op would look like a + successful run that wrote nothing. + """ + with pytest.raises(ValueError): + AggregationTier("houry") + + +class TestTheLockIsPerTier: + def test_every_tier_gets_its_own_key(self) -> None: + """The two schedules collide at the top of every hour. One global key and + whichever fired first would hold it while the other returned lock_held — + so the slower tier could be starved indefinitely. + """ + keys = {_aggregation_lock_key(t) for t in AggregationTier} + assert len(keys) == len(list(AggregationTier)) + + def test_the_key_names_the_tier(self) -> None: + for tier in AggregationTier: + assert _aggregation_lock_key(tier).endswith(tier.value) diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py index 85ea407899..7e1257fd60 100644 --- a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -6,7 +6,12 @@ the whole failure mode this file exists for: a schedule changed on Beat but not on PG means the task silently runs on a different cadence the moment the flag flips. -DB-free — both migration modules are imported and their declared specs compared directly, +``0005_split_aggregation_schedule`` (UN-3974) then splits the aggregation into two rows by +tier. It writes both scheduler tables from one spec, so the new row cannot drift by +construction — but it also rewrites an existing row, and *how* it does that is load-bearing. +The last section covers both. + +DB-free — the migration modules are imported and their declared specs compared directly, so this runs in the unit tier rather than needing a migrated database. """ @@ -14,6 +19,7 @@ import importlib import json +from typing import Any import pytest @@ -65,7 +71,7 @@ def beat_specs() -> dict[str, dict]: captured: dict[str, dict] = {} class _Apps: - def get_model(self, _app, model): + def get_model(self, _app: str, model: str) -> type: if model == "PeriodicTask": return type("PT", (), {"objects": _FakeQuerySet(captured)}) return type("S", (), {"objects": _FakeQuerySet({})}) @@ -111,3 +117,110 @@ def test_no_spec_presets_a_run_time(self, pg_specs): for spec in pg_specs.values(): assert "next_run_at" not in spec assert "last_run_at" not in spec + + +_SPLIT_MIGRATION = "dashboard_metrics.migrations.0005_split_aggregation_schedule" +_NEW_ROW = "dashboard_metrics_aggregate_daily_monthly" +_EXISTING_ROW = "dashboard_metrics_aggregate_from_sources" + + +class _SplitRecorder: + """Captures what 0005 does to one scheduler table, keeping creates and updates apart. + + The distinction is the point: creating a row writes every default, updating one writes + only the named fields. Conflating them is exactly the bug this guards. + """ + + def __init__(self) -> None: + self.created: dict[str, dict[str, Any]] = {} + self.updated: dict[str, dict[str, Any]] = {} + self._filtered_on: str = "" + + def filter(self, name: str = "", **_kw: Any) -> _SplitRecorder: + self._filtered_on = name + return self + + def update(self, **kwargs: Any) -> int: + self.updated[self._filtered_on] = kwargs + return 1 + + def update_or_create( + self, name: str = "", defaults: dict[str, Any] | None = None, **_kw: Any + ) -> tuple[dict[str, Any], bool]: + self.created[name] = defaults or {} + return self.created[name], True + + def get_or_create(self, **kwargs: Any) -> tuple[dict[str, Any], bool]: + return kwargs, True + + def delete(self) -> tuple[int, dict[str, Any]]: + return (0, {}) + + +@pytest.fixture(scope="module") +def split() -> dict[str, _SplitRecorder]: + """Run 0005's forward function against fakes and capture both tables.""" + mod = importlib.import_module(_SPLIT_MIGRATION) + beat, pg, crontab = _SplitRecorder(), _SplitRecorder(), _SplitRecorder() + + class _Apps: + def get_model(self, _app: str, model: str) -> type: + table = {"PeriodicTask": beat, "PgPeriodicTask": pg}.get(model, crontab) + return type("M", (), {"objects": table}) + + mod.split_schedules(_Apps(), None) + return {"beat": beat, "pg": pg} + + +class TestTheSplitAddsOneRowAndRewritesOne: + def test_only_the_daily_monthly_row_is_created(self, split: dict[str, _SplitRecorder]) -> None: + for table in ("beat", "pg"): + assert set(split[table].created) == {_NEW_ROW} + + def test_only_the_existing_aggregate_row_is_updated(self, split: dict[str, _SplitRecorder]) -> None: + for table in ("beat", "pg"): + assert set(split[table].updated) == {_EXISTING_ROW} + + def test_the_new_row_is_declared_the_same_on_both_tables(self, split: dict[str, _SplitRecorder]) -> None: + beat, pg = split["beat"].created[_NEW_ROW], split["pg"].created[_NEW_ROW] + assert pg["task_name"] == beat["task"] + assert pg["queue"] == beat["queue"] + assert pg["task_kwargs"] == json.loads(beat["kwargs"]) + + def test_the_new_row_runs_hourly_on_both_tables(self, split: dict[str, _SplitRecorder]) -> None: + assert split["pg"].created[_NEW_ROW]["cron_string"] == "0 * * * *" + crontab = split["beat"].created[_NEW_ROW]["crontab"] + assert (crontab["minute"], crontab["hour"]) == ("0", "*") + + def test_the_new_row_is_seeded_inert_on_the_pg_side(self, split: dict[str, _SplitRecorder]) -> None: + """Same reason as 0004's rows: a PG row that is pg_owned before the scheduler + has adopted it would fire alongside its Beat twin. + """ + assert split["pg"].created[_NEW_ROW]["pg_owned"] is False + + def test_the_two_rows_ask_for_different_tiers(self, split: dict[str, _SplitRecorder]) -> None: + new = split["pg"].created[_NEW_ROW]["task_kwargs"]["tier"] + existing = split["pg"].updated[_EXISTING_ROW]["task_kwargs"]["tier"] + assert new != existing + + +class TestTheRewriteLeavesSchedulerOwnershipAlone: + """The existing row may already be owned by the PG scheduler, with its Beat twin + disabled by converge_pg_scheduler. Rewriting `pg_owned` or `enabled` here would + hand it back — and since the Beat twin stays disabled, the aggregation would be + left with no firer at all. Only the payload may change. + """ + + def test_the_pg_update_touches_only_the_kwargs(self, split: dict[str, _SplitRecorder]) -> None: + assert set(split["pg"].updated[_EXISTING_ROW]) == {"task_kwargs"} + + def test_the_beat_update_does_not_re_enable_the_row(self, split: dict[str, _SplitRecorder]) -> None: + assert "enabled" not in split["beat"].updated[_EXISTING_ROW] + + def test_the_existing_row_keeps_its_cadence(self, split: dict[str, _SplitRecorder]) -> None: + """Only the daily/monthly half moves to hourly; the hourly tier stays at 15 + minutes, which is the first half of the ticket's acceptance criteria. + """ + for table in ("beat", "pg"): + update = split[table].updated[_EXISTING_ROW] + assert not {"crontab", "interval", "cron_string"} & set(update) diff --git a/backend/dashboard_metrics/tests/test_tier_split_equivalence.py b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py new file mode 100644 index 0000000000..ed9b08c975 --- /dev/null +++ b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py @@ -0,0 +1,181 @@ +"""The split preserves every figure it used to write (UN-3974, AC-2). + +AC-2 is an equivalence claim — "hourly figures unchanged; daily and monthly lag by at +most one hour" — so it is settled by running the real aggregation and diffing what lands +in the metrics tables, not by reasoning about the gating predicates. Those are pinned +separately in test_aggregation_tier.py; this is the outcome they are supposed to produce. + +Two properties, and both matter: + +- the `hourly` schedule reproduces what the single pre-split run wrote to + EventMetricsHourly, exactly — that is the "unchanged" half +- `hourly` and `daily_monthly` together reproduce every row the pre-split run wrote to + any table — that is the "nothing is lost" half, which the AC assumes rather than states + +DB-bound, so conftest marks it integration. +""" + +from __future__ import annotations + +import uuid +from datetime import timedelta +from typing import Any + +from account_v2.models import Organization +from django.db import connection +from django.test import TestCase +from django.utils import timezone +from workflow_manager.workflow_v2.models.workflow import Workflow + +from dashboard_metrics.models import ( + EventMetricsDaily, + EventMetricsHourly, + EventMetricsMonthly, +) +from dashboard_metrics.tasks import AggregationTier, _run_aggregation + +# (model, the column naming its period) — the period field differs per tier. +_TIERS = [ + (EventMetricsHourly, "timestamp"), + (EventMetricsDaily, "date"), + (EventMetricsMonthly, "month"), +] +_FIELDS = ["metric_name", "metric_type", "metric_value", "metric_count"] + + +class TestTheSplitPreservesEveryFigure(TestCase): + def setUp(self) -> None: + self.org = Organization.objects.create( + organization_id="tier-split-org", name="tier-split", display_name="Tier Split" + ) + self.workflow = Workflow.objects.create( + workflow_name="tier-split-wf", organization=self.org + ) + now = timezone.now() + # One row per window the aggregation reads — last 24h for the hourly tier, last + # 7 days for daily, inside the previous month for monthly. The two recent ones + # also make the org visible to the active-org prefilter, without which nothing + # runs at all. + windows = [ + now - timedelta(hours=2), + now - timedelta(hours=5), + now - timedelta(days=3), + now - timedelta(days=25), + ] + executions = self._add_executions(windows) + # Both aggregation paths have to be exercised: the per-metric queries go through + # _aggregate_single_metric and the four LLM metrics through + # _aggregate_llm_combined, and each gates on the tier separately. A fixture + # producing only LLM figures leaves half the split unverified. + self._add_file_executions(executions) + self._add_llm_usage(windows) + + def _add_executions(self, timestamps: list[Any]) -> list[tuple[Any, Any]]: + """Raw insert so created_at is ours; the model sets it with auto_now_add.""" + created = [] + with connection.cursor() as cur: + for ts in timestamps: + execution_id = uuid.uuid4() + cur.execute( + "INSERT INTO workflow_execution (id, created_at, modified_at, " + "workflow_id, execution_mode, execution_method, execution_type, " + "execution_log_id, status, error_message, attempts, execution_time, " + "result_acknowledged, total_files) " + "VALUES (%s, %s, %s, %s, 'INSTANT', 'DIRECT', 'COMPLETE', '', " + "'COMPLETED', '', 0, 1.0, false, 1)", + [execution_id, ts, ts, self.workflow.id], + ) + created.append((execution_id, ts)) + return created + + def _add_file_executions(self, executions: list[tuple[Any, Any]]) -> None: + """Feeds documents_processed, which runs through _aggregate_single_metric.""" + with connection.cursor() as cur: + for execution_id, ts in executions: + cur.execute( + "INSERT INTO workflow_file_execution (id, created_at, modified_at, " + "file_name, status, workflow_execution_id) " + "VALUES (%s, %s, %s, 'doc.pdf', 'COMPLETED', %s)", + [uuid.uuid4(), ts, ts, execution_id], + ) + + def _add_llm_usage(self, timestamps: list[Any]) -> None: + """LLM metrics need no joins, so they are the cheapest way to put a real figure + in all three tiers.""" + with connection.cursor() as cur: + for ts in timestamps: + cur.execute( + "INSERT INTO usage (id, created_at, modified_at, adapter_instance_id, " + "usage_type, llm_usage_reason, model_name, embedding_tokens, " + "prompt_tokens, completion_tokens, total_tokens, cost_in_dollars, " + "organization_id) " + "VALUES (%s, %s, %s, 'test-adapter', 'llm', 'extraction', 'test-model', " + "0, 100, 50, 150, 0.25, %s)", + [uuid.uuid4(), ts, ts, self.org.id], + ) + + def _snapshot(self) -> dict[str, set[tuple[Any, ...]]]: + return { + model.__name__: set( + model._base_manager.values_list("organization_id", period, *_FIELDS) + ) + for model, period in _TIERS + } + + def _clear(self) -> None: + for model, _ in _TIERS: + model._base_manager.all().delete() + + def _run(self, tier: AggregationTier) -> dict[str, set[tuple[Any, ...]]]: + self._clear() + _run_aggregation(tier) + return self._snapshot() + + def test_the_pre_split_run_writes_all_three_tiers(self) -> None: + """Guards the tests below from passing vacuously: an equivalence between two + empty sets proves nothing. + """ + every_tier = self._run(AggregationTier.ALL) + for name, rows in every_tier.items(): + assert rows, f"{name} is empty — the fixture produces no metrics to compare" + + def test_the_fixture_exercises_both_aggregation_paths(self) -> None: + """The other way these tests can go quietly vacuous. The tier is checked + separately in _aggregate_single_metric and in _aggregate_llm_combined, so a + fixture yielding only one kind of metric verifies only half the split — which + is exactly what a mutation test caught here. + """ + every_tier = self._run(AggregationTier.ALL) + for table, rows in every_tier.items(): + names = {row[2] for row in rows} + assert "documents_processed" in names, f"{table}: no per-metric figure" + assert "llm_calls" in names, f"{table}: no combined-LLM figure" + + def test_hourly_reproduces_the_pre_split_hourly_figures(self) -> None: + """The "figures unchanged" half of AC-2, row for row rather than in aggregate.""" + before = self._run(AggregationTier.ALL)["EventMetricsHourly"] + after = self._run(AggregationTier.HOURLY)["EventMetricsHourly"] + assert after == before + + def test_the_two_schedules_together_lose_nothing(self) -> None: + """Every row the single pre-split run wrote is still written by one of the two + schedules, and neither invents one. + """ + every_tier = self._run(AggregationTier.ALL) + hourly = self._run(AggregationTier.HOURLY) + daily_monthly = self._run(AggregationTier.DAILY_MONTHLY) + + for name in every_tier: + combined = hourly[name] | daily_monthly[name] + assert combined == every_tier[name], f"{name} differs after the split" + + def test_neither_schedule_writes_the_other_tiers_tables(self) -> None: + """If they overlapped, the two schedules would duplicate work every hour on the + hour — harmless thanks to the upserts, but not free. + """ + hourly = self._run(AggregationTier.HOURLY) + assert not hourly["EventMetricsDaily"] + assert not hourly["EventMetricsMonthly"] + + daily_monthly = self._run(AggregationTier.DAILY_MONTHLY) + assert not daily_monthly["EventMetricsHourly"] diff --git a/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py b/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py new file mode 100644 index 0000000000..93317989eb --- /dev/null +++ b/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py @@ -0,0 +1,110 @@ +"""Guard: ``we_created_at_idx`` keeps the shape that makes it safe to deploy. + +The backend suite runs with ``--no-migrations``, so migration 0029 never executes in +CI. Regenerating it with ``makemigrations``, or dropping ``atomic = False`` while +tidying, lands a plain ``AddIndex`` — which holds a SHARE lock for the whole build and +blocks every in-flight execution on a multi-million-row table — with every other test +still green. These assert the properties that keep that from happening. + +Model and migration introspection only, no test database, so this runs in the unit tier +alongside ``test_active_execution_index.py`` and ``test_undispatched_execution_index.py``. +""" + +from __future__ import annotations + +import importlib +import os +import re +from pathlib import Path +from typing import Any, cast + +import django +from django.apps import apps +from django.db import migrations, models + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +INDEX_NAME = "we_created_at_idx" +_MIGRATION_FILE = ( + Path(__file__).resolve().parent.parent / "migrations" / "0029_we_created_at_idx.py" +) +_MIGRATION_MODULE = "workflow_manager.workflow_v2.migrations.0029_we_created_at_idx" + + +def _model_index() -> models.Index | None: + model = apps.get_model("workflow_v2", "WorkflowExecution") + return next((i for i in model._meta.indexes if i.name == INDEX_NAME), None) + + +def _operations() -> list[Any]: + return cast( + list[Any], importlib.import_module(_MIGRATION_MODULE).Migration.operations + ) + + +class TestTheModelDeclaresIt: + def test_it_is_keyed_on_created_at_alone(self) -> None: + """A bare created_at range with no leading column value is the whole point — + the composite indexes lead with workflow_id / pipeline_id and are date-ordered + only within one workflow or pipeline. + """ + index = _model_index() + assert index is not None, f"{INDEX_NAME} is missing from WorkflowExecution.Meta" + assert index.fields == ["created_at"] + + def test_it_carries_no_condition(self) -> None: + """A partial index would not serve the prefilter, which bounds nothing but the + date. we_undispatched_dispatch_idx is the partial one and is a different index. + """ + assert getattr(_model_index(), "condition", None) is None + + +class TestTheMigrationIsSafeToDeploy: + def test_it_is_non_atomic(self) -> None: + """CREATE/DROP INDEX CONCURRENTLY cannot run inside a transaction block, so + without this the migration cannot run at all. + """ + assert re.search( + r"^\s*atomic\s*=\s*False", _MIGRATION_FILE.read_text(), re.MULTILINE + ) + + def test_it_builds_and_drops_concurrently(self) -> None: + """Both directions: a plain DROP INDEX takes an ACCESS EXCLUSIVE lock, so a + rollback would block writes just as a plain build would. + """ + sql = _MIGRATION_FILE.read_text() + assert "CREATE INDEX CONCURRENTLY IF NOT EXISTS" in sql + assert "DROP INDEX CONCURRENTLY IF EXISTS" in sql + + def test_it_guards_against_a_leftover_invalid_index(self) -> None: + """An interrupted CONCURRENTLY build leaves an INVALID index that costs on every + write and is never read. IF NOT EXISTS would keep it while Django recorded the + migration as applied — green, and permanently slower. + """ + sql = _MIGRATION_FILE.read_text() + assert "RAISE EXCEPTION" in sql + assert "indisvalid" in sql + + def test_add_index_is_state_only(self) -> None: + """The failure mode this whole file exists for. AddIndex outside + state_operations is a real lock-taking build; inside, it only keeps Django's + model state in step so makemigrations does not re-add the index. + """ + ops = _operations() + assert len(ops) == 1 + wrapper = ops[0] + assert isinstance(wrapper, migrations.SeparateDatabaseAndState) + assert all( + isinstance(op, migrations.RunSQL) for op in wrapper.database_operations + ) + assert [type(op) for op in wrapper.state_operations] == [migrations.AddIndex] + + def test_the_migration_and_the_model_agree(self) -> None: + """Two declarations of one index; they must not drift.""" + index = _model_index() + assert index is not None + add_index = _operations()[0].state_operations[0] + assert add_index.index.name == INDEX_NAME + assert add_index.index.fields == index.fields diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py index ce8ff853ac..055a7725d5 100644 --- a/workers/tests/test_dashboard_metrics_tasks.py +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -67,6 +67,27 @@ def test_aggregate_posts_to_the_aggregate_endpoint(self): dmt.dashboard_metrics_aggregate() assert call.call_args[0][0] == "v1/dashboard-metrics/aggregate/" + @pytest.mark.parametrize("tier", ["hourly", "daily_monthly", "all"]) + def test_aggregate_forwards_the_tier_from_the_schedule_row(self, tier): + """UN-3974: the PG scheduler hands a row's task_kwargs over as **kwargs, so the + tier arrives here and has to reach the backend in the request body. + + This is the leg that fails quietly. Drop the forwarding and every schedule still + fires, the endpoint still returns 200, and every other test here still passes — + but both rows run the default tier, so daily and monthly quietly go back to + being recomputed every 15 minutes. + """ + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate(tier=tier) + assert call.call_args.kwargs["body"] == {"tier": tier} + + def test_aggregate_omits_the_body_when_no_tier_is_given(self): + # Pre-0005 rows carry no tier kwarg; the backend's default then applies, which + # is every tier rather than none. + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_args.kwargs["body"] is None + @pytest.mark.parametrize( "func,path", [ From 01e01f0873d9495a6452d1847bd7058759ba0762 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 1 Sep 2026 19:20:50 +0530 Subject: [PATCH 09/14] UN-3973 Renumber the reconciliation migration to 0005 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit UN-3445 landed 0004_pg_periodic_tasks on main after this branch was cut, leaving dashboard_metrics with two 0004s depending on 0003 and nothing depending on either. Django saw two leaf nodes and refused to build the graph, so `migrate` failed before applying anything — every app, not just this one. Depend on 0004_pg_periodic_tasks and renumber to match, so the short prefix form stays usable for a rollback. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- ...d_reconciliation_task.py => 0005_add_reconciliation_task.py} | 2 +- backend/dashboard_metrics/tests/test_tasks.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) rename backend/dashboard_metrics/migrations/{0004_add_reconciliation_task.py => 0005_add_reconciliation_task.py} (95%) diff --git a/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py b/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py similarity index 95% rename from backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py rename to backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py index 09667e867f..8a44759eed 100644 --- a/backend/dashboard_metrics/migrations/0004_add_reconciliation_task.py +++ b/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py @@ -49,7 +49,7 @@ def remove_reconciliation_task(apps, schema_editor): class Migration(migrations.Migration): dependencies = [ - ("dashboard_metrics", "0003_alter_eventmetricsdaily_organization_and_more"), + ("dashboard_metrics", "0004_pg_periodic_tasks"), ("django_celery_beat", "0018_improve_crontab_helptext"), ] diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index 6cb9316d1c..822c197a63 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -495,7 +495,7 @@ class TestReconciliationSchedule(TestCase): def setUp(self): """Load the data migration module.""" self.migration = import_module( - "dashboard_metrics.migrations.0004_add_reconciliation_task" + "dashboard_metrics.migrations.0005_add_reconciliation_task" ) def _task(self): From cc060938f46aaef705685efcbd116b5909c5cd1e Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Tue, 1 Sep 2026 19:21:05 +0530 Subject: [PATCH 10/14] UN-3974 Renumber the schedule-split migration to 0006 behind UN-3973's 0005 UN-3445's 0004_pg_periodic_tasks is the parent of both this migration and UN-3973's reconciliation migration, so landing both would leave dashboard_metrics with two leaf nodes and no applicable graph. Depend on 0005_add_reconciliation_task instead, which puts the intended merge order (UN-3973 then UN-3974) in the graph rather than in the merge queue. This branch cannot migrate on its own until UN-3973 lands; its tests are unaffected, since the suite runs with --no-migrations. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- ...egation_schedule.py => 0006_split_aggregation_schedule.py} | 2 +- backend/dashboard_metrics/tests/test_aggregation_dispatch.py | 2 +- .../tests/test_pg_periodic_task_declarations.py | 4 ++-- 3 files changed, 4 insertions(+), 4 deletions(-) rename backend/dashboard_metrics/migrations/{0005_split_aggregation_schedule.py => 0006_split_aggregation_schedule.py} (98%) diff --git a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py similarity index 98% rename from backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py rename to backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py index 9d6b40454f..a9cbf1aabf 100644 --- a/backend/dashboard_metrics/migrations/0005_split_aggregation_schedule.py +++ b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py @@ -127,7 +127,7 @@ def merge_schedules(apps, schema_editor): class Migration(migrations.Migration): dependencies = [ - ("dashboard_metrics", "0004_pg_periodic_tasks"), + ("dashboard_metrics", "0005_add_reconciliation_task"), ("django_celery_beat", "0018_improve_crontab_helptext"), ("pg_queue", "0003_pgperiodictask"), ] diff --git a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py index c4892ea59d..060f43e88e 100644 --- a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py +++ b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py @@ -38,7 +38,7 @@ aggregate_metrics_from_sources, ) -_SPLIT_MIGRATION = "dashboard_metrics.migrations.0005_split_aggregation_schedule" +_SPLIT_MIGRATION = "dashboard_metrics.migrations.0006_split_aggregation_schedule" _ENDPOINT = "/internal/v1/dashboard-metrics/aggregate/" diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py index 7e1257fd60..ee90e93375 100644 --- a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -6,7 +6,7 @@ the whole failure mode this file exists for: a schedule changed on Beat but not on PG means the task silently runs on a different cadence the moment the flag flips. -``0005_split_aggregation_schedule`` (UN-3974) then splits the aggregation into two rows by +``0006_split_aggregation_schedule`` (UN-3974) then splits the aggregation into two rows by tier. It writes both scheduler tables from one spec, so the new row cannot drift by construction — but it also rewrites an existing row, and *how* it does that is load-bearing. The last section covers both. @@ -119,7 +119,7 @@ def test_no_spec_presets_a_run_time(self, pg_specs): assert "last_run_at" not in spec -_SPLIT_MIGRATION = "dashboard_metrics.migrations.0005_split_aggregation_schedule" +_SPLIT_MIGRATION = "dashboard_metrics.migrations.0006_split_aggregation_schedule" _NEW_ROW = "dashboard_metrics_aggregate_daily_monthly" _EXISTING_ROW = "dashboard_metrics_aggregate_from_sources" From b974eaea5cdf8db30519eb19414167148e0ee374 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 2 Sep 2026 16:49:15 +0530 Subject: [PATCH 11/14] UN-3973 [FIX] Address review: PG transport, scoped orphan sweep, retry posture MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The reconciliation row could not run on the PG transport — two functions share the task name dashboard_metrics.aggregate_from_sources and the worker one took no arguments, so the mirrored row dispatched source_window_days into a zero-arg function and the message was dropped. The worker proxy and the internal endpoint now plumb it, and 0005 declares the PG twin rather than leaving the mirror to invent one. The orphan sweep is scoped to the (organization, month) partitions the rollup actually produced: an incomplete daily tier passed the empty-tier guard and deleted monthly rows it could not vouch for. Its deletion count now reaches the task result and a WARNING. DatabaseError and OperationalError propagate from the monthly rollup so the configured autoretry fires, instead of being logged once behind success: True. The prefilter is never narrower than the query window, so a widened source_window_days cannot skip the orgs it exists to repair. bulk_create takes an explicit batch_size. The Beat/PG drift guard named 0002 and 0004, so it kept comparing three schedules against three while this PR added a fourth. It now discovers every migration in the app, replays their RunPython forwards in order, derives the Beat cadence from the schedule row, binds every declared kwarg to its task signature, and asserts every post-install Beat write bumps PeriodicTasks.last_update. The reconcile row moves to 04:40 UTC: */15 fires at minute 0, and the per-tier lock keys do not block each other, so 04:00 started two full aggregations at once. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/README.md | 28 +- backend/dashboard_metrics/internal_views.py | 13 +- .../management/commands/backfill_metrics.py | 25 +- .../0005_add_reconciliation_task.py | 99 ++++- backend/dashboard_metrics/tasks.py | 115 ++++-- .../test_pg_periodic_task_declarations.py | 279 +++++++++---- backend/dashboard_metrics/tests/test_tasks.py | 368 ++++++++++++++++-- workers/scheduler/dashboard_metrics_tasks.py | 17 +- workers/tests/test_dashboard_metrics_tasks.py | 27 ++ 9 files changed, 802 insertions(+), 169 deletions(-) diff --git a/backend/dashboard_metrics/README.md b/backend/dashboard_metrics/README.md index fc40e9f55b..d4f77b9eb7 100644 --- a/backend/dashboard_metrics/README.md +++ b/backend/dashboard_metrics/README.md @@ -46,7 +46,8 @@ celery -A backend beat -l info ### Celery Tasks & Schedule | Task | Schedule | What It Does | |------|----------|--------------| -| `aggregate_from_sources` | Every 15 min | Aggregates source → hourly/daily/monthly | +| `aggregate_from_sources` | Every 15 min | Aggregates source → hourly/daily; rolls monthly up from daily | +| `aggregate_from_sources` (reconcile) | Daily 4:40 AM | Same task over a 7-day source window, to repair gaps after downtime | | `cleanup_hourly_data` | Daily 2 AM | Deletes hourly data > 30 days | | `cleanup_daily_data` | Weekly Sun 3 AM | Deletes daily data > 365 days | @@ -109,7 +110,7 @@ celery -A backend beat -l info │ EventMetrics │ │ EventMetrics │ │ EventMetrics │ │ Hourly │ │ Daily │ │ Monthly │ │ │ │ │ │ │ -│ • 24h query │ │ • 7 day query │ │ • 2 month query │ +│ • 24h query │ │ • 2 day query │ │ • from daily │ │ • 30 day retain │ │ • 365 day retain│ │ • No cleanup │ └────────┬────────┘ └────────┬────────┘ └────────┬────────┘ │ │ │ @@ -163,6 +164,7 @@ The dashboard reads from **pre-aggregated tables** (`event_metrics_hourly`, `eve **Failure Resilience:** - If the aggregation task fails, the dashboard shows stale data (up to 15 minutes old) rather than crashing. +- A daily 04:40 UTC reconciliation pass reruns the same task over a 7-day source window, so a gap shorter than that repairs itself without a manual backfill. - Celery tasks have `max_retries=3` with exponential backoff. - Cleanup tasks (hourly: 30-day retention, daily: 365-day retention) prevent unbounded table growth. @@ -298,8 +300,8 @@ cost = (input_cost_per_token × input_tokens) + (output_cost_per_token × output | Table | Model | Time Column | Granularity | Query Window | Retention | |-------|-------|-------------|-------------|--------------|-----------| | `event_metrics_hourly` | `EventMetricsHourly` | `timestamp` | Hour | Last 24 hours | 30 days | -| `event_metrics_daily` | `EventMetricsDaily` | `date` | Day | Last 7 days | 365 days | -| `event_metrics_monthly` | `EventMetricsMonthly` | `month` | Month | Last 2 months | Forever | +| `event_metrics_daily` | `EventMetricsDaily` | `date` | Day | Last 2 days (7 on the daily reconciliation pass) | 365 days | +| `event_metrics_monthly` | `EventMetricsMonthly` | `month` | Month | Rolled up from the daily tier, current + previous month | Forever | ### Table Schema @@ -340,6 +342,7 @@ Located in `tasks.py`: | Task Name | Celery Name | Schedule | Queue | Purpose | |-----------|-------------|----------|-------|---------| | `aggregate_metrics_from_sources` | `dashboard_metrics.aggregate_from_sources` | Every 15 min | `dashboard_metric_events` | Aggregate from source tables | +| `aggregate_metrics_from_sources` | `dashboard_metrics.aggregate_from_sources` | Daily 4:40 AM UTC | `dashboard_metric_events` | Reconciliation pass, `source_window_days=7` | | `cleanup_hourly_metrics` | `dashboard_metrics.cleanup_hourly_data` | Daily 2:00 AM UTC | `dashboard_metric_events` | Delete hourly data >30 days | | `cleanup_daily_metrics` | `dashboard_metrics.cleanup_daily_data` | Weekly Sun 3:00 AM UTC | `dashboard_metric_events` | Delete daily data >365 days | @@ -378,16 +381,23 @@ The `aggregate_metrics_from_sources` task: 2. **For each metric**: - Queries source table with `MetricsQueryService` - Groups by time period (hour/day/month) -3. **Upserts results** into aggregated tables using `update_or_create` -4. **Uses `_base_manager`** to bypass Django's organization filter in Celery context +3. **Upserts results** into the hourly and daily tables +4. **Rolls monthly up from the daily tier** in one statement for all orgs, dropping + monthly rows the daily tier no longer produces (scoped to the organization/month + partitions the rollup covered) +5. **Uses `_base_manager`** to bypass Django's organization filter in Celery context ```python # Query windows -hourly_start = end_date - timedelta(hours=24) # Last 24 hours -daily_start = end_date - timedelta(days=7) # Last 7 days -monthly_start = first_of_previous_month # Last 2 months +hourly_start = end_date - timedelta(hours=24) # Last 24 hours +daily_start = truncate_to_day(end_date - source_window_days) # 2 days, 7 on reconcile +monthly_start = first_of_previous_month # summed from daily ``` +The monthly tier has no source queries of its own. `backfill_metrics` still computes +monthly from source, so within the rollup window (current + previous month) its output +is recomputed within 15 minutes — see that command's help text. + --- ## API Endpoints diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index f776633944..0ecc6759f2 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -34,6 +34,7 @@ from utils.local_context import StateStore from dashboard_metrics.tasks import ( + DASHBOARD_SOURCE_WINDOW_DAYS, aggregate_metrics_from_sources, cleanup_daily_metrics, cleanup_hourly_metrics, @@ -94,7 +95,17 @@ class AggregateMetricsAPIView(_MetricsTaskAPIView): """ def post(self, request: Request) -> Response: - return self._run(aggregate_metrics_from_sources) + """``source_window_days`` is optional: the 15-minute schedule omits it and + gets the task's default, the daily reconciliation pass widens it. + """ + body = request.data if isinstance(request.data, dict) else {} + if "source_window_days" not in body: + return self._run(aggregate_metrics_from_sources) + try: + days = _int_arg(request, "source_window_days", DASHBOARD_SOURCE_WINDOW_DAYS) + except ValueError as exc: + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + return self._run(aggregate_metrics_from_sources, source_window_days=days) class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView): diff --git a/backend/dashboard_metrics/management/commands/backfill_metrics.py b/backend/dashboard_metrics/management/commands/backfill_metrics.py index 9c4d82baca..d598e366bb 100644 --- a/backend/dashboard_metrics/management/commands/backfill_metrics.py +++ b/backend/dashboard_metrics/management/commands/backfill_metrics.py @@ -3,6 +3,12 @@ This command populates EventMetricsHourly, EventMetricsDaily, and EventMetricsMonthly tables from historical data in source tables (Usage, PageUsage, WorkflowExecution, etc.) +The current and previous month are owned exclusively by the aggregation task, which +derives them from the daily tier and drops monthly rows the daily tier no longer +produces. Inside that window this command's monthly output is recomputed within 15 +minutes, so --skip-monthly is a no-op and --skip-daily leaves monthly to be rebuilt +from a tier this run did not populate. Backfill both, or neither. + Usage: python manage.py backfill_metrics --days=30 python manage.py backfill_metrics --days=90 --org-id=5 @@ -92,12 +98,18 @@ def add_arguments(self, parser): parser.add_argument( "--skip-daily", action="store_true", - help="Skip daily aggregation", + help=( + "Skip daily aggregation. Unsafe for the current and previous month: " + "the aggregation task rebuilds monthly from daily there." + ), ) parser.add_argument( "--skip-monthly", action="store_true", - help="Skip monthly aggregation", + help=( + "Skip monthly aggregation. A no-op for the current and previous " + "month, which the aggregation task owns." + ), ) parser.add_argument( "--active-only", @@ -123,6 +135,15 @@ def handle(self, *args, **options): self.stdout.write(f"Backfill period: {start_date.date()} to {end_date.date()}") self.stdout.write(f"Days: {days}") + if skip_daily and not skip_monthly: + self.stdout.write( + self.style.WARNING( + "--skip-daily without --skip-monthly: the aggregation task " + "rebuilds the current and previous month from the daily tier " + "and will drop monthly rows this run writes there." + ) + ) + if dry_run: self.stdout.write(self.style.WARNING("DRY RUN - no changes will be made")) diff --git a/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py b/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py index 8a44759eed..d7bf134aab 100644 --- a/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py +++ b/backend/dashboard_metrics/migrations/0005_add_reconciliation_task.py @@ -3,21 +3,47 @@ The 15-minute aggregation reads a narrow source window, which cannot repair gaps left by cron downtime. This runs the same task once a day at a wider window to backfill them. + +Declared for **both** transports, like 0002/0004: Beat reads +``django_celery_beat_periodictask``, the PG scheduler reads ``pg_periodic_task``, +and a schedule present on one only stops firing the moment the flag flips. +``kwargs`` is a JSON string on Beat and a JSONField on PG — same value, two +encodings. """ from django.db import migrations +from django.utils import timezone RECONCILE_TASK_NAME = "dashboard_metrics_reconcile_source_window" +RECONCILE_DESCRIPTION = ( + "Re-aggregate metrics over a 7 day source window to repair " + "daily-tier gaps left by cron downtime" +) + +# Single source for both directions, and importable by the drift test. +PG_PERIODIC_TASKS = [ + { + "name": RECONCILE_TASK_NAME, + "task_name": "dashboard_metrics.aggregate_from_sources", + "queue": "dashboard_metric_events", + "task_args": [], + "task_kwargs": {"source_window_days": 7}, + # Beat: CrontabSchedule(minute=40, hour=4, every day) UTC — clear of the + # 2:00 and 3:00 cleanup tasks, and off the aggregation's */15 grid + # (:00 :15 :30 :45) so the two never start together. + "cron_string": "40 4 * * *", + }, +] def create_reconciliation_task(apps, schema_editor): - """Create the once-daily reconciliation periodic task.""" + """Create the once-daily reconciliation periodic task on both transports.""" crontab_model = apps.get_model("django_celery_beat", "CrontabSchedule") periodic_task_model = apps.get_model("django_celery_beat", "PeriodicTask") + pg_periodic_task_model = apps.get_model("pg_queue", "PgPeriodicTask") - # 4:00 AM UTC — clear of the 2:00 and 3:00 cleanup tasks schedule_4am, _ = crontab_model.objects.get_or_create( - minute="0", + minute="40", hour="4", day_of_week="*", day_of_month="*", @@ -25,32 +51,65 @@ def create_reconciliation_task(apps, schema_editor): defaults={"timezone": "UTC"}, ) - periodic_task_model.objects.update_or_create( - name=RECONCILE_TASK_NAME, - defaults={ - "task": "dashboard_metrics.aggregate_from_sources", - "crontab": schedule_4am, - "queue": "dashboard_metric_events", - "kwargs": '{"source_window_days": 7}', - "enabled": True, - "description": ( - "Re-aggregate metrics over a 7 day source window to repair " - "daily-tier gaps left by cron downtime" - ), - }, - ) + for spec in PG_PERIODIC_TASKS: + periodic_task_model.objects.update_or_create( + name=spec["name"], + defaults={ + "task": spec["task_name"], + "crontab": schedule_4am, + "queue": spec["queue"], + "kwargs": '{"source_window_days": 7}', + "enabled": True, + "description": RECONCILE_DESCRIPTION, + }, + ) + pg_periodic_task_model.objects.update_or_create( + name=spec["name"], + defaults={ + "task_name": spec["task_name"], + "queue": spec["queue"], + "task_args": spec["task_args"], + "task_kwargs": spec["task_kwargs"], + "cron_string": spec["cron_string"], + "org_id": "", + "enabled": True, + # Inert until the rollout flag decides otherwise. + "pg_owned": False, + }, + ) + + _bump_beat_change_tracker(apps) def remove_reconciliation_task(apps, schema_editor): - """Remove the reconciliation periodic task on rollback.""" - periodic_task_model = apps.get_model("django_celery_beat", "PeriodicTask") - periodic_task_model.objects.filter(name=RECONCILE_TASK_NAME).delete() + """Remove the reconciliation periodic task from both transports.""" + names = [spec["name"] for spec in PG_PERIODIC_TASKS] + apps.get_model("django_celery_beat", "PeriodicTask").objects.filter( + name__in=names + ).delete() + apps.get_model("pg_queue", "PgPeriodicTask").objects.filter(name__in=names).delete() + _bump_beat_change_tracker(apps) + + +def _bump_beat_change_tracker(apps): + """Make a running Beat reload instead of missing the new schedule. + + django-celery-beat's post_save receiver binds the concrete PeriodicTask, so + writes through a historical model never bump PeriodicTasks.last_update and + DatabaseScheduler keeps its stale in-memory copy. Same fix and reason as + scheduler/ownership.py and mirror_pg_periodic_tasks.py. + """ + periodic_tasks_model = apps.get_model("django_celery_beat", "PeriodicTasks") + periodic_tasks_model.objects.update_or_create( + ident=1, defaults={"last_update": timezone.now()} + ) class Migration(migrations.Migration): dependencies = [ ("dashboard_metrics", "0004_pg_periodic_tasks"), ("django_celery_beat", "0018_improve_crontab_helptext"), + ("pg_queue", "0003_pgperiodictask"), ] operations = [ diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index c0b1b1a9b8..332cb623ef 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -32,6 +32,10 @@ logger = logging.getLogger(__name__) +# Django 4.2's PostgreSQL backend does not override bulk_batch_size, so an +# unbatched bulk_create emits one statement whose size scales with tenant count. +MONTHLY_ROLLUP_BATCH_SIZE = 1000 + # Retention periods for metrics cleanup DASHBOARD_HOURLY_METRICS_RETENTION_DAYS = 30 DASHBOARD_DAILY_METRICS_RETENTION_DAYS = 365 @@ -40,7 +44,9 @@ # created_at -> terminal-status lag. DASHBOARD_SOURCE_WINDOW_DAYS = 2 -# Wider lookback for the once-daily reconciliation pass. +# Wider lookback for the once-daily reconciliation pass. A migration must not +# import live app code, so 0005_add_reconciliation_task carries this as a literal +# in the schedule row's kwargs — editing this constant does not move the schedule. DASHBOARD_RECONCILE_WINDOW_DAYS = 7 # Wider than the source window: metrics keyed on another column @@ -179,39 +185,54 @@ def _bulk_upsert_daily(aggregations: dict) -> int: return len(objects) -def _delete_orphan_monthly(month_start: date, fresh_keys: set[tuple]) -> int: - """Drop monthly rows from month_start that the rollup no longer produces. +def _delete_orphan_monthly(objects: list[EventMetricsMonthly]) -> int: + """Drop stale monthly rows inside the partitions the rollup actually covered. Monthly derives from the daily tier, so a key with no daily rows left must - not survive as a stale total. + not survive as a stale total. An (organization, month) the rollup produced + nothing for is left alone instead: absent daily rows there mean an + incomplete tier, not a metric that went to zero. """ - stale_pks = [ - row["pk"] - for row in EventMetricsMonthly._base_manager.filter( - month__gte=month_start - ).values("pk", "organization_id", "month", "metric_name", "project", "tag") - if ( - row["organization_id"], - row["month"], - row["metric_name"], - row["project"], - row["tag"], - ) - not in fresh_keys - ] - if not stale_pks: - return 0 + fresh_keys = { + (o.organization_id, o.month, o.metric_name, o.project, o.tag) for o in objects + } + covered: dict[date, set] = {} + for o in objects: + covered.setdefault(o.month, set()).add(o.organization_id) + + deleted = 0 + for month, org_ids in covered.items(): + stale_pks = [ + row["pk"] + for row in EventMetricsMonthly._base_manager.filter( + month=month, organization_id__in=org_ids + ).values("pk", "organization_id", "month", "metric_name", "project", "tag") + if ( + row["organization_id"], + row["month"], + row["metric_name"], + row["project"], + row["tag"], + ) + not in fresh_keys + ] + if not stale_pks: + continue + removed, _ = EventMetricsMonthly._base_manager.filter(pk__in=stale_pks).delete() + deleted += removed - deleted, _ = EventMetricsMonthly._base_manager.filter(pk__in=stale_pks).delete() return deleted -def _rollup_monthly_from_daily(month_start: date) -> int: +def _rollup_monthly_from_daily(month_start: date) -> tuple[int, int]: """Sum the daily tier from month_start into monthly, for all orgs at once. metric_type is aggregated rather than grouped: it is not part of unique_monthly_metric, so grouping on it could yield two rows for one conflict target. + + Returns: + (rows upserted, rows deleted as orphans) """ rows = ( EventMetricsDaily._base_manager.filter(date__gte=month_start) @@ -238,14 +259,11 @@ def _rollup_monthly_from_daily(month_start: date) -> int: for row in rows ] - # An empty daily tier means the source is gone, not that every month is - # zero — leave existing rows alone. + # Nothing to write and nothing covered, so nothing to sweep. The scoping in + # _delete_orphan_monthly already makes this safe; returning early just skips a + # pointless transaction. if not objects: - return 0 - - fresh_keys = { - (o.organization_id, o.month, o.metric_name, o.project, o.tag) for o in objects - } + return 0, 0 with transaction.atomic(): EventMetricsMonthly._base_manager.bulk_create( @@ -253,10 +271,11 @@ def _rollup_monthly_from_daily(month_start: date) -> int: update_conflicts=True, unique_fields=["organization", "month", "metric_name", "project", "tag"], update_fields=["metric_type", "metric_value", "metric_count"], + batch_size=MONTHLY_ROLLUP_BATCH_SIZE, ) - _delete_orphan_monthly(month_start, fresh_keys) + deleted = _delete_orphan_monthly(objects) - return len(objects) + return len(objects), deleted AGGREGATION_LOCK_KEY = "dashboard_metrics:aggregation_lock" @@ -525,12 +544,19 @@ def _aggregate_org( stats["orgs_processed"] += 1 -def _active_org_ids(end_date: datetime) -> set: - """Organizations with execution activity in the prefilter lookback.""" +def _active_org_ids(end_date: datetime, window_start: datetime) -> set: + """Organizations with execution activity in the prefilter lookback. + + Never narrower than the caller's own query window: a widened + source_window_days must not be prefiltered back down to the default + lookback, or the reconciliation pass skips the orgs it exists to repair. + """ + cutoff = min( + window_start, + end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), + ) return set( - WorkflowExecution.objects.filter( - created_at__gte=end_date - timedelta(days=DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS), - ) + WorkflowExecution.objects.filter(created_at__gte=cutoff) .values_list("workflow__organization_id", flat=True) .distinct() ) @@ -579,13 +605,13 @@ def _run_aggregation( stats = { "hourly": {"upserted": 0}, "daily": {"upserted": 0}, - "monthly": {"upserted": 0}, + "monthly": {"upserted": 0, "deleted": 0}, "errors": 0, "orgs_processed": 0, } # Pre-filter to orgs with recent activity to reduce DB load. - active_org_ids = _active_org_ids(end_date) + active_org_ids = _active_org_ids(end_date, daily_start) logger.info( "Aggregation: %d active orgs out of %d total", len(active_org_ids), @@ -614,7 +640,17 @@ def _run_aggregation( stats["errors"] += 1 try: - stats["monthly"]["upserted"] = _rollup_monthly_from_daily(monthly_start) + upserted, deleted = _rollup_monthly_from_daily(monthly_start) + stats["monthly"]["upserted"] = upserted + stats["monthly"]["deleted"] = deleted + if deleted: + logger.warning( + "Monthly rollup deleted %d orphan row(s) from %s", deleted, monthly_start + ) + except (DatabaseError, OperationalError): + # Configured on the task for autoretry — swallowing them here would + # leave monthly permanently stale behind successful-looking runs. + raise except Exception: logger.exception("Error rolling up monthly metrics from %s", monthly_start) stats["errors"] += 1 @@ -624,6 +660,7 @@ def _run_aggregation( f"hourly={stats['hourly']['upserted']}, " f"daily={stats['daily']['upserted']}, " f"monthly={stats['monthly']['upserted']}, " + f"monthly_deleted={stats['monthly']['deleted']}, " f"errors={stats['errors']}" ) diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py index 85ea407899..d83e2c2c69 100644 --- a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -1,113 +1,260 @@ """Drift guard between the Beat and PG declarations of the metrics periodics (UN-3796). -Two migrations declare the same three schedules — ``0002_setup_periodic_tasks`` for Celery -Beat and ``0004_pg_periodic_tasks`` for the PG scheduler. They are separate rows in -separate tables, so nothing stops someone editing one and forgetting the other. That is -the whole failure mode this file exists for: a schedule changed on Beat but not on PG means -the task silently runs on a different cadence the moment the flag flips. - -DB-free — both migration modules are imported and their declared specs compared directly, -so this runs in the unit tier rather than needing a migrated database. +Every schedule in this app is declared twice — once in +``django_celery_beat_periodictask`` for Celery Beat, once in ``pg_periodic_task`` for the +PG scheduler. They are separate rows in separate tables, so nothing stops someone editing +one and forgetting the other. That is the whole failure mode this file exists for: a +schedule changed on Beat but not on PG means the task silently runs on a different cadence +— or not at all — the moment the flag flips. + +**Every data migration in the app is replayed**, not a named pair. Naming modules is how +the guard went stale before: a schedule added in a later migration kept comparing the +original three against three and stayed green while the invariant it names was violated. +Migrations are run in order against fake models, so rows a later migration rewrites are +compared in their final state. + +DB-free — nothing here touches a database. """ from __future__ import annotations import importlib +import inspect import json +import re +from pathlib import Path +from types import SimpleNamespace import pytest +from django.db import migrations -_BEAT_MIGRATION = "dashboard_metrics.migrations.0002_setup_periodic_tasks" -_PG_MIGRATION = "dashboard_metrics.migrations.0004_pg_periodic_tasks" +from dashboard_metrics.tasks import ( + aggregate_metrics_from_sources, + cleanup_daily_metrics, + cleanup_hourly_metrics, +) -# Cron equivalent of each Beat schedule, asserted against what the Beat migration builds. -# Written out rather than derived: deriving it from the same code under test would make -# the comparison vacuous. +_MIGRATIONS_PKG = "dashboard_metrics.migrations" +_MIGRATIONS_DIR = Path(__file__).resolve().parent.parent / "migrations" + +# Cron equivalent of each schedule, written out rather than derived: an anchor that a +# reviewer reads, and that an edit to both declarations at once still has to touch. _EXPECTED_CRON = { "dashboard_metrics_aggregate_from_sources": "*/15 * * * *", "dashboard_metrics_cleanup_hourly": "0 2 * * *", "dashboard_metrics_cleanup_daily": "0 3 * * 0", + "dashboard_metrics_reconcile_source_window": "40 4 * * *", } -@pytest.fixture(scope="module") -def pg_specs() -> dict[str, dict]: - mod = importlib.import_module(_PG_MIGRATION) - return {spec["name"]: spec for spec in mod.PG_PERIODIC_TASKS} +class _Schedule: + """Stands in for an Interval/CrontabSchedule row, carrying its own cron string.""" + def __init__(self, **kwargs): + self.kwargs = kwargs -class _FakeQuerySet: - """Captures update_or_create calls from the Beat migration without a database.""" + @property + def cron_string(self) -> str: + k = self.kwargs + if "period" in k: + every, period = k["every"], k["period"] + if period == "minutes": + return f"*/{every} * * * *" + if period == "hours": + return f"0 */{every} * * *" + raise AssertionError(f"unhandled interval period: {period}") + return " ".join( + str(k[f]) + for f in ("minute", "hour", "day_of_month", "month_of_year", "day_of_week") + ) - def __init__(self, sink: dict): - self._sink = sink + +class _Rows: + """Captures a migration's writes to one model without a database.""" + + def __init__(self, factory=None): + self.rows: dict[str, dict] = {} + self.writes = 0 + self._factory = factory + self._selected: list[str] = [] def get_or_create(self, **kwargs): - # Schedule rows (Interval/Crontab) — return the kwargs so the PeriodicTask - # call can be inspected for which schedule it was given. - return kwargs, True + kwargs.pop("defaults", None) + return (self._factory(**kwargs) if self._factory else kwargs), True - def update_or_create(self, name=None, defaults=None, **_kw): - self._sink[name] = defaults or {} - return defaults, True + def update_or_create(self, name=None, defaults=None, **kwargs): + self.writes += 1 + if name is None: # e.g. PeriodicTasks(ident=1) — not a schedule row + return defaults, True + self.rows.setdefault(name, {}).update(defaults or {}) + return self.rows[name], True - def filter(self, *_a, **_k): + def filter(self, name=None, name__in=None, **_kwargs): + self._selected = [name] if name is not None else list(name__in or []) return self + def update(self, **kwargs): + self.writes += 1 + for name in self._selected: + self.rows.setdefault(name, {}).update(kwargs) + return len(self._selected) + def delete(self): + self.writes += 1 + for name in self._selected: + self.rows.pop(name, None) return (0, {}) +class _Apps: + def __init__(self): + self.beat = _Rows() + self.pg = _Rows() + self.schedules = _Rows(factory=_Schedule) + self.tracker = _Rows() + self.other = _Rows() + + def get_model(self, app_label, model_name): + target = { + ("django_celery_beat", "PeriodicTask"): self.beat, + ("pg_queue", "PgPeriodicTask"): self.pg, + ("django_celery_beat", "CrontabSchedule"): self.schedules, + ("django_celery_beat", "IntervalSchedule"): self.schedules, + ("django_celery_beat", "PeriodicTasks"): self.tracker, + }.get((app_label, model_name), self.other) + return type("_M", (), {"objects": target}) + + +def _migration_modules() -> list[str]: + names = sorted( + p.stem for p in _MIGRATIONS_DIR.glob("*.py") if re.match(r"^\d{4}_", p.stem) + ) + assert names, "no migrations discovered — the glob is wrong, not the app" + return [f"{_MIGRATIONS_PKG}.{name}" for name in names] + + @pytest.fixture(scope="module") -def beat_specs() -> dict[str, dict]: - """Run the Beat migration's forward function against fakes and capture what it declares.""" - mod = importlib.import_module(_BEAT_MIGRATION) - captured: dict[str, dict] = {} +def declared() -> SimpleNamespace: + """Replay every data migration in order and capture what it declares.""" + apps = _Apps() + for dotted in _migration_modules(): + for op in importlib.import_module(dotted).Migration.operations: + if isinstance(op, migrations.RunPython): + op.code(apps, None) + return SimpleNamespace(beat=apps.beat.rows, pg=apps.pg.rows) - class _Apps: - def get_model(self, _app, model): - if model == "PeriodicTask": - return type("PT", (), {"objects": _FakeQuerySet(captured)}) - return type("S", (), {"objects": _FakeQuerySet({})}) - mod.create_periodic_tasks(_Apps(), None) - return captured +def _beat_cron(row: dict) -> str: + schedule = row.get("crontab") or row.get("interval") + assert schedule is not None, "Beat row declares neither a crontab nor an interval" + return schedule.cron_string class TestDeclarationsAgree: - def test_same_set_of_schedules(self, beat_specs, pg_specs): + def test_same_set_of_schedules(self, declared): # A schedule added to Beat but not PG stops firing the moment the flag flips; # the reverse fires something Beat never knew about. - assert set(beat_specs) == set(pg_specs) + assert set(declared.beat) == set(declared.pg) + + def test_every_known_schedule_is_declared(self, declared): + # Guards the guard: a replay that silently captured nothing would pass the + # set comparison above with two empty sets. + assert set(declared.beat) == set(_EXPECTED_CRON) - @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON)) - def test_task_path_and_queue_match(self, beat_specs, pg_specs, name): - assert pg_specs[name]["task_name"] == beat_specs[name]["task"] - assert pg_specs[name]["queue"] == beat_specs[name]["queue"] + def test_task_path_and_queue_match(self, declared): + for name, beat in declared.beat.items(): + assert declared.pg[name]["task_name"] == beat["task"], name + assert declared.pg[name]["queue"] == beat["queue"], name - @pytest.mark.parametrize("name", sorted(_EXPECTED_CRON)) - def test_kwargs_match_once_decoded(self, beat_specs, pg_specs, name): + def test_kwargs_match_once_decoded(self, declared): # Beat stores kwargs as a JSON *string*; PgPeriodicTask.task_kwargs is a - # JSONField. A mismatch here means the cleanup runs with the wrong retention. - beat_kwargs = json.loads(beat_specs[name].get("kwargs") or "{}") - assert pg_specs[name]["task_kwargs"] == beat_kwargs + # JSONField. A mismatch means the task runs with the wrong arguments. + for name, beat in declared.beat.items(): + assert declared.pg[name]["task_kwargs"] == json.loads( + beat.get("kwargs") or "{}" + ), name + + def test_cadence_matches_across_transports(self, declared): + # Derived from the Beat schedule row rather than from a table, so a cadence + # changed on one transport only fails here whatever its name. + for name, beat in declared.beat.items(): + assert declared.pg[name]["cron_string"] == _beat_cron(beat), name - @pytest.mark.parametrize("name,cron", sorted(_EXPECTED_CRON.items())) - def test_cron_matches_the_beat_cadence(self, pg_specs, name, cron): - assert pg_specs[name]["cron_string"] == cron + def test_cadence_matches_the_written_anchor(self, declared): + for name, cron in _EXPECTED_CRON.items(): + assert declared.pg[name]["cron_string"] == cron + + +# 0002 seeds at install time, when Beat has never started and has nothing stale to +# reload. Every migration after it rewrites a schedule a running Beat already holds. +_INSTALL_MIGRATION = "0002_setup_periodic_tasks" + + +class TestRunningBeatIsToldToReload: + """Historical models fire no post_save, so DatabaseScheduler never reloads. + + Without an explicit ``PeriodicTasks.last_update`` bump a live Beat keeps firing its + in-memory copy: rows this migration adds never fire, rows it rewrites keep their old + arguments. Nothing errors, and the whole change silently does not happen. + """ + + def test_every_post_install_beat_write_bumps_the_change_tracker(self): + checked = 0 + for dotted in _migration_modules(): + if dotted.endswith(_INSTALL_MIGRATION): + continue + for op in importlib.import_module(dotted).Migration.operations: + if not isinstance(op, migrations.RunPython): + continue + for direction in (op.code, op.reverse_code): + if direction is None: + continue + apps = _Apps() + direction(apps, None) + if not apps.beat.writes: + continue + checked += 1 + assert apps.tracker.writes, f"{dotted}.{direction.__name__}" + assert checked, "no post-install Beat writes found — the discovery is broken" + + +class TestDeclaredKwargsAreCallable: + """A schedule row carrying a kwarg its task cannot bind raises TypeError per tick. + + TypeError is not in ``autoretry_for``, and the PG leg drops the message at + MAX_ATTEMPTS=1 — so the schedule silently never runs. Enumerating every declared + row rather than one migration's own spec is the point: the rows are added by + different migrations, and each new one is exactly the case that escapes a guard + scoped to a single module. + """ + + _TASKS = { + task.name: task + for task in ( + aggregate_metrics_from_sources, + cleanup_hourly_metrics, + cleanup_daily_metrics, + ) + } + + def test_every_declared_kwarg_set_binds_to_the_task_signature(self, declared): + for name, row in declared.pg.items(): + task = self._TASKS.get(row["task_name"]) + assert task is not None, f"{name} schedules an unknown task" + inspect.signature(task).bind(**row["task_kwargs"]) class TestSeededInert: - """Applying the migration must not cause anything to fire.""" - - def test_no_spec_declares_itself_pg_owned(self, pg_specs): - # pg_owned is set to False in the migration's defaults, never from the spec — - # this pins that no spec can smuggle ownership in. - assert not any("pg_owned" in spec for spec in pg_specs.values()) - - def test_no_spec_presets_a_run_time(self, pg_specs): - # A non-NULL next_run_at in the past would read as "overdue" and fire a burst - # of catch-up runs the moment the flag is enabled. - for spec in pg_specs.values(): - assert "next_run_at" not in spec - assert "last_run_at" not in spec + """Applying the migrations must not cause anything to fire.""" + + def test_nothing_is_declared_pg_owned(self, declared): + # pg_owned=True would hand the row to the PG scheduler before the rollout + # flag decides, and disable its Beat twin. + assert not any(row.get("pg_owned") for row in declared.pg.values()) + + def test_no_row_presets_a_run_time(self, declared): + # A non-NULL next_run_at in the past reads as "overdue" and fires a burst of + # catch-up runs the moment the flag is enabled. + for row in declared.pg.values(): + assert "next_run_at" not in row + assert "last_run_at" not in row diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index 822c197a63..e85acd203b 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -3,14 +3,16 @@ import json from datetime import date, datetime, timedelta from importlib import import_module +from types import SimpleNamespace from unittest.mock import patch from django.apps import apps from django.db import connection +from django.db.utils import DatabaseError from django.test import TestCase from django.test.utils import CaptureQueriesContext from django.utils import timezone -from django_celery_beat.models import PeriodicTask +from django_celery_beat.models import PeriodicTask, PeriodicTasks from account_v2.models import Organization from dashboard_metrics.models import ( @@ -19,13 +21,16 @@ EventMetricsMonthly, MetricType, ) +from pg_queue.models import PgPeriodicTask from workflow_manager.file_execution.models import WorkflowFileExecution from workflow_manager.workflow_v2.enums import ExecutionStatus from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow +from dashboard_metrics.internal_views import AggregateMetricsAPIView from dashboard_metrics.tasks import ( DASHBOARD_RECONCILE_WINDOW_DAYS, DASHBOARD_SOURCE_WINDOW_DAYS, + _active_org_ids, _rollup_monthly_from_daily, _run_aggregation, _truncate_to_day, @@ -226,12 +231,20 @@ def setUp(self): organization_id="rollup-org", name="rollup-org", display_name="Rollup Org" ) - def _daily(self, day, value, count=1, metric_type=MetricType.COUNTER): - """Create a daily metric row for the fixture org.""" + def _daily( + self, + day, + value, + count=1, + metric_type=MetricType.COUNTER, + metric_name="documents_processed", + org=None, + ): + """Create a daily metric row, defaulting to the fixture org and metric.""" EventMetricsDaily.objects.create( - organization=self.org, + organization=org or self.org, date=day, - metric_name="documents_processed", + metric_name=metric_name, metric_type=metric_type, metric_value=value, metric_count=count, @@ -239,15 +252,19 @@ def _daily(self, day, value, count=1, metric_type=MetricType.COUNTER): ) def _monthly_rows(self): - """Read back monthly rows ordered by month.""" - return list(EventMetricsMonthly._base_manager.order_by("month")) + """Read back monthly rows in a stable order.""" + return list( + EventMetricsMonthly._base_manager.order_by( + "month", "organization_id", "metric_name" + ) + ) def test_sums_daily_rows_into_month_bucket(self): """Daily rows within a month sum into a single monthly row.""" self._daily(date(2024, 3, 5), value=10, count=2) self._daily(date(2024, 3, 18), value=32, count=4) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) rows = self._monthly_rows() assert len(rows) == 1 @@ -262,7 +279,7 @@ def test_month_boundary_keeps_months_separate(self): self._daily(date(2024, 2, 1), value=100) self._daily(date(2024, 2, 2), value=200) - assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 2 + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == (2, 0) rows = self._monthly_rows() assert [r.month for r in rows] == [date(2024, 1, 1), date(2024, 2, 1)] @@ -273,7 +290,7 @@ def test_excludes_months_before_the_window(self): self._daily(date(2023, 12, 15), value=999) self._daily(date(2024, 1, 15), value=5) - assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 1 + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == (1, 0) rows = self._monthly_rows() assert len(rows) == 1 @@ -297,30 +314,104 @@ def test_mixed_metric_type_within_a_month_yields_one_row(self): self._daily(date(2024, 3, 5), value=10, metric_type=MetricType.HISTOGRAM) self._daily(date(2024, 3, 6), value=5, metric_type=MetricType.COUNTER) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) rows = self._monthly_rows() assert len(rows) == 1 assert rows[0].metric_value == 15 - def test_no_daily_rows_upserts_nothing(self): - """An empty daily tier is a no-op, not an error.""" - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 0 - assert not self._monthly_rows() + def test_an_empty_daily_tier_leaves_existing_monthly_rows_alone(self): + """An empty tier means the source is gone, not that every month is zero. - def test_monthly_row_is_dropped_once_its_daily_rows_are_gone(self): - """A month whose daily rows were deleted must not keep a stale total.""" + Seeding a monthly row first is what makes the failure reachable at all: with + an empty table a sweep that deletes everything and one that deletes nothing + both leave an empty table, and the assertion passes either way. + """ + EventMetricsMonthly._base_manager.create( + organization=self.org, + month=date(2024, 3, 1), + metric_name="documents_processed", + metric_type=MetricType.COUNTER, + metric_value=42, + metric_count=6, + project="default", + ) + + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (0, 0) + + rows = self._monthly_rows() + assert len(rows) == 1 + assert rows[0].metric_value == 42 + + def test_metric_is_dropped_once_its_daily_rows_are_gone(self): + """A metric with no daily rows left must not keep a stale monthly total.""" + self._daily(date(2024, 3, 5), value=10) + self._daily(date(2024, 3, 6), value=7, metric_name="pages_processed") + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) + + EventMetricsDaily._base_manager.filter(metric_name="pages_processed").delete() + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 1) + + rows = self._monthly_rows() + assert [r.metric_name for r in rows] == ["documents_processed"] + + def test_a_month_with_no_daily_rows_of_its_own_is_left_alone(self): + """An incomplete daily tier must read as "unknown", not as "deleted". + + A partially populated tier passes the empty-tier guard, so without scoping the + sweep to the (organization, month) partitions the rollup actually produced, + one missing month wipes that month's monthly rows for every org. + """ self._daily(date(2024, 3, 5), value=10) self._daily(date(2024, 4, 5), value=7) - _rollup_monthly_from_daily(date(2024, 3, 1)) - assert len(self._monthly_rows()) == 2 + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) EventMetricsDaily._base_manager.filter(date=date(2024, 3, 5)).delete() - _rollup_monthly_from_daily(date(2024, 3, 1)) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) + + months = [row.month for row in self._monthly_rows()] + assert months == [date(2024, 3, 1), date(2024, 4, 1)] + + def test_the_sweep_is_scoped_per_organization(self): + """The sweep bypasses the org-scoped manager, so the org half of the key is + load-bearing: drop it and one org's daily rows vouch for another's monthly.""" + other = Organization.objects.create( + organization_id="rollup-org-2", name="rollup-org-2", display_name="Other" + ) + self._daily(date(2024, 3, 5), value=10) + self._daily(date(2024, 3, 6), value=20, org=other) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) + + EventMetricsDaily._base_manager.filter(organization=other).delete() + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) rows = self._monthly_rows() - assert len(rows) == 1 - assert rows[0].month == date(2024, 4, 1) + assert [(r.organization_id, r.metric_value) for r in rows] == [ + (self.org.id, 10), + (other.id, 20), + ] + + def test_a_stale_row_is_dropped_for_one_organization_only(self): + """Two orgs in one month: only the org whose metric vanished loses its row.""" + other = Organization.objects.create( + organization_id="rollup-org-3", name="rollup-org-3", display_name="Other" + ) + self._daily(date(2024, 3, 5), value=10) + self._daily(date(2024, 3, 5), value=1, metric_name="pages_processed") + self._daily(date(2024, 3, 6), value=20, org=other) + self._daily(date(2024, 3, 6), value=2, metric_name="pages_processed", org=other) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (4, 0) + + EventMetricsDaily._base_manager.filter( + organization=other, metric_name="pages_processed" + ).delete() + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (3, 1) + + survivors = { + (r.organization_id, r.metric_name) for r in self._monthly_rows() + } + assert (self.org.id, "pages_processed") in survivors + assert (other.id, "pages_processed") not in survivors def test_months_before_the_window_are_left_alone(self): """Orphan cleanup must not reach outside the rebuilt window.""" @@ -366,6 +457,181 @@ def test_monthly_rollup_never_touches_source_tables(self): assert source_table not in sql, f"monthly rollup read {source_table}" +class TestActiveOrgPrefilter(TestCase): + """The prefilter must never be narrower than the window it is filtering for.""" + + def setUp(self): + self.org = Organization.objects.create( + organization_id="prefilter-org", name="prefilter", display_name="Prefilter" + ) + workflow = Workflow.objects.create( + workflow_name="prefilter-wf", organization=self.org + ) + self.now = timezone.now() + execution = WorkflowExecution.objects.create( + workflow_id=workflow.id, status=ExecutionStatus.COMPLETED + ) + WorkflowExecution.objects.filter(pk=execution.pk).update( + created_at=self.now - timedelta(days=10) + ) + + def test_an_org_outside_the_default_lookback_is_filtered_out(self): + """The default lookback is the cheap case and stays exactly as wide as before.""" + window_start = self.now - timedelta(days=DASHBOARD_SOURCE_WINDOW_DAYS) + assert self.org.id not in _active_org_ids(self.now, window_start) + + def test_a_widened_window_widens_the_prefilter_with_it(self): + """Otherwise a long-outage repair queries 30 days for orgs active in 7, and + reports errors: 0 having skipped every org it exists to repair.""" + window_start = self.now - timedelta(days=30) + assert self.org.id in _active_org_ids(self.now, window_start) + + +class TestMonthlyRollupFailurePosture(TestCase): + """The rollup's errors must reach the task's retry, not a stats counter.""" + + def test_a_database_error_propagates_instead_of_reporting_success(self): + """DatabaseError/OperationalError are what autoretry_for is configured for. + + Swallowed here they become one INFO line and success: True, and a persistent + fault leaves monthly permanently stale behind 96 clean-looking runs a day. + """ + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [1] + with patch( + "dashboard_metrics.tasks._rollup_monthly_from_daily", + side_effect=DatabaseError("lock timeout"), + ): + with self.assertRaises(DatabaseError): + _run_aggregation() + + def test_an_unexpected_error_is_still_counted_rather_than_fatal(self): + """Everything outside the retry set keeps the previous non-fatal posture.""" + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [1] + with patch( + "dashboard_metrics.tasks._rollup_monthly_from_daily", + side_effect=ValueError("bad row"), + ): + result = _run_aggregation() + assert result["success"] is True + assert result["errors"] == 1 + + +class TestInternalAggregateEndpoint(TestCase): + """The PG transport reaches the task through this view, not through Celery.""" + + def _post(self, data): + return AggregateMetricsAPIView().post(SimpleNamespace(data=data)) + + def test_the_source_window_reaches_the_task(self): + with patch( + "dashboard_metrics.internal_views.aggregate_metrics_from_sources", + return_value={"success": True}, + ) as task: + self._post({"source_window_days": 7}) + assert task.call_args.kwargs == {"source_window_days": 7} + + def test_omitting_it_leaves_the_task_default_in_charge(self): + with patch( + "dashboard_metrics.internal_views.aggregate_metrics_from_sources", + return_value={"success": True}, + ) as task: + self._post({}) + assert task.call_args.kwargs == {} + + def test_a_non_integer_window_is_a_400(self): + with patch( + "dashboard_metrics.internal_views.aggregate_metrics_from_sources" + ) as task: + response = self._post({"source_window_days": "seven"}) + assert response.status_code == 400 + task.assert_not_called() + + +class TestMonthlyThroughTheTask(TestCase): + """The rollup as the task actually runs it, not via the helper directly. + + Every other rollup test calls ``_rollup_monthly_from_daily`` with a hand-chosen + ``month_start``. Nothing exercised the arithmetic that computes it, nor the sweep + running against a monthly table that already holds rows from earlier runs — so a + regression to "first of the current month" would silently drop last month's rows + with the whole rollup suite still green. + """ + + def setUp(self): + self.org = Organization.objects.create( + organization_id="entry-org", name="entry-org", display_name="Entry Org" + ) + now = timezone.now() + self.this_month = _truncate_to_month(now).date() + self.last_month = _truncate_to_month( + _truncate_to_month(now) - timedelta(days=1) + ).date() + self.before_window = _truncate_to_month( + _truncate_to_month(now - timedelta(days=1)) - timedelta(days=40) + ).date() + + def _daily(self, day, value, metric_name="documents_processed"): + EventMetricsDaily._base_manager.create( + organization=self.org, + date=day, + metric_name=metric_name, + metric_type=MetricType.COUNTER, + metric_value=value, + metric_count=1, + project="default", + tag="", + ) + + def _monthly(self, month, value, metric_name="documents_processed"): + EventMetricsMonthly._base_manager.create( + organization=self.org, + month=month, + metric_name=metric_name, + metric_type=MetricType.COUNTER, + metric_value=value, + metric_count=1, + project="default", + tag="", + ) + + def _run(self, **kwargs): + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [self.org.id] + return _run_aggregation(**kwargs) + + def test_the_window_covers_the_previous_month_and_spares_what_precedes_it(self): + """monthly_start is the first of the *previous* month, and the sweep stops there.""" + self._daily(self.this_month, value=10) + self._daily(self.last_month, value=20) + self._monthly(self.before_window, value=999) + + result = self._run() + + assert result["period"]["monthly"]["start"] == self.last_month.isoformat() + assert result["monthly"] == {"upserted": 2, "deleted": 0} + + rows = EventMetricsMonthly._base_manager.order_by("month") + assert [r.month for r in rows] == [ + self.before_window, + self.last_month, + self.this_month, + ] + + def test_the_run_reports_what_it_deleted(self): + """The only destructive write in the task has to reach the result.""" + self._daily(self.this_month, value=10) + self._daily(self.this_month, value=1, metric_name="pages_processed") + assert self._run()["monthly"] == {"upserted": 2, "deleted": 0} + + EventMetricsDaily._base_manager.filter(metric_name="pages_processed").delete() + assert self._run()["monthly"] == {"upserted": 1, "deleted": 1} + + class TestSourceWindow(TestCase): """Tests for the per-run source window and the reconciliation pass.""" @@ -403,12 +669,25 @@ def test_reconciliation_window_widens_the_daily_query(self): assert result["period"]["daily"]["start"] == expected.isoformat() def test_task_passes_the_window_through(self): - """The scheduled task forwards its kwarg, defaulting to the per-run window.""" - with patch("dashboard_metrics.tasks._run_aggregation") as mock_run: + """The scheduled task forwards its kwarg, defaulting to the per-run window. + + The lock is patched out: acquiring it for real takes — and then releases in the + task's ``finally`` — the shared Redis key a live local aggregation may be + holding. + """ + with ( + patch("dashboard_metrics.tasks._acquire_aggregation_lock", return_value=True), + patch("dashboard_metrics.tasks.cache"), + patch("dashboard_metrics.tasks._run_aggregation") as mock_run, + ): aggregate_metrics_from_sources() mock_run.assert_called_once_with(DASHBOARD_SOURCE_WINDOW_DAYS) - with patch("dashboard_metrics.tasks._run_aggregation") as mock_run: + with ( + patch("dashboard_metrics.tasks._acquire_aggregation_lock", return_value=True), + patch("dashboard_metrics.tasks.cache"), + patch("dashboard_metrics.tasks._run_aggregation") as mock_run, + ): aggregate_metrics_from_sources(source_window_days=7) mock_run.assert_called_once_with(7) @@ -486,7 +765,7 @@ def test_gap_older_than_the_reconcile_window_needs_a_manual_backfill(self): class TestReconciliationSchedule(TestCase): - """Migration 0004 schedules the once-daily reconciliation pass. + """Migration 0005 schedules the once-daily reconciliation pass on both transports. The suite runs with --no-migrations, so the migration's function is called directly rather than relying on it having been applied. @@ -501,8 +780,8 @@ def setUp(self): def _task(self): return PeriodicTask.objects.get(name=self.migration.RECONCILE_TASK_NAME) - def test_migration_schedules_the_pass_at_0400_with_a_7_day_window(self): - """The beat row lands enabled, at 04:00 UTC, carrying the wider window.""" + def test_migration_schedules_the_pass_at_0440_with_a_7_day_window(self): + """The beat row lands enabled, at 04:40 UTC, carrying the wider window.""" self.migration.create_reconciliation_task(apps, None) task = self._task() @@ -512,7 +791,35 @@ def test_migration_schedules_the_pass_at_0400_with_a_7_day_window(self): assert json.loads(task.kwargs) == { "source_window_days": DASHBOARD_RECONCILE_WINDOW_DAYS } - assert (task.crontab.hour, task.crontab.minute) == ("4", "0") + assert (task.crontab.hour, task.crontab.minute) == ("4", "40") + + def test_the_pg_twin_lands_with_the_same_cadence_and_kwargs(self): + """A Beat-only row stops firing the moment the PG scheduler takes over.""" + self.migration.create_reconciliation_task(apps, None) + + row = PgPeriodicTask.objects.get(name=self.migration.RECONCILE_TASK_NAME) + assert row.task_name == "dashboard_metrics.aggregate_from_sources" + assert row.queue == "dashboard_metric_events" + assert row.cron_string == "40 4 * * *" + assert row.task_kwargs == { + "source_window_days": DASHBOARD_RECONCILE_WINDOW_DAYS + } + assert row.enabled + # Inert until the rollout flag decides otherwise. + assert not row.pg_owned + assert row.next_run_at is None + + def test_a_running_beat_is_told_to_reload(self): + """Historical models fire no post_save, so the tracker has to be bumped by hand. + + Without it a live Beat never adopts the new schedule and the reconciliation + pass simply never runs — no error, nothing logged. + """ + before = timezone.now() + self.migration.create_reconciliation_task(apps, None) + + tracker = PeriodicTasks.objects.get(ident=1) + assert tracker.last_update >= before def test_migration_is_idempotent_and_reversible(self): """Re-running leaves one row; the reverse function removes it.""" @@ -530,3 +837,6 @@ def test_migration_is_idempotent_and_reversible(self): assert not PeriodicTask.objects.filter( name=self.migration.RECONCILE_TASK_NAME ).exists() + assert not PgPeriodicTask.objects.filter( + name=self.migration.RECONCILE_TASK_NAME + ).exists() diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 44bbe50440..07787aa9f3 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -105,9 +105,20 @@ def _log_if_skipped(name: str, result: dict[str, Any]) -> None: @worker_task(name="dashboard_metrics.aggregate_from_sources") -def dashboard_metrics_aggregate() -> dict[str, Any]: - """Aggregate source tables into the hourly/daily/monthly metrics tables.""" - result = _call_internal(_AGGREGATE_PATH) +def dashboard_metrics_aggregate( + source_window_days: int | None = None, +) -> dict[str, Any]: + """Aggregate source tables into the hourly/daily/monthly metrics tables. + + The daily reconciliation schedule dispatches a wider ``source_window_days``; + omitting it applies the backend task's own default. + """ + body = ( + {"source_window_days": source_window_days} + if source_window_days is not None + else None + ) + result = _call_internal(_AGGREGATE_PATH, body=body) _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) return result diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py index ce8ff853ac..b78cbedb01 100644 --- a/workers/tests/test_dashboard_metrics_tasks.py +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -8,6 +8,7 @@ from __future__ import annotations +import inspect import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -61,12 +62,38 @@ def test_task_is_registered_under_the_beat_name(self, name, func): assert getattr(dmt, func).name == name +# The kwargs the dashboard_metrics migrations declare on schedule rows for +# dashboard_metrics.aggregate_from_sources. Beat dispatches straight to the Django +# task, the PG scheduler dispatches to the proxy below — so a kwarg the proxy cannot +# bind raises TypeError per tick, is not covered by autoretry_for, and is dropped at +# MAX_ATTEMPTS=1. Kept in step by dashboard_metrics/tests/test_pg_periodic_task_declarations.py +# on the Django side; this is the half that lives outside Django. +_DECLARED_AGGREGATE_KWARGS = [{}, {"source_window_days": 7}] + + class TestCallContract: def test_aggregate_posts_to_the_aggregate_endpoint(self): with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: dmt.dashboard_metrics_aggregate() assert call.call_args[0][0] == "v1/dashboard-metrics/aggregate/" + @pytest.mark.parametrize("kwargs", _DECLARED_AGGREGATE_KWARGS) + def test_every_scheduled_kwarg_set_binds_to_the_proxy(self, kwargs): + inspect.signature(dmt.dashboard_metrics_aggregate).bind(**kwargs) + + def test_aggregate_passes_the_source_window_through(self): + # The 4 AM reconciliation row carries this; dropping it here silently reverts + # the pass to the narrow 15-minute window it exists to widen. + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate(source_window_days=7) + assert call.call_args.kwargs["body"] == {"source_window_days": 7} + + def test_aggregate_omits_body_when_no_window_given(self): + # The backend then applies the task's own default rather than one invented here. + with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_args.kwargs["body"] is None + @pytest.mark.parametrize( "func,path", [ From 327c883afbbd3b137bd61d732804bc29ecef766d Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 2 Sep 2026 16:49:17 +0530 Subject: [PATCH 12/14] UN-3974 [FIX] Address review: Beat reload, inherited ownership, boundary validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Rewriting live Beat rows through historical models fires no post_save, so DatabaseScheduler never reloaded: the existing row kept firing with no tier and the new row never fired at all. 0006 now bumps PeriodicTasks.last_update in both directions, as scheduler/ownership.py and mirror_pg_periodic_tasks.py already do. The new row inherits pg_owned and both enabled flags from the row it is split from instead of hardcoding Beat. In a PG-adopted environment the daily and monthly tiers had no firer at all while the hourly run still returned success. It also moves to minute 20. Minute 0 collides with */15 — and so does the suggested minute 30, since */15 fires at :00 :15 :30 :45 — and the per-tier locks are built so the two runs cannot block each other. An unrecognised tier is now rejected in post(), and the blanket except ValueError in _run is gone, so a ValueError from inside the aggregation reaches the logged 500 path rather than reading as a bad request body. Tests: the JSON round-trip assertion was a stdlib tautology that never read what the migration writes — replaced with an assertion on the updated row, the one firing the hourly tier in production. The ALL default is pinned off inspect.signature. The RunSQL table and column are derived from the model rather than grepped. The planner-choice assertion is deleted: a cost model on 12,000 rows is not production evidence. Also drops a full Organization count that ran on every tier for one log field, and gives two test modules the Django bootstrap their siblings carry. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/README.md | 10 +- backend/dashboard_metrics/internal_views.py | 21 +++- .../0006_split_aggregation_schedule.py | 62 ++++++++- backend/dashboard_metrics/tasks.py | 10 +- .../tests/test_active_org_prefilter.py | 38 +++--- .../tests/test_aggregation_dispatch.py | 27 ++-- .../tests/test_aggregation_tier.py | 22 +++- .../test_pg_periodic_task_declarations.py | 119 ++++++++++++++++-- .../tests/test_tier_split_equivalence.py | 23 ++-- .../tests/test_we_created_at_idx.py | 15 +++ 10 files changed, 274 insertions(+), 73 deletions(-) diff --git a/backend/dashboard_metrics/README.md b/backend/dashboard_metrics/README.md index fc40e9f55b..aa1d39d2e3 100644 --- a/backend/dashboard_metrics/README.md +++ b/backend/dashboard_metrics/README.md @@ -14,7 +14,7 @@ This module provides a metrics dashboard for monitoring document processing, API ### Data Flow ``` Source Tables (usage_v2, page_usage, workflow_execution, workflow_file_execution) - ↓ [Celery task every 15 min] + ↓ [Celery: hourly tier every 15 min, daily+monthly hourly at :20] Aggregated Tables (EventMetricsHourly → Daily → Monthly) ↓ API Endpoints (/overview/, /summary/, /series/) @@ -46,7 +46,8 @@ celery -A backend beat -l info ### Celery Tasks & Schedule | Task | Schedule | What It Does | |------|----------|--------------| -| `aggregate_from_sources` | Every 15 min | Aggregates source → hourly/daily/monthly | +| `aggregate_from_sources` | Every 15 min | Aggregates source → **hourly tier only** (`tier=hourly`) | +| `aggregate_daily_monthly` | Hourly at :20 | Aggregates source → **daily + monthly tiers** (`tier=daily_monthly`) | | `cleanup_hourly_data` | Daily 2 AM | Deletes hourly data > 30 days | | `cleanup_daily_data` | Weekly Sun 3 AM | Deletes daily data > 365 days | @@ -162,7 +163,7 @@ The dashboard reads from **pre-aggregated tables** (`event_metrics_hourly`, `eve - Source table performance is unaffected by the dashboard feature. If the aggregation task is slow or fails, source tables continue working normally. **Failure Resilience:** -- If the aggregation task fails, the dashboard shows stale data (up to 15 minutes old) rather than crashing. +- If the aggregation task fails, the dashboard shows stale data rather than crashing — up to 15 minutes old for hourly figures, up to an hour for daily and monthly. - Celery tasks have `max_retries=3` with exponential backoff. - Cleanup tasks (hourly: 30-day retention, daily: 365-day retention) prevent unbounded table growth. @@ -339,7 +340,8 @@ Located in `tasks.py`: | Task Name | Celery Name | Schedule | Queue | Purpose | |-----------|-------------|----------|-------|---------| -| `aggregate_metrics_from_sources` | `dashboard_metrics.aggregate_from_sources` | Every 15 min | `dashboard_metric_events` | Aggregate from source tables | +| `aggregate_metrics_from_sources` | `dashboard_metrics.aggregate_from_sources` | Every 15 min | `dashboard_metric_events` | Aggregate the hourly tier (`tier=hourly`) | +| `aggregate_metrics_from_sources` | `dashboard_metrics.aggregate_from_sources` | Hourly at :20 UTC | `dashboard_metric_events` | Aggregate the daily and monthly tiers (`tier=daily_monthly`) | | `cleanup_hourly_metrics` | `dashboard_metrics.cleanup_hourly_data` | Daily 2:00 AM UTC | `dashboard_metric_events` | Delete hourly data >30 days | | `cleanup_daily_metrics` | `dashboard_metrics.cleanup_daily_data` | Weekly Sun 3:00 AM UTC | `dashboard_metric_events` | Delete daily data >365 days | diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index 529e8f34af..64bf048691 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -34,6 +34,7 @@ from utils.local_context import StateStore from dashboard_metrics.tasks import ( + AggregationTier, aggregate_metrics_from_sources, cleanup_daily_metrics, cleanup_hourly_metrics, @@ -74,11 +75,12 @@ class _MetricsTaskAPIView(APIView): """Shared plumbing: clear org context, run, translate errors.""" def _run(self, fn, *args: Any, **kwargs: Any) -> Response: + """Run one task body. Every view validates its own body first, so anything + raising in here is an internal fault and belongs on the logged 500 path. + """ _clear_org_context() try: return Response(fn(*args, **kwargs)) - except ValueError as exc: # bad request body - return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) except Exception as exc: logger.error("dashboard-metrics internal call failed: %s", exc, exc_info=True) return Response( @@ -93,13 +95,24 @@ class AggregateMetricsAPIView(_MetricsTaskAPIView): only because the PG consumer has no Django, not to change what the job does. Optional ``tier`` in the body selects which metric tiers to write; omitting it - writes all of them. An unrecognised value is a 400, not a silent no-op. + writes all of them. An unrecognised value is a 400 raised here at the boundary, so + a ValueError from inside the aggregation stays a logged 500 rather than reading as + a bad request. """ def post(self, request: Request) -> Response: - tier = request.data.get("tier") if isinstance(request.data, dict) else None + body = request.data if isinstance(request.data, dict) else {} + tier = body.get("tier") if tier is None: return self._run(aggregate_metrics_from_sources) + try: + tier = AggregationTier(tier) + except ValueError: + valid = [member.value for member in AggregationTier] + return Response( + {"error": f"tier must be one of {valid}, got {tier!r}"}, + status=status.HTTP_400_BAD_REQUEST, + ) return self._run(aggregate_metrics_from_sources, tier=tier) diff --git a/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py index a9cbf1aabf..9b60398d19 100644 --- a/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py +++ b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py @@ -7,12 +7,23 @@ Beat and PG rows are declared here from one spec, so this pair cannot drift the way 0002 and 0004 can. Beat stores kwargs as a JSON string, PgPeriodicTask decoded. The -new PG row lands inert (``pg_owned=False``) like 0004's. +new row inherits whichever scheduler owns the row it is split from, rather than +hardcoding Beat: it is one half of that row, and the same process should fire it. + +The new row runs at minute 20 — off the ``*/15`` grid — so it never starts alongside +the hourly-tier run, whose per-tier lock is deliberately unable to block it. + +**Rolling back the code past this release requires reversing this migration too.** +After it runs both scheduler rows carry a ``tier`` kwarg that the previous release's +zero-argument signatures reject with ``TypeError``, which ``autoretry_for`` does not +cover — aggregation would stop for both tiers until ``migrate dashboard_metrics 0005`` +restores the single row. """ import json from django.db import migrations +from django.utils import timezone AGGREGATE_TASK_NAME = "dashboard_metrics.aggregate_from_sources" AGGREGATE_QUEUE = "dashboard_metric_events" @@ -40,8 +51,11 @@ { "name": "dashboard_metrics_aggregate_daily_monthly", "tier": TIER_DAILY_MONTHLY, - "cron_string": "0 * * * *", - "crontab": {"minute": "0", "hour": "*"}, + # Off the */15 grid (:00 :15 :30 :45): the per-tier locks are built so the + # two runs cannot block each other, so a shared minute means two full + # prefilter scans and two per-org loops at once. Same cadence, no overlap. + "cron_string": "20 * * * *", + "crontab": {"minute": "20", "hour": "*"}, "description": ( "Aggregate the daily and monthly dashboard metrics tiers from source " "tables — hourly, since these figures do not need 15-minute freshness" @@ -51,10 +65,27 @@ ] +def _inherited_ownership(periodic_task_model, pg_periodic_task_model): + """Which scheduler fires the row being split, so its other half matches. + + Hardcoding Beat would leave the daily/monthly tier with no firer wherever the + metrics periodics are already PG-adopted: the adopted row's Beat twin is disabled + and Beat may not be running at all. + """ + beat = periodic_task_model.objects.filter(name=EXISTING_AGGREGATE_ROW).first() + pg = pg_periodic_task_model.objects.filter(name=EXISTING_AGGREGATE_ROW).first() + return { + "beat_enabled": True if beat is None else beat.enabled, + "pg_enabled": True if pg is None else pg.enabled, + "pg_owned": False if pg is None else pg.pg_owned, + } + + def split_schedules(apps, schema_editor): CrontabSchedule = apps.get_model("django_celery_beat", "CrontabSchedule") PeriodicTask = apps.get_model("django_celery_beat", "PeriodicTask") PgPeriodicTask = apps.get_model("pg_queue", "PgPeriodicTask") + owner = _inherited_ownership(PeriodicTask, PgPeriodicTask) for spec in AGGREGATION_SCHEDULES: kwargs = {"tier": spec["tier"]} @@ -84,7 +115,7 @@ def split_schedules(apps, schema_editor): "crontab": schedule, "queue": AGGREGATE_QUEUE, "kwargs": json.dumps(kwargs), - "enabled": True, + "enabled": owner["beat_enabled"], "description": spec["description"], }, ) @@ -97,11 +128,13 @@ def split_schedules(apps, schema_editor): "task_kwargs": kwargs, "cron_string": spec["cron_string"], "org_id": "", - "enabled": True, - "pg_owned": False, + "enabled": owner["pg_enabled"], + "pg_owned": owner["pg_owned"], }, ) + _bump_beat_change_tracker(apps) + def merge_schedules(apps, schema_editor): """Restore the single every-15-minutes row that writes all three tiers. @@ -123,6 +156,23 @@ def merge_schedules(apps, schema_editor): ), ) PgPeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update(task_kwargs={}) + _bump_beat_change_tracker(apps) + + +def _bump_beat_change_tracker(apps): + """Make a running Beat reload instead of keeping the pre-split schedule. + + django-celery-beat's post_save receiver binds the concrete PeriodicTask, so writes + through a historical model never bump PeriodicTasks.last_update and + DatabaseScheduler keeps its in-memory copy: the existing row would go on firing + with no tier and the new row would never fire at all — the whole saving silently + not happening. Same fix and reason as scheduler/ownership.py and + mirror_pg_periodic_tasks.py. + """ + periodic_tasks_model = apps.get_model("django_celery_beat", "PeriodicTasks") + periodic_tasks_model.objects.update_or_create( + ident=1, defaults={"last_update": timezone.now()} + ) class Migration(migrations.Migration): diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 8af0cba530..1c8a447ffb 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -628,13 +628,9 @@ def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, A .values_list("workflow__organization_id", flat=True) .distinct() ) - total_orgs = Organization.objects.count() - logger.info( - "Aggregation (%s): %d active orgs out of %d total", - tier.value, - len(active_org_ids), - total_orgs, - ) + # No total_orgs here: a full count of the organization table, on every run of + # every tier, whose only consumer was this log line. + logger.info("Aggregation (%s): %d active orgs", tier.value, len(active_org_ids)) if not active_org_ids: return { diff --git a/backend/dashboard_metrics/tests/test_active_org_prefilter.py b/backend/dashboard_metrics/tests/test_active_org_prefilter.py index a1dbe082ad..6cb6263764 100644 --- a/backend/dashboard_metrics/tests/test_active_org_prefilter.py +++ b/backend/dashboard_metrics/tests/test_active_org_prefilter.py @@ -16,21 +16,36 @@ composite (workflow_id, created_at DESC) index instead, which is an artefact of the fixture rather than anything about the query. +**Not production evidence.** A few thousand rows in an otherwise-empty table on a +locally-configured Postgres is not the production planner's input: index-vs-seq-scan at +this selectivity is a cost-model output, sensitive to the PG major version, +`random_page_cost`, `effective_cache_size` and parallel workers, none of which are +pinned here. What the plan assertion below rules out is the *regression* — a prefilter +that has to read the executions table whatever the costs say. Whether production picks +the index is measured on production, and belongs to AC-3. + DB-bound, so conftest marks it integration. """ from __future__ import annotations +import os + +import django +from django.apps import apps -from account_v2.models import Organization -from django.db import connection -from django.test import TestCase -from django.test.utils import CaptureQueriesContext -from workflow_manager.workflow_v2.models.workflow import Workflow +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() -from dashboard_metrics.tasks import AggregationTier, _run_aggregation +from account_v2.models import Organization # noqa: E402 +from django.db import connection # noqa: E402 +from django.test import TestCase # noqa: E402 +from django.test.utils import CaptureQueriesContext # noqa: E402 +from workflow_manager.workflow_v2.models.workflow import Workflow # noqa: E402 + +from dashboard_metrics.tasks import AggregationTier, _run_aggregation # noqa: E402 -INDEX_NAME = "we_created_at_idx" _ROWS = 12000 _SPAN_DAYS = 255 @@ -96,15 +111,6 @@ def test_the_window_is_the_selectivity_the_index_is_for(self) -> None: share = cur.fetchone()[0] assert 0 < share < 0.10 - def test_the_planner_reaches_for_the_index(self) -> None: - """The whole point of 2a. An index that exists but is never chosen costs on - every insert and buys nothing. - """ - with connection.cursor() as cur: - cur.execute("EXPLAIN " + self._prefilter_sql()) - plan = "\n".join(row[0] for row in cur.fetchall()) - assert INDEX_NAME in plan, f"expected {INDEX_NAME} in:\n{plan}" - def test_the_prefilter_does_not_scan_the_executions_table(self) -> None: """The regression the index is meant to remove.""" with connection.cursor() as cur: diff --git a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py index 060f43e88e..c971ce3059 100644 --- a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py +++ b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py @@ -17,7 +17,6 @@ import importlib import inspect -import json import os from typing import Any from unittest import mock @@ -43,7 +42,11 @@ def _post(body: dict[str, Any]) -> tuple[int, Any]: - """POST to the aggregate endpoint with the task mocked; return status and its kwargs.""" + """POST to the aggregate endpoint with the task mocked. + + Returns the status and the kwargs the task was called with, or ``None`` if it was + never reached — which is what a rejected body has to look like. + """ view = internal_views.AggregateMetricsAPIView.as_view() request = APIRequestFactory().post(_ENDPOINT, body, format="json") with mock.patch.object( @@ -73,14 +76,13 @@ def test_an_omitted_tier_leaves_the_task_default_in_place(self) -> None: def test_an_unrecognised_tier_is_rejected_rather_than_ignored(self) -> None: """A silent no-op would look like a successful run that wrote nothing. - Runs against the real task, not the mock: the 400 comes from the ValueError the - task raises, and a mock would accept anything and return 200. The tier is - validated on the task's first line, so nothing touches the database. + The 400 is raised at the boundary, before the task is entered, so it cannot be + confused with a ValueError from inside the aggregation — that one belongs on + the logged 500 path. """ - view = internal_views.AggregateMetricsAPIView.as_view() - response = view(APIRequestFactory().post(_ENDPOINT, {"tier": "houry"}, format="json")) - assert response.status_code == 400 - assert "houry" in str(response.data) + status, called_with = _post({"tier": "houry"}) + assert status == 400 + assert called_with is None class TestTheBeatLegCarriesTheTier: @@ -114,10 +116,3 @@ def test_every_declared_tier_is_a_real_tier( there raises inside the task on every single run.""" for kwargs in declared_kwargs.values(): AggregationTier(kwargs["tier"]) - - def test_beat_stores_the_kwargs_as_json_the_task_can_receive(self) -> None: - """Beat's kwargs column is a JSON *string*; PgPeriodicTask's is a JSONField. - The Beat side has to round-trip back to the same mapping.""" - mod = importlib.import_module(_SPLIT_MIGRATION) - for spec in mod.AGGREGATION_SCHEDULES: - assert json.loads(json.dumps({"tier": spec["tier"]})) == {"tier": spec["tier"]} diff --git a/backend/dashboard_metrics/tests/test_aggregation_tier.py b/backend/dashboard_metrics/tests/test_aggregation_tier.py index 10c9fb6183..4cb3cfbc5e 100644 --- a/backend/dashboard_metrics/tests/test_aggregation_tier.py +++ b/backend/dashboard_metrics/tests/test_aggregation_tier.py @@ -10,6 +10,7 @@ from __future__ import annotations +import inspect import os import django @@ -25,6 +26,7 @@ _aggregation_lock_key, _writes_daily_monthly, _writes_hourly, + aggregate_metrics_from_sources, ) @@ -60,11 +62,23 @@ def test_no_tier_is_written_by_both_schedules(self) -> None: class TestTheDefaultIsAll: - def test_omitting_the_tier_writes_everything_rather_than_nothing(self) -> None: - """Between the code deploying and migration 0005 running, the schedule row - still carries no tier kwarg. Defaulting to anything narrower would stop writing - tiers during that window; defaulting to none would stop writing entirely. + """The property that keeps the deploy window safe, pinned at the signature. + + Between the code deploying and migration 0006 running, the schedule row still + carries no tier kwarg. Every other test in the suite passes a tier explicitly or + mocks the task, so none of them can see what the default actually is. + """ + + def test_the_signature_default_is_all(self) -> None: + """Narrower and daily/monthly stop being written for the whole window; none + and nothing is written at all. Both look like successful runs. """ + default = inspect.signature(aggregate_metrics_from_sources).parameters[ + "tier" + ].default + assert default == AggregationTier.ALL + + def test_the_default_writes_everything_rather_than_nothing(self) -> None: assert AggregationTier("all") is AggregationTier.ALL assert _writes_hourly(AggregationTier.ALL) assert _writes_daily_monthly(AggregationTier.ALL) diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py index ee90e93375..728a1def82 100644 --- a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -19,6 +19,7 @@ import importlib import json +from types import SimpleNamespace from typing import Any import pytest @@ -131,15 +132,20 @@ class _SplitRecorder: only the named fields. Conflating them is exactly the bug this guards. """ - def __init__(self) -> None: + def __init__(self, existing: Any = None) -> None: self.created: dict[str, dict[str, Any]] = {} self.updated: dict[str, dict[str, Any]] = {} + self.bumps = 0 + self._existing = existing self._filtered_on: str = "" def filter(self, name: str = "", **_kw: Any) -> _SplitRecorder: self._filtered_on = name return self + def first(self) -> Any: + return self._existing + def update(self, **kwargs: Any) -> int: self.updated[self._filtered_on] = kwargs return 1 @@ -147,6 +153,9 @@ def update(self, **kwargs: Any) -> int: def update_or_create( self, name: str = "", defaults: dict[str, Any] | None = None, **_kw: Any ) -> tuple[dict[str, Any], bool]: + if not name: # PeriodicTasks(ident=1) — the Beat reload tracker + self.bumps += 1 + return defaults or {}, True self.created[name] = defaults or {} return self.created[name], True @@ -157,19 +166,33 @@ def delete(self) -> tuple[int, dict[str, Any]]: return (0, {}) -@pytest.fixture(scope="module") -def split() -> dict[str, _SplitRecorder]: - """Run 0005's forward function against fakes and capture both tables.""" +def _run_split(beat_row: Any = None, pg_row: Any = None) -> dict[str, _SplitRecorder]: + """Run 0006's forward function against fakes and capture every table it writes.""" mod = importlib.import_module(_SPLIT_MIGRATION) - beat, pg, crontab = _SplitRecorder(), _SplitRecorder(), _SplitRecorder() + beat = _SplitRecorder(existing=beat_row) + pg = _SplitRecorder(existing=pg_row) + crontab, tracker = _SplitRecorder(), _SplitRecorder() class _Apps: def get_model(self, _app: str, model: str) -> type: - table = {"PeriodicTask": beat, "PgPeriodicTask": pg}.get(model, crontab) + table = { + "PeriodicTask": beat, + "PgPeriodicTask": pg, + "PeriodicTasks": tracker, + }.get(model, crontab) return type("M", (), {"objects": table}) mod.split_schedules(_Apps(), None) - return {"beat": beat, "pg": pg} + return {"beat": beat, "pg": pg, "tracker": tracker} + + +@pytest.fixture(scope="module") +def split() -> dict[str, _SplitRecorder]: + """The default case: a Beat-owned row, as every environment ships today.""" + return _run_split( + beat_row=SimpleNamespace(enabled=True), + pg_row=SimpleNamespace(enabled=True, pg_owned=False), + ) class TestTheSplitAddsOneRowAndRewritesOne: @@ -188,9 +211,18 @@ def test_the_new_row_is_declared_the_same_on_both_tables(self, split: dict[str, assert pg["task_kwargs"] == json.loads(beat["kwargs"]) def test_the_new_row_runs_hourly_on_both_tables(self, split: dict[str, _SplitRecorder]) -> None: - assert split["pg"].created[_NEW_ROW]["cron_string"] == "0 * * * *" + assert split["pg"].created[_NEW_ROW]["cron_string"] == "20 * * * *" crontab = split["beat"].created[_NEW_ROW]["crontab"] - assert (crontab["minute"], crontab["hour"]) == ("0", "*") + assert (crontab["minute"], crontab["hour"]) == ("20", "*") + + def test_the_two_rows_never_start_together(self, split: dict[str, _SplitRecorder]) -> None: + """The per-tier locks are built so the two runs cannot block each other, so a + shared minute is two full prefilter scans and two per-org loops at once — on a + change whose object is flattening cron load. + """ + fires_at = {0, 15, 30, 45} # the existing row's */15 + minute = int(split["beat"].created[_NEW_ROW]["crontab"]["minute"]) + assert minute not in fires_at def test_the_new_row_is_seeded_inert_on_the_pg_side(self, split: dict[str, _SplitRecorder]) -> None: """Same reason as 0004's rows: a PG row that is pg_owned before the scheduler @@ -198,12 +230,81 @@ def test_the_new_row_is_seeded_inert_on_the_pg_side(self, split: dict[str, _Spli """ assert split["pg"].created[_NEW_ROW]["pg_owned"] is False + def test_the_rewritten_row_carries_the_same_kwargs_on_both_tables( + self, split: dict[str, _SplitRecorder] + ) -> None: + """The row firing the hourly tier every 15 minutes in production. + + Beat's ``kwargs`` is a TextField it parses with ``json.loads``; writing the + mapping rather than its JSON encoding stores a Python repr, ``ModelEntry`` + raises, and the hourly aggregation silently stops firing. + """ + beat = split["beat"].updated[_EXISTING_ROW] + assert json.loads(beat["kwargs"]) == split["pg"].updated[_EXISTING_ROW][ + "task_kwargs" + ] + def test_the_two_rows_ask_for_different_tiers(self, split: dict[str, _SplitRecorder]) -> None: new = split["pg"].created[_NEW_ROW]["task_kwargs"]["tier"] existing = split["pg"].updated[_EXISTING_ROW]["task_kwargs"]["tier"] assert new != existing +class TestTheNewRowInheritsWhoeverFiresTheRowItSplitsFrom: + """Hardcoding Beat leaves the daily/monthly tier with no firer in a PG-adopted + environment: the adopted row's Beat twin is disabled and Beat may be scaled to + zero, so the sole writer of those figures never runs and the hourly run still + reports success. + """ + + def test_a_pg_adopted_row_hands_its_new_half_to_the_pg_scheduler(self) -> None: + split = _run_split( + beat_row=SimpleNamespace(enabled=False), + pg_row=SimpleNamespace(enabled=True, pg_owned=True), + ) + assert split["pg"].created[_NEW_ROW]["pg_owned"] is True + assert split["pg"].created[_NEW_ROW]["enabled"] is True + assert split["beat"].created[_NEW_ROW]["enabled"] is False + + def test_a_disabled_row_does_not_come_back_as_an_enabled_half(self) -> None: + split = _run_split( + beat_row=SimpleNamespace(enabled=False), + pg_row=SimpleNamespace(enabled=False, pg_owned=False), + ) + assert split["beat"].created[_NEW_ROW]["enabled"] is False + assert split["pg"].created[_NEW_ROW]["enabled"] is False + + def test_a_missing_row_falls_back_to_beat(self) -> None: + """A fresh install applies 0002/0004 first, so this is defensive only.""" + split = _run_split(beat_row=None, pg_row=None) + assert split["beat"].created[_NEW_ROW]["enabled"] is True + assert split["pg"].created[_NEW_ROW]["pg_owned"] is False + + +class TestARunningBeatIsToldToReload: + """Historical models fire no post_save, so DatabaseScheduler never reloads. + + Without an explicit bump the existing row keeps firing with no tier and the new + row never fires at all — no error, nothing logged, and the whole saving silently + does not happen. + """ + + def test_the_forward_direction_bumps_the_change_tracker(self, split) -> None: + assert split["tracker"].bumps == 1 + + def test_the_reverse_direction_bumps_it_too(self) -> None: + mod = importlib.import_module(_SPLIT_MIGRATION) + tracker, other = _SplitRecorder(), _SplitRecorder() + + class _Apps: + def get_model(self, _app: str, model: str) -> type: + table = tracker if model == "PeriodicTasks" else other + return type("M", (), {"objects": table}) + + mod.merge_schedules(_Apps(), None) + assert tracker.bumps == 1 + + class TestTheRewriteLeavesSchedulerOwnershipAlone: """The existing row may already be owned by the PG scheduler, with its Beat twin disabled by converge_pg_scheduler. Rewriting `pg_owned` or `enabled` here would diff --git a/backend/dashboard_metrics/tests/test_tier_split_equivalence.py b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py index ed9b08c975..13599deb2c 100644 --- a/backend/dashboard_metrics/tests/test_tier_split_equivalence.py +++ b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py @@ -21,18 +21,27 @@ from datetime import timedelta from typing import Any -from account_v2.models import Organization -from django.db import connection -from django.test import TestCase -from django.utils import timezone -from workflow_manager.workflow_v2.models.workflow import Workflow +import os -from dashboard_metrics.models import ( +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from account_v2.models import Organization # noqa: E402 +from django.db import connection # noqa: E402 +from django.test import TestCase # noqa: E402 +from django.utils import timezone # noqa: E402 +from workflow_manager.workflow_v2.models.workflow import Workflow # noqa: E402 + +from dashboard_metrics.models import ( # noqa: E402 EventMetricsDaily, EventMetricsHourly, EventMetricsMonthly, ) -from dashboard_metrics.tasks import AggregationTier, _run_aggregation +from dashboard_metrics.tasks import AggregationTier, _run_aggregation # noqa: E402 # (model, the column naming its period) — the period field differs per tier. _TIERS = [ diff --git a/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py b/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py index 93317989eb..f4f1d2d759 100644 --- a/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py +++ b/backend/workflow_manager/workflow_v2/tests/test_we_created_at_idx.py @@ -78,6 +78,21 @@ def test_it_builds_and_drops_concurrently(self) -> None: assert "CREATE INDEX CONCURRENTLY IF NOT EXISTS" in sql assert "DROP INDEX CONCURRENTLY IF EXISTS" in sql + def test_the_statement_that_runs_names_the_model_table_and_column(self) -> None: + """Everything else here reads the ``AddIndex`` state operation, which by + construction never reaches the database, or greps the source for + ``CONCURRENTLY``. The ``RunSQL`` is the only statement production executes and + its table and column were cross-checked against nothing — so the index could be + built on the wrong column while Django's model state claimed otherwise. + """ + index = _model_index() + assert index is not None + model = apps.get_model("workflow_v2", "WorkflowExecution") + expected = f"{model._meta.db_table} ({', '.join(index.fields)})" + + create = _operations()[0].database_operations[0] + assert expected in create.sql, f"expected {expected!r} in {create.sql!r}" + def test_it_guards_against_a_leftover_invalid_index(self) -> None: """An interrupted CONCURRENTLY build leaves an INVALID index that costs on every write and is never read. IF NOT EXISTS would keep it while Django recorded the From 2398a4332e5c57a9b45da1d03cc7f64dfe059bc7 Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 2 Sep 2026 17:26:38 +0530 Subject: [PATCH 13/14] UN-3973 [FIX] Address Athul's review: upsert-only monthly, per-schedule lock, honest success MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The orphan delete is removed. The design agreed on this ticket (comments 44768/45016) is a pure INSERT ... ON CONFLICT DO UPDATE; the delete was scope beyond it, and it converts a recoverable undercount into unrecoverable loss — the daily rows that would rebuild a deleted monthly row are exactly the ones that were missing. A stale total is recoverable with backfill_metrics. The reconciliation pass no longer shares a lock key with the 15-minute schedule. The 15-minute row is an IntervalSchedule and drifts against a fixed crontab, so on a shared key the once-daily repair loses the race roughly one day in seven, returns skipped=True and is never retried. A run in which every metric for every org failed no longer reports success: True. The result's success now reflects the error count, the completion log rises to WARNING, and the worker-side guard reads skipped_reason and errors as well as skipped — it saw none of these three did-nothing shapes before. A failed monthly rollup is distinguishable from an empty one: upserted=0 collided with the legitimate no-op and the no_active_orgs return. source_window_days is validated and bounded. It arrives as JSON from a Beat row that is editable in the admin: negative puts the window in the future, 0 never refreshes yesterday, 365 restores the multi-month scan this ticket exists to remove. Tests: a golden test seeds source rows, lets the real aggregation populate daily, and compares the rolled-up monthly against the pre-change derivation computed independently from get_documents_processed — AC-4 was claimed Met and had no equivalence assertion. Fixture offsets derive from the month boundary rather than fixed day counts, which land in the wrong month for the last days of any month. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/README.md | 11 +- .../management/commands/backfill_metrics.py | 17 +- backend/dashboard_metrics/tasks.py | 176 +++++----- backend/dashboard_metrics/tests/test_tasks.py | 311 ++++++++++++++---- workers/scheduler/dashboard_metrics_tasks.py | 19 +- 5 files changed, 373 insertions(+), 161 deletions(-) diff --git a/backend/dashboard_metrics/README.md b/backend/dashboard_metrics/README.md index d4f77b9eb7..8c3fee6ce8 100644 --- a/backend/dashboard_metrics/README.md +++ b/backend/dashboard_metrics/README.md @@ -382,9 +382,10 @@ The `aggregate_metrics_from_sources` task: - Queries source table with `MetricsQueryService` - Groups by time period (hour/day/month) 3. **Upserts results** into the hourly and daily tables -4. **Rolls monthly up from the daily tier** in one statement for all orgs, dropping - monthly rows the daily tier no longer produces (scoped to the organization/month - partitions the rollup covered) +4. **Rolls monthly up from the daily tier** in one statement for all orgs. Upsert-only: + a monthly row the daily tier no longer produces is left in place. A stale total is + recoverable with `backfill_metrics`; a deleted one is not, because the daily rows + that would rebuild it are exactly what is missing 5. **Uses `_base_manager`** to bypass Django's organization filter in Celery context ```python @@ -396,7 +397,9 @@ monthly_start = first_of_previous_month # summed from dail The monthly tier has no source queries of its own. `backfill_metrics` still computes monthly from source, so within the rollup window (current + previous month) its output -is recomputed within 15 minutes — see that command's help text. +is overwritten within 15 minutes by the sum of the daily tier — see that command's help +text. **Backfill daily before relying on monthly:** the rollup writes whatever daily +holds, so a month whose daily tier is short produces an under-counted monthly total. --- diff --git a/backend/dashboard_metrics/management/commands/backfill_metrics.py b/backend/dashboard_metrics/management/commands/backfill_metrics.py index d598e366bb..a3af666ede 100644 --- a/backend/dashboard_metrics/management/commands/backfill_metrics.py +++ b/backend/dashboard_metrics/management/commands/backfill_metrics.py @@ -3,11 +3,11 @@ This command populates EventMetricsHourly, EventMetricsDaily, and EventMetricsMonthly tables from historical data in source tables (Usage, PageUsage, WorkflowExecution, etc.) -The current and previous month are owned exclusively by the aggregation task, which -derives them from the daily tier and drops monthly rows the daily tier no longer -produces. Inside that window this command's monthly output is recomputed within 15 -minutes, so --skip-monthly is a no-op and --skip-daily leaves monthly to be rebuilt -from a tier this run did not populate. Backfill both, or neither. +The current and previous month are recomputed from the daily tier by the aggregation +task every 15 minutes, so inside that window this command's monthly output is +overwritten and --skip-monthly is a no-op. --skip-daily is worse than useless there: +monthly is rebuilt from a tier this run did not populate, producing an under-count. +Backfill both, or neither. Usage: python manage.py backfill_metrics --days=30 @@ -100,7 +100,8 @@ def add_arguments(self, parser): action="store_true", help=( "Skip daily aggregation. Unsafe for the current and previous month: " - "the aggregation task rebuilds monthly from daily there." + "the aggregation task rebuilds monthly from daily there, so monthly " + "ends up under-counted." ), ) parser.add_argument( @@ -139,8 +140,8 @@ def handle(self, *args, **options): self.stdout.write( self.style.WARNING( "--skip-daily without --skip-monthly: the aggregation task " - "rebuilds the current and previous month from the daily tier " - "and will drop monthly rows this run writes there." + "rebuilds the current and previous month from the daily tier, " + "so monthly will be overwritten with an under-count." ) ) diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 332cb623ef..f9b896f68a 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -14,7 +14,6 @@ from account_v2.models import Organization from celery import shared_task from django.core.cache import cache -from django.db import transaction from django.db.models import Min, Sum from django.db.models.functions import TruncMonth from django.db.utils import DatabaseError, OperationalError @@ -49,8 +48,11 @@ # in the schedule row's kwargs — editing this constant does not move the schedule. DASHBOARD_RECONCILE_WINDOW_DAYS = 7 -# Wider than the source window: metrics keyed on another column -# (e.g. approved_at) can land for an org whose executions are older. +# Floor on the prefilter lookback: metrics keyed on another column (e.g. +# approved_at) can land for an org whose executions are older. _active_org_ids +# takes the wider of this and the run's own window, so the prefilter is never +# narrower than what is being queried — at the 7-day reconciliation window the +# two are equal rather than this one being wider. DASHBOARD_ACTIVE_ORG_LOOKBACK_DAYS = 7 @@ -185,54 +187,17 @@ def _bulk_upsert_daily(aggregations: dict) -> int: return len(objects) -def _delete_orphan_monthly(objects: list[EventMetricsMonthly]) -> int: - """Drop stale monthly rows inside the partitions the rollup actually covered. - - Monthly derives from the daily tier, so a key with no daily rows left must - not survive as a stale total. An (organization, month) the rollup produced - nothing for is left alone instead: absent daily rows there mean an - incomplete tier, not a metric that went to zero. - """ - fresh_keys = { - (o.organization_id, o.month, o.metric_name, o.project, o.tag) for o in objects - } - covered: dict[date, set] = {} - for o in objects: - covered.setdefault(o.month, set()).add(o.organization_id) - - deleted = 0 - for month, org_ids in covered.items(): - stale_pks = [ - row["pk"] - for row in EventMetricsMonthly._base_manager.filter( - month=month, organization_id__in=org_ids - ).values("pk", "organization_id", "month", "metric_name", "project", "tag") - if ( - row["organization_id"], - row["month"], - row["metric_name"], - row["project"], - row["tag"], - ) - not in fresh_keys - ] - if not stale_pks: - continue - removed, _ = EventMetricsMonthly._base_manager.filter(pk__in=stale_pks).delete() - deleted += removed - - return deleted - - -def _rollup_monthly_from_daily(month_start: date) -> tuple[int, int]: +def _rollup_monthly_from_daily(month_start: date) -> int: """Sum the daily tier from month_start into monthly, for all orgs at once. + Upsert-only, per the design agreed on UN-3973: a monthly row the daily tier + no longer produces is left in place rather than deleted. A stale total is + recoverable with backfill_metrics; a deleted one is not, because the daily + rows that would rebuild it are exactly what is missing. + metric_type is aggregated rather than grouped: it is not part of unique_monthly_metric, so grouping on it could yield two rows for one conflict target. - - Returns: - (rows upserted, rows deleted as orphans) """ rows = ( EventMetricsDaily._base_manager.filter(date__gte=month_start) @@ -259,30 +224,36 @@ def _rollup_monthly_from_daily(month_start: date) -> tuple[int, int]: for row in rows ] - # Nothing to write and nothing covered, so nothing to sweep. The scoping in - # _delete_orphan_monthly already makes this safe; returning early just skips a - # pointless transaction. if not objects: - return 0, 0 - - with transaction.atomic(): - EventMetricsMonthly._base_manager.bulk_create( - objects, - update_conflicts=True, - unique_fields=["organization", "month", "metric_name", "project", "tag"], - update_fields=["metric_type", "metric_value", "metric_count"], - batch_size=MONTHLY_ROLLUP_BATCH_SIZE, - ) - deleted = _delete_orphan_monthly(objects) + return 0 - return len(objects), deleted + EventMetricsMonthly._base_manager.bulk_create( + objects, + update_conflicts=True, + unique_fields=["organization", "month", "metric_name", "project", "tag"], + update_fields=["metric_type", "metric_value", "metric_count"], + batch_size=MONTHLY_ROLLUP_BATCH_SIZE, + ) + return len(objects) -AGGREGATION_LOCK_KEY = "dashboard_metrics:aggregation_lock" +AGGREGATION_LOCK_KEY_PREFIX = "dashboard_metrics:aggregation_lock" AGGREGATION_LOCK_TIMEOUT = 900 # 15 minutes (matches task schedule) -def _acquire_aggregation_lock() -> bool: +def _aggregation_lock_key(source_window_days: int) -> str: + """One key per schedule, not one key for the task. + + The 15-minute row is an IntervalSchedule and drifts against the reconciliation + pass's fixed crontab, so on a shared key the once-daily repair would lose the + race and return skipped=True — never retried, and the only mechanism that + repairs the narrowed window. Distinct windows are distinct jobs; both are + idempotent upserts, so the once-a-day overlap costs duplicated work at worst. + """ + return f"{AGGREGATION_LOCK_KEY_PREFIX}:{source_window_days}d" + + +def _acquire_aggregation_lock(lock_key: str) -> bool: """Acquire the distributed aggregation lock with self-healing. Stores a Unix timestamp as the lock value. If a previous run crashed @@ -295,22 +266,22 @@ def _acquire_aggregation_lock() -> bool: now = time.time() # Fast path: lock is free - if cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT): + if cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT): return True # Lock exists — check if it's stale (previous run died without releasing) - lock_value = cache.get(AGGREGATION_LOCK_KEY) + lock_value = cache.get(lock_key) if lock_value is None: # Expired between our check and get — lock is now free, try to acquire it - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) try: lock_time = float(lock_value) except (TypeError, ValueError): # Corrupted value (e.g. old "running" string) — reclaim it logger.warning("Reclaiming aggregation lock with invalid value: %s", lock_value) - cache.delete(AGGREGATION_LOCK_KEY) - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + cache.delete(lock_key) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) age = now - lock_time if age > AGGREGATION_LOCK_TIMEOUT: @@ -319,8 +290,8 @@ def _acquire_aggregation_lock() -> bool: age, AGGREGATION_LOCK_TIMEOUT, ) - cache.delete(AGGREGATION_LOCK_KEY) - return cache.add(AGGREGATION_LOCK_KEY, str(now), AGGREGATION_LOCK_TIMEOUT) + cache.delete(lock_key) + return cache.add(lock_key, str(now), AGGREGATION_LOCK_TIMEOUT) return False @@ -350,14 +321,26 @@ def aggregate_metrics_from_sources( Returns: Dict with aggregation summary for all three tiers """ - if not _acquire_aggregation_lock(): - logger.info("Skipping aggregation — another run is in progress") - return {"success": True, "skipped": True, "reason": "lock_held"} + source_window_days = _validate_source_window(source_window_days) + lock_key = _aggregation_lock_key(source_window_days) + + if not _acquire_aggregation_lock(lock_key): + logger.warning( + "Skipping the %d-day aggregation — another run of the same schedule is " + "in progress", + source_window_days, + ) + return { + "success": True, + "skipped": True, + "reason": "lock_held", + "source_window_days": source_window_days, + } try: return _run_aggregation(source_window_days) finally: - cache.delete(AGGREGATION_LOCK_KEY) + cache.delete(lock_key) def _aggregate_single_metric( @@ -572,7 +555,10 @@ def _build_result( ) -> dict[str, Any]: """Shape the task's return value from the accumulated stats.""" result = { - "success": True, + # Not a literal: every metric for every org can fail while each exception is + # caught per-metric, and the run would otherwise report 200 / success with + # zero rows written and the dashboard frozen. + "success": stats["errors"] == 0, "organizations_processed": stats["orgs_processed"], "hourly": stats["hourly"], "daily": stats["daily"], @@ -589,10 +575,33 @@ def _build_result( return result +# A negative window puts daily_start in the future so nothing matches; 0 never +# refreshes yesterday; an unbounded one restores the multi-month per-org scan this +# change exists to remove, past soft_time_limit. +MAX_SOURCE_WINDOW_DAYS = 90 + + +def _validate_source_window(source_window_days: int) -> int: + """Coerce and bound the window. It arrives as JSON from an editable Beat row.""" + try: + days = int(source_window_days) + except (TypeError, ValueError) as exc: + raise ValueError( + f"source_window_days must be an integer, got {source_window_days!r}" + ) from exc + if not 1 <= days <= MAX_SOURCE_WINDOW_DAYS: + raise ValueError( + f"source_window_days must be between 1 and {MAX_SOURCE_WINDOW_DAYS}, " + f"got {days}" + ) + return days + + def _run_aggregation( source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, ) -> dict[str, Any]: """Execute the aggregation, separately from the task's lock handling.""" + source_window_days = _validate_source_window(source_window_days) end_date = timezone.now() # Monthly spans the current and previous month. @@ -605,7 +614,7 @@ def _run_aggregation( stats = { "hourly": {"upserted": 0}, "daily": {"upserted": 0}, - "monthly": {"upserted": 0, "deleted": 0}, + "monthly": {"upserted": 0, "failed": False}, "errors": 0, "orgs_processed": 0, } @@ -640,27 +649,24 @@ def _run_aggregation( stats["errors"] += 1 try: - upserted, deleted = _rollup_monthly_from_daily(monthly_start) - stats["monthly"]["upserted"] = upserted - stats["monthly"]["deleted"] = deleted - if deleted: - logger.warning( - "Monthly rollup deleted %d orphan row(s) from %s", deleted, monthly_start - ) + stats["monthly"]["upserted"] = _rollup_monthly_from_daily(monthly_start) except (DatabaseError, OperationalError): # Configured on the task for autoretry — swallowing them here would # leave monthly permanently stale behind successful-looking runs. raise except Exception: + # upserted stays 0, which is also the legitimate empty-rollup value, so + # mark the failure explicitly rather than letting the two collapse. logger.exception("Error rolling up monthly metrics from %s", monthly_start) + stats["monthly"]["failed"] = True stats["errors"] += 1 - logger.info( + log = logger.warning if stats["errors"] else logger.info + log( f"Aggregation completed: {stats['orgs_processed']} orgs, " f"hourly={stats['hourly']['upserted']}, " f"daily={stats['daily']['upserted']}, " f"monthly={stats['monthly']['upserted']}, " - f"monthly_deleted={stats['monthly']['deleted']}, " f"errors={stats['errors']}" ) diff --git a/backend/dashboard_metrics/tests/test_tasks.py b/backend/dashboard_metrics/tests/test_tasks.py index e85acd203b..e7911226d6 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -1,12 +1,14 @@ """Unit tests for Dashboard Metrics Celery tasks.""" import json +import time from datetime import date, datetime, timedelta from importlib import import_module from types import SimpleNamespace from unittest.mock import patch from django.apps import apps +from django.core.cache import cache from django.db import connection from django.db.utils import DatabaseError from django.test import TestCase @@ -19,6 +21,7 @@ EventMetricsDaily, EventMetricsHourly, EventMetricsMonthly, + Granularity, MetricType, ) from pg_queue.models import PgPeriodicTask @@ -27,15 +30,20 @@ from workflow_manager.workflow_v2.models.execution import WorkflowExecution from workflow_manager.workflow_v2.models.workflow import Workflow from dashboard_metrics.internal_views import AggregateMetricsAPIView +from dashboard_metrics.services import MetricsQueryService from dashboard_metrics.tasks import ( + AGGREGATION_LOCK_TIMEOUT, DASHBOARD_RECONCILE_WINDOW_DAYS, DASHBOARD_SOURCE_WINDOW_DAYS, + _acquire_aggregation_lock, _active_org_ids, + _aggregation_lock_key, _rollup_monthly_from_daily, _run_aggregation, _truncate_to_day, _truncate_to_hour, _truncate_to_month, + _validate_source_window, aggregate_metrics_from_sources, cleanup_daily_metrics, cleanup_hourly_metrics, @@ -264,7 +272,7 @@ def test_sums_daily_rows_into_month_bucket(self): self._daily(date(2024, 3, 5), value=10, count=2) self._daily(date(2024, 3, 18), value=32, count=4) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 rows = self._monthly_rows() assert len(rows) == 1 @@ -279,7 +287,7 @@ def test_month_boundary_keeps_months_separate(self): self._daily(date(2024, 2, 1), value=100) self._daily(date(2024, 2, 2), value=200) - assert _rollup_monthly_from_daily(date(2024, 1, 1)) == (2, 0) + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 2 rows = self._monthly_rows() assert [r.month for r in rows] == [date(2024, 1, 1), date(2024, 2, 1)] @@ -290,7 +298,7 @@ def test_excludes_months_before_the_window(self): self._daily(date(2023, 12, 15), value=999) self._daily(date(2024, 1, 15), value=5) - assert _rollup_monthly_from_daily(date(2024, 1, 1)) == (1, 0) + assert _rollup_monthly_from_daily(date(2024, 1, 1)) == 1 rows = self._monthly_rows() assert len(rows) == 1 @@ -314,7 +322,7 @@ def test_mixed_metric_type_within_a_month_yields_one_row(self): self._daily(date(2024, 3, 5), value=10, metric_type=MetricType.HISTOGRAM) self._daily(date(2024, 3, 6), value=5, metric_type=MetricType.COUNTER) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 rows = self._monthly_rows() assert len(rows) == 1 @@ -324,8 +332,8 @@ def test_an_empty_daily_tier_leaves_existing_monthly_rows_alone(self): """An empty tier means the source is gone, not that every month is zero. Seeding a monthly row first is what makes the failure reachable at all: with - an empty table a sweep that deletes everything and one that deletes nothing - both leave an empty table, and the assertion passes either way. + an empty table an implementation that wipes and one that writes nothing both + leave an empty table, and the assertion passes either way. """ EventMetricsMonthly._base_manager.create( organization=self.org, @@ -337,53 +345,58 @@ def test_an_empty_daily_tier_leaves_existing_monthly_rows_alone(self): project="default", ) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (0, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 0 rows = self._monthly_rows() assert len(rows) == 1 assert rows[0].metric_value == 42 - def test_metric_is_dropped_once_its_daily_rows_are_gone(self): - """A metric with no daily rows left must not keep a stale monthly total.""" + def test_a_metric_whose_daily_rows_are_gone_keeps_its_last_total(self): + """Upsert-only, per the design agreed on UN-3973. + + A stale total is recoverable — backfill_metrics rewrites it. A deleted row is + not, because the daily rows that would rebuild it are exactly what is missing. + """ self._daily(date(2024, 3, 5), value=10) self._daily(date(2024, 3, 6), value=7, metric_name="pages_processed") - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 2 EventMetricsDaily._base_manager.filter(metric_name="pages_processed").delete() - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 1) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 rows = self._monthly_rows() - assert [r.metric_name for r in rows] == ["documents_processed"] + assert [r.metric_name for r in rows] == ["documents_processed", "pages_processed"] - def test_a_month_with_no_daily_rows_of_its_own_is_left_alone(self): - """An incomplete daily tier must read as "unknown", not as "deleted". + def test_a_partially_repopulated_month_is_overwritten_not_accumulated(self): + """The realistic post-downtime shape: the daily tier comes back short. - A partially populated tier passes the empty-tier guard, so without scoping the - sweep to the (organization, month) partitions the rollup actually produced, - one missing month wipes that month's monthly rows for every org. + The total tracks whatever the daily tier currently holds, so repairing daily + repairs monthly on the next run — which is what makes upsert-only recoverable. """ self._daily(date(2024, 3, 5), value=10) - self._daily(date(2024, 4, 5), value=7) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) + self._daily(date(2024, 3, 6), value=32) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 + assert self._monthly_rows()[0].metric_value == 42 - EventMetricsDaily._base_manager.filter(date=date(2024, 3, 5)).delete() - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) + EventMetricsDaily._base_manager.filter(date=date(2024, 3, 6)).delete() + _rollup_monthly_from_daily(date(2024, 3, 1)) + assert self._monthly_rows()[0].metric_value == 10 - months = [row.month for row in self._monthly_rows()] - assert months == [date(2024, 3, 1), date(2024, 4, 1)] + self._daily(date(2024, 3, 6), value=32) + _rollup_monthly_from_daily(date(2024, 3, 1)) + assert self._monthly_rows()[0].metric_value == 42 - def test_the_sweep_is_scoped_per_organization(self): - """The sweep bypasses the org-scoped manager, so the org half of the key is - load-bearing: drop it and one org's daily rows vouch for another's monthly.""" + def test_rows_for_other_organizations_are_never_touched(self): + """The rollup goes through _base_manager, bypassing the org-scoped default.""" other = Organization.objects.create( organization_id="rollup-org-2", name="rollup-org-2", display_name="Other" ) self._daily(date(2024, 3, 5), value=10) self._daily(date(2024, 3, 6), value=20, org=other) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (2, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 2 EventMetricsDaily._base_manager.filter(organization=other).delete() - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (1, 0) + assert _rollup_monthly_from_daily(date(2024, 3, 1)) == 1 rows = self._monthly_rows() assert [(r.organization_id, r.metric_value) for r in rows] == [ @@ -391,28 +404,6 @@ def test_the_sweep_is_scoped_per_organization(self): (other.id, 20), ] - def test_a_stale_row_is_dropped_for_one_organization_only(self): - """Two orgs in one month: only the org whose metric vanished loses its row.""" - other = Organization.objects.create( - organization_id="rollup-org-3", name="rollup-org-3", display_name="Other" - ) - self._daily(date(2024, 3, 5), value=10) - self._daily(date(2024, 3, 5), value=1, metric_name="pages_processed") - self._daily(date(2024, 3, 6), value=20, org=other) - self._daily(date(2024, 3, 6), value=2, metric_name="pages_processed", org=other) - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (4, 0) - - EventMetricsDaily._base_manager.filter( - organization=other, metric_name="pages_processed" - ).delete() - assert _rollup_monthly_from_daily(date(2024, 3, 1)) == (3, 1) - - survivors = { - (r.organization_id, r.metric_name) for r in self._monthly_rows() - } - assert (self.org.id, "pages_processed") in survivors - assert (other.id, "pages_processed") not in survivors - def test_months_before_the_window_are_left_alone(self): """Orphan cleanup must not reach outside the rebuilt window.""" self._daily(date(2024, 1, 10), value=99) @@ -506,8 +497,9 @@ def test_a_database_error_propagates_instead_of_reporting_success(self): with self.assertRaises(DatabaseError): _run_aggregation() - def test_an_unexpected_error_is_still_counted_rather_than_fatal(self): - """Everything outside the retry set keeps the previous non-fatal posture.""" + def test_an_unexpected_error_is_counted_but_does_not_abort_the_run(self): + """Everything outside the retry set stays non-fatal — the hourly and daily + tiers this run already wrote are kept — but it is not reported as success.""" with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: prefilter = mock_execution.objects.filter.return_value prefilter.values_list.return_value.distinct.return_value = [1] @@ -516,8 +508,9 @@ def test_an_unexpected_error_is_still_counted_rather_than_fatal(self): side_effect=ValueError("bad row"), ): result = _run_aggregation() - assert result["success"] is True + assert result["success"] is False assert result["errors"] == 1 + assert result["monthly"]["failed"] is True class TestInternalAggregateEndpoint(TestCase): @@ -613,7 +606,7 @@ def test_the_window_covers_the_previous_month_and_spares_what_precedes_it(self): result = self._run() assert result["period"]["monthly"]["start"] == self.last_month.isoformat() - assert result["monthly"] == {"upserted": 2, "deleted": 0} + assert result["monthly"] == {"upserted": 2, "failed": False} rows = EventMetricsMonthly._base_manager.order_by("month") assert [r.month for r in rows] == [ @@ -622,14 +615,212 @@ def test_the_window_covers_the_previous_month_and_spares_what_precedes_it(self): self.this_month, ] - def test_the_run_reports_what_it_deleted(self): - """The only destructive write in the task has to reach the result.""" + def test_a_failed_rollup_is_not_reported_as_nothing_to_do(self): + """upserted stays 0 on failure, which is also the legitimate empty value. + + Three states used to collapse into one alongside success: True — failed, + empty, and no active orgs. + """ self._daily(self.this_month, value=10) - self._daily(self.this_month, value=1, metric_name="pages_processed") - assert self._run()["monthly"] == {"upserted": 2, "deleted": 0} + with patch( + "dashboard_metrics.tasks._rollup_monthly_from_daily", + side_effect=ValueError("bad row"), + ): + result = self._run() - EventMetricsDaily._base_manager.filter(metric_name="pages_processed").delete() - assert self._run()["monthly"] == {"upserted": 1, "deleted": 1} + assert result["monthly"] == {"upserted": 0, "failed": True} + assert result["success"] is False + + +class TestMonthlyMatchesTheOldDerivation(TestCase): + """AC-4: the new monthly figures equal the ones the source queries produced. + + Every other monthly test feeds hand-written daily rows in and checks the sum of + what it just wrote — self-consistency, not equivalence. This one seeds *source* + rows, lets the real aggregation populate the daily tier from them, and compares + the rolled-up monthly against the pre-change derivation computed independently: + `get_documents_processed` at DAY granularity, bucketed by month in Python. + + The window is deliberately wide enough to cover both months, which is the state + `backfill_metrics` establishes before this change is deployed. + """ + + def setUp(self): + self.org = Organization.objects.create( + organization_id="golden-org", name="golden-org", display_name="Golden Org" + ) + self.workflow = Workflow.objects.create( + workflow_name="golden-wf", organization=self.org + ) + # Offsets are derived from the month boundary, never fixed day counts: on the + # 25th of a month a hardcoded "25 days ago" lands in the current month and the + # cross-boundary coverage silently disappears. + now = timezone.now() + first_of_this_month = _truncate_to_month(now) + self.days_to_last_month_end = (now - first_of_this_month).days + 1 + self.days_to_last_month_start = ( + now - _truncate_to_month(first_of_this_month - timedelta(days=1)) + ).days + + def _seed(self, days_ago: int, count: int) -> None: + """Seed `count` completed file executions dated `days_ago`.""" + stamp = timezone.now() - timedelta(days=days_ago) + for n in range(count): + execution = WorkflowExecution.objects.create( + workflow=self.workflow, status=ExecutionStatus.COMPLETED + ) + file_execution = WorkflowFileExecution.objects.create( + workflow_execution=execution, + file_name=f"{days_ago}-{n}.pdf", + status=ExecutionStatus.COMPLETED.value, + ) + WorkflowFileExecution.objects.filter(pk=file_execution.pk).update( + created_at=stamp + ) + WorkflowExecution.objects.filter(pk=execution.pk).update(created_at=stamp) + + def _written(self) -> dict: + """Monthly totals as the rollup wrote them.""" + return { + row.month: row.metric_value + for row in EventMetricsMonthly._base_manager.filter( + metric_name="documents_processed" + ) + } + + def _oracle(self, monthly_start, end_date) -> dict: + """Monthly totals the way the code derived them before this change.""" + rows = MetricsQueryService.get_documents_processed( + organization_id=str(self.org.id), # tasks.py passes the numeric PK + start_date=monthly_start, + end_date=end_date, + granularity=Granularity.DAY, + ) + totals: dict = {} + for row in rows: + month = _truncate_to_month(row["period"]).date() + totals[month] = totals.get(month, 0) + row["value"] + return totals + + def test_monthly_equals_the_pre_change_figures_across_a_month_boundary(self): + self._seed(days_ago=0, count=3) + self._seed(days_ago=self.days_to_last_month_end, count=2) + self._seed(days_ago=self.days_to_last_month_start, count=4) + + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [self.org.id] + result = _run_aggregation( + source_window_days=self.days_to_last_month_start + 1 + ) + + monthly_start = date.fromisoformat(result["period"]["monthly"]["start"]) + end_date = datetime.fromisoformat(result["period"]["monthly"]["end"]) + expected = self._oracle( + datetime.combine(monthly_start, datetime.min.time(), tzinfo=end_date.tzinfo), + end_date, + ) + + assert len(expected) == 2, f"fixture must straddle a month boundary: {expected}" + assert self._written() == expected + + def test_the_comparison_can_fail_when_the_daily_tier_is_wrong(self): + """Guards the test above: an oracle that always matches proves nothing. + + Monthly is the sum of whatever the daily tier holds, so corrupting a day has + to move the monthly total away from the source-derived figure. Corrupting + rather than deleting is the point — under upsert-only a *deleted* day leaves + the previous monthly total in place, which is the recoverable state the + design accepts and is covered by TestMonthlyRollup. + """ + self._seed(days_ago=0, count=3) + self._seed(days_ago=self.days_to_last_month_end, count=2) + + with patch("dashboard_metrics.tasks.WorkflowExecution") as mock_execution: + prefilter = mock_execution.objects.filter.return_value + prefilter.values_list.return_value.distinct.return_value = [self.org.id] + result = _run_aggregation( + source_window_days=self.days_to_last_month_start + 1 + ) + + monthly_start = date.fromisoformat(result["period"]["monthly"]["start"]) + end_date = datetime.fromisoformat(result["period"]["monthly"]["end"]) + expected = self._oracle( + datetime.combine(monthly_start, datetime.min.time(), tzinfo=end_date.tzinfo), + end_date, + ) + assert self._written() == expected + + last_month_day = ( + timezone.now() - timedelta(days=self.days_to_last_month_end) + ).date() + corrupted = EventMetricsDaily._base_manager.filter( + date=last_month_day, metric_name="documents_processed" + ).update(metric_value=99) + assert corrupted, "fixture wrote no daily row for the previous month" + + _rollup_monthly_from_daily(monthly_start) + assert self._written() != expected + + +class TestTheLockIsPerSchedule(TestCase): + """The reconciliation pass must not lose a race it is never retried after.""" + + def test_the_two_schedules_take_different_keys(self): + assert _aggregation_lock_key( + DASHBOARD_SOURCE_WINDOW_DAYS + ) != _aggregation_lock_key(DASHBOARD_RECONCILE_WINDOW_DAYS) + + def test_a_held_key_does_not_block_the_other_schedule(self): + cache.clear() + assert _acquire_aggregation_lock( + _aggregation_lock_key(DASHBOARD_SOURCE_WINDOW_DAYS) + ) + # Same schedule: excluded, which is what the lock is for. + assert not _acquire_aggregation_lock( + _aggregation_lock_key(DASHBOARD_SOURCE_WINDOW_DAYS) + ) + # The reconciliation pass proceeds regardless. + assert _acquire_aggregation_lock( + _aggregation_lock_key(DASHBOARD_RECONCILE_WINDOW_DAYS) + ) + cache.clear() + + def test_a_stale_lock_is_reclaimed(self): + cache.clear() + key = _aggregation_lock_key(DASHBOARD_SOURCE_WINDOW_DAYS) + cache.set(key, str(time.time() - AGGREGATION_LOCK_TIMEOUT - 1), 3600) + assert _acquire_aggregation_lock(key) + cache.clear() + + def test_a_corrupted_lock_value_is_reclaimed(self): + cache.clear() + key = _aggregation_lock_key(DASHBOARD_SOURCE_WINDOW_DAYS) + cache.set(key, "running", 3600) + assert _acquire_aggregation_lock(key) + cache.clear() + + +class TestSourceWindowValidation(TestCase): + """The window arrives as JSON from a Beat row editable in the admin.""" + + def test_a_sane_window_passes_through(self): + assert _validate_source_window(7) == 7 + assert _validate_source_window("7") == 7 + + def test_a_window_that_would_query_nothing_is_rejected(self): + # Negative puts daily_start in the future; 0 never refreshes yesterday. + for bad in (-1, 0): + with self.assertRaises(ValueError): + _validate_source_window(bad) + + def test_a_window_that_restores_the_multi_month_scan_is_rejected(self): + with self.assertRaises(ValueError): + _validate_source_window(365) + + def test_a_non_integer_window_is_rejected(self): + with self.assertRaises(ValueError): + _validate_source_window("seven") class TestSourceWindow(TestCase): diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 07787aa9f3..90e016682a 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -92,16 +92,27 @@ def _call_internal( def _log_if_skipped(name: str, result: dict[str, Any]) -> None: - """Surface a lock-held no-op. + """Surface a run that did nothing, whatever shape the backend reported it in. - The backend returns success with ``skipped=True`` when the Redis lock is held. That - is correct behaviour, but left at INFO a permanently leaked lock looks like 96 - successful runs a day that did nothing. + Three of them, and only the first sets ``skipped``: the Redis lock was held + (``skipped``/``reason``), no organisation had recent activity + (``skipped_reason``), or every metric raised and was caught per-metric + (``errors``). Each is correct behaviour in isolation, but left at INFO a leaked + lock or a frozen source table looks like 96 successful runs a day. """ if result.get("skipped"): logger.warning( "%s did no work: %s", name, result.get("reason", "reported skipped=True") ) + elif result.get("skipped_reason"): + logger.warning("%s did no work: %s", name, result["skipped_reason"]) + elif result.get("errors"): + logger.warning( + "%s completed with %s error(s) across %s organisation(s)", + name, + result["errors"], + result.get("organizations_processed", "?"), + ) @worker_task(name="dashboard_metrics.aggregate_from_sources") From 846f3da9446b8ed87d170802d7b4112a41fe3e4b Mon Sep 17 00:00:00 2001 From: kirtimanmishrazipstack Date: Wed, 2 Sep 2026 17:27:23 +0530 Subject: [PATCH 14/14] UN-3974 [FIX] Address Athul's review: lock covers what is written, both kwargs, graph guard MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The lock is keyed by granularity written, not by enum member. ALL took a third key that excluded nothing, so an ALL run and the scheduled hourly run wrote EventMetricsHourly concurrently — reachable from the documented manual trigger and from the endpoint's own "omit tier" contract. ALL now takes both keys and releases whatever it took if it cannot take them all. Keys are namespaced by source window so the once-daily reconciliation pass, which is never retried, is not starved by the 15-minute schedule. source_window_days is accepted on all three legs. 0006 hard-depends on 0005, so the reconciliation row is a certainty rather than a hypothetical, and this branch's signatures rejected the kwarg it dispatches. The tier predicates come from one membership table, so a member added without an entry raises instead of acquiring the lock, iterating every org, writing nothing and returning success. The migration's bulk updates check their row counts. A filtered update matching nothing reported success while leaving the old row on kwargs="{}" — every tier every 15 minutes — alongside the new hourly row: strictly more load than before, silently. tier is validated at the request boundary with a warning log, and an explicit null is treated as omitted. New test_migration_graph.py builds the migration graph, which is what catches 0006's dependency on a node that is not on this branch; --no-migrations means nothing else does. Lock behaviour is now exercised rather than its key string asserted, merge_schedules has coverage at all, the equivalence file carries one absolute expectation and a frozen clock, and the prefilter asserts the index is usable under enable_seqscan=off rather than that the planner chose it on 12,000 synthetic rows. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_01KGZBF68CShem3pbUJM2tBc --- backend/dashboard_metrics/internal_views.py | 42 ++++--- .../0006_split_aggregation_schedule.py | 37 +++++- backend/dashboard_metrics/tasks.py | 112 ++++++++++++++++-- .../tests/test_active_org_prefilter.py | 30 ++++- .../tests/test_aggregation_dispatch.py | 25 +++- .../tests/test_aggregation_tier.py | 100 ++++++++++++++-- .../tests/test_migration_graph.py | 53 +++++++++ .../test_pg_periodic_task_declarations.py | 87 +++++++++++++- .../tests/test_tier_split_equivalence.py | 50 +++++++- workers/scheduler/dashboard_metrics_tasks.py | 18 ++- workers/tests/test_dashboard_metrics_tasks.py | 37 +++++- 11 files changed, 527 insertions(+), 64 deletions(-) create mode 100644 backend/dashboard_metrics/tests/test_migration_graph.py diff --git a/backend/dashboard_metrics/internal_views.py b/backend/dashboard_metrics/internal_views.py index 64bf048691..fce4fc978d 100644 --- a/backend/dashboard_metrics/internal_views.py +++ b/backend/dashboard_metrics/internal_views.py @@ -34,6 +34,7 @@ from utils.local_context import StateStore from dashboard_metrics.tasks import ( + DASHBOARD_SOURCE_WINDOW_DAYS, AggregationTier, aggregate_metrics_from_sources, cleanup_daily_metrics, @@ -59,6 +60,15 @@ def _clear_org_context() -> None: StateStore.clear(Account.ORGANIZATION_ID) +def _tier_arg(raw: Any) -> AggregationTier: + """Coerce a request body's tier. An explicit null means "omitted", not "none".""" + try: + return AggregationTier(raw) + except ValueError as exc: + valid = [member.value for member in AggregationTier] + raise ValueError(f"tier must be one of {valid}, got {raw!r}") from exc + + def _int_arg(request: Request, key: str, default: int) -> int: """Read an optional positive integer from the request body.""" raw = request.data.get(key, default) if isinstance(request.data, dict) else default @@ -94,26 +104,28 @@ class AggregateMetricsAPIView(_MetricsTaskAPIView): Calls the Celery task body verbatim, Redis lock included — this endpoint exists only because the PG consumer has no Django, not to change what the job does. - Optional ``tier`` in the body selects which metric tiers to write; omitting it - writes all of them. An unrecognised value is a 400 raised here at the boundary, so - a ValueError from inside the aggregation stays a logged 500 rather than reading as - a bad request. + Two optional body fields, both validated here at the boundary so a ValueError + from inside the ten-minute aggregation stays a logged 500 rather than reading as + a bad request: ``tier`` selects which tiers to write, ``source_window_days`` + widens the daily lookback for the reconciliation pass. Omitting either — or + sending it as ``null`` — applies the task's own default. """ def post(self, request: Request) -> Response: body = request.data if isinstance(request.data, dict) else {} - tier = body.get("tier") - if tier is None: - return self._run(aggregate_metrics_from_sources) + kwargs: dict[str, Any] = {} try: - tier = AggregationTier(tier) - except ValueError: - valid = [member.value for member in AggregationTier] - return Response( - {"error": f"tier must be one of {valid}, got {tier!r}"}, - status=status.HTTP_400_BAD_REQUEST, - ) - return self._run(aggregate_metrics_from_sources, tier=tier) + if body.get("tier") is not None: + kwargs["tier"] = _tier_arg(body["tier"]) + if body.get("source_window_days") is not None: + kwargs["source_window_days"] = _int_arg( + request, "source_window_days", DASHBOARD_SOURCE_WINDOW_DAYS + ) + except ValueError as exc: + # The one branch _run no longer covers, so it is logged here or nowhere. + logger.warning("dashboard-metrics aggregate rejected: %s", exc) + return Response({"error": str(exc)}, status=status.HTTP_400_BAD_REQUEST) + return self._run(aggregate_metrics_from_sources, **kwargs) class CleanupHourlyMetricsAPIView(_MetricsTaskAPIView): diff --git a/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py index 9b60398d19..21fda1bff2 100644 --- a/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py +++ b/backend/dashboard_metrics/migrations/0006_split_aggregation_schedule.py @@ -28,8 +28,12 @@ AGGREGATE_TASK_NAME = "dashboard_metrics.aggregate_from_sources" AGGREGATE_QUEUE = "dashboard_metric_events" -# Frozen literals — migrations must not import app enums. Kept in step with -# dashboard_metrics.tasks.AggregationTier. +# Frozen wire values, not a copy to keep in step with the enum. These are written +# into rows this migration never re-runs against, so editing them to follow a rename +# of AggregationTier changes nothing in production and every test stays green while +# live rows still carry the old string — the task then raises ValueError on every +# run. Renaming an AggregationTier value needs a NEW data migration that rewrites the +# rows; this file is a record of what was written on the day it ran. TIER_HOURLY = "hourly" TIER_DAILY_MONTHLY = "daily_monthly" @@ -94,10 +98,23 @@ def split_schedules(apps, schema_editor): # Payload only. `enabled` and `pg_owned` say which scheduler fires this # row and belong to converge_pg_scheduler; rewriting them here can leave # an adopted row with no firer. Its cadence does not change. - PeriodicTask.objects.filter(name=spec["name"]).update( + # + # The counts are checked rather than discarded: a bulk update matching no + # row reports success having changed nothing, leaving the old row on + # kwargs="{}" — which defaults to every tier every 15 minutes — while the + # new hourly row also fires. Strictly more load than before, silently. + beat_updated = PeriodicTask.objects.filter(name=spec["name"]).update( kwargs=json.dumps(kwargs), description=spec["description"] ) - PgPeriodicTask.objects.filter(name=spec["name"]).update(task_kwargs=kwargs) + pg_updated = PgPeriodicTask.objects.filter(name=spec["name"]).update( + task_kwargs=kwargs + ) + if not beat_updated or not pg_updated: + raise RuntimeError( + f"{spec['name']}: expected a row on both schedulers to split, " + f"found beat={beat_updated} pg={pg_updated}. Apply 0002 and 0004 " + "first, or restore the row before re-running." + ) continue schedule, _ = CrontabSchedule.objects.get_or_create( @@ -148,14 +165,22 @@ def merge_schedules(apps, schema_editor): PeriodicTask.objects.filter(name__in=added).delete() PgPeriodicTask.objects.filter(name__in=added).delete() - PeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update( + beat_restored = PeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update( kwargs="{}", description=( "Aggregate metrics from source tables (Usage, PageUsage, etc.) " "into hourly, daily, and monthly metrics tables" ), ) - PgPeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update(task_kwargs={}) + pg_restored = PgPeriodicTask.objects.filter(name=EXISTING_AGGREGATE_ROW).update( + task_kwargs={} + ) + if not beat_restored or not pg_restored: + raise RuntimeError( + f"{EXISTING_AGGREGATE_ROW}: expected a row on both schedulers to restore, " + f"found beat={beat_restored} pg={pg_restored}. The rollback would leave " + "the daily and monthly tiers with no schedule." + ) _bump_beat_change_tracker(apps) diff --git a/backend/dashboard_metrics/tasks.py b/backend/dashboard_metrics/tasks.py index 1c8a447ffb..48a742357f 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -217,20 +217,75 @@ class AggregationTier(StrEnum): # Keyed per tier: the two schedules collide hourly and must not block each other. +# Daily-tier source lookback. The reconciliation schedule 0005 declares dispatches a +# wider one, so this signature has to accept it — see workers/scheduler and +# internal_views for the other two legs of the same path. +DASHBOARD_SOURCE_WINDOW_DAYS = 7 +MAX_SOURCE_WINDOW_DAYS = 90 + AGGREGATION_LOCK_KEY_PREFIX = "dashboard_metrics:aggregation_lock" AGGREGATION_LOCK_TIMEOUT = 900 # 15 minutes (matches the fastest task schedule) +# Which granularities each tier writes. One table rather than a predicate per +# granularity: add a member without an entry here and _tiers_written raises on the +# first run, instead of the run acquiring its lock, iterating every org, writing +# nothing and returning success. +_TIER_WRITES: dict[AggregationTier, frozenset[str]] = { + AggregationTier.HOURLY: frozenset({AggregationTier.HOURLY.value}), + AggregationTier.DAILY_MONTHLY: frozenset({AggregationTier.DAILY_MONTHLY.value}), + AggregationTier.ALL: frozenset( + {AggregationTier.HOURLY.value, AggregationTier.DAILY_MONTHLY.value} + ), +} + + +def _tiers_written(tier: AggregationTier) -> frozenset[str]: + """The granularities one tier writes. Unhandled members raise rather than no-op.""" + try: + return _TIER_WRITES[tier] + except KeyError: + raise AssertionError(f"Unhandled AggregationTier: {tier!r}") from None + + def _writes_hourly(tier: AggregationTier) -> bool: - return tier in (AggregationTier.HOURLY, AggregationTier.ALL) + return AggregationTier.HOURLY.value in _tiers_written(tier) def _writes_daily_monthly(tier: AggregationTier) -> bool: - return tier in (AggregationTier.DAILY_MONTHLY, AggregationTier.ALL) + return AggregationTier.DAILY_MONTHLY.value in _tiers_written(tier) + + +def _aggregation_lock_keys(tier: AggregationTier, source_window_days: int) -> list[str]: + """One key per granularity written, namespaced by source window. + + Per granularity, not per enum member: keying on the label alone gives ALL a third + key that excludes nothing, so an ALL run and the scheduled hourly run would write + EventMetricsHourly concurrently. Taking one key per granularity restores exclusion + exactly where writes collide, and the two scheduled tiers still never block. + + Per window because a wider window is a different job. The reconciliation pass runs + once a day on a fixed crontab against a drifting 15-minute interval; on a shared + key it would lose the race, return skipped=True and never be retried — and it is + the only thing that repairs the narrowed window. Both are idempotent upserts, so + that once-a-day overlap costs duplicated work at worst. + """ + return [ + f"{AGGREGATION_LOCK_KEY_PREFIX}:{source_window_days}d:{granularity}" + for granularity in sorted(_tiers_written(tier)) + ] -def _aggregation_lock_key(tier: AggregationTier) -> str: - return f"{AGGREGATION_LOCK_KEY_PREFIX}:{tier.value}" +def _acquire_aggregation_locks(lock_keys: list[str]) -> list[str]: + """Take every key or none; returns the keys taken, empty if the run must skip.""" + taken: list[str] = [] + for key in lock_keys: + if not _acquire_aggregation_lock(key): + for held in taken: + cache.delete(held) + return [] + taken.append(key) + return taken def _acquire_aggregation_lock(lock_key: str) -> bool: @@ -293,6 +348,7 @@ def _acquire_aggregation_lock(lock_key: str) -> bool: ) def aggregate_metrics_from_sources( tier: str = AggregationTier.ALL, + source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, ) -> dict[str, Any]: """Aggregate metrics from source tables into the hourly/daily/monthly tables. @@ -311,29 +367,41 @@ def aggregate_metrics_from_sources( Args: tier: An AggregationTier value. Defaults to all, so a caller that omits it writes every tier rather than none. + source_window_days: Daily-tier source lookback. The once-daily + reconciliation schedule declared by 0005 dispatches a wider one. Returns: Dict with aggregation summary for the tiers that ran Raises: - ValueError: tier is not a recognised AggregationTier + ValueError: tier is not a recognised AggregationTier, or the window is + not an integer between 1 and MAX_SOURCE_WINDOW_DAYS """ tier = AggregationTier(tier) - lock_key = _aggregation_lock_key(tier) + source_window_days = _validate_source_window(source_window_days) + lock_keys = _aggregation_lock_keys(tier, source_window_days) - if not _acquire_aggregation_lock(lock_key): - logger.info("Skipping %s aggregation — another run is in progress", tier.value) + held = _acquire_aggregation_locks(lock_keys) + if not held: + logger.warning( + "Skipping the %s aggregation over %d day(s) — another run writing the " + "same tier is in progress", + tier.value, + source_window_days, + ) return { "success": True, "skipped": True, "reason": "lock_held", "tier": tier.value, + "source_window_days": source_window_days, } try: - return _run_aggregation(tier) + return _run_aggregation(tier, source_window_days) finally: - cache.delete(lock_key) + for key in held: + cache.delete(key) def _aggregate_single_metric( @@ -578,11 +646,31 @@ def _aggregate_org( stats["errors"] += 1 -def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, Any]: +def _validate_source_window(source_window_days: int) -> int: + """Coerce and bound the window. It arrives as JSON from an editable Beat row.""" + try: + days = int(source_window_days) + except (TypeError, ValueError) as exc: + raise ValueError( + f"source_window_days must be an integer, got {source_window_days!r}" + ) from exc + if not 1 <= days <= MAX_SOURCE_WINDOW_DAYS: + raise ValueError( + f"source_window_days must be between 1 and {MAX_SOURCE_WINDOW_DAYS}, " + f"got {days}" + ) + return days + + +def _run_aggregation( + tier: AggregationTier = AggregationTier.ALL, + source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, +) -> dict[str, Any]: """Execute the actual aggregation logic. Separated from the task function to keep the lock management clean. """ + source_window_days = _validate_source_window(source_window_days) end_date = timezone.now() # Query windows for each granularity @@ -590,7 +678,7 @@ def _run_aggregation(tier: AggregationTier = AggregationTier.ALL) -> dict[str, A # - Daily: Last 7 days (ensures we capture late-arriving data) # - Monthly: Last 2 months (current + previous, ensures month transitions are captured) hourly_start = end_date - timedelta(hours=24) - daily_start = _truncate_to_day(end_date - timedelta(days=7)) + daily_start = _truncate_to_day(end_date - timedelta(days=source_window_days)) # Include previous month to handle month boundaries if end_date.month == 1: monthly_start = end_date.replace( diff --git a/backend/dashboard_metrics/tests/test_active_org_prefilter.py b/backend/dashboard_metrics/tests/test_active_org_prefilter.py index 6cb6263764..d7893c653d 100644 --- a/backend/dashboard_metrics/tests/test_active_org_prefilter.py +++ b/backend/dashboard_metrics/tests/test_active_org_prefilter.py @@ -46,6 +46,8 @@ from dashboard_metrics.tasks import AggregationTier, _run_aggregation # noqa: E402 +INDEX_NAME = "we_created_at_idx" + _ROWS = 12000 _SPAN_DAYS = 255 @@ -96,7 +98,14 @@ def _prefilter_sql(self) -> str: and "DISTINCT" in q["sql"].upper() and "created_at" in q["sql"] ] + # The run issues nine further queries against this table, several of them + # joining it and filtering created_at. Index 0 is right today only by execution + # order, which nothing here states — so require the shape to be unambiguous. assert candidates, "the aggregation issued no active-org prefilter query" + assert len(candidates) == 1, ( + f"{len(candidates)} queries match the prefilter shape; the match is no " + f"longer distinguishing:\n" + "\n\n".join(candidates) + ) return str(candidates[0]) def test_the_window_is_the_selectivity_the_index_is_for(self) -> None: @@ -111,9 +120,22 @@ def test_the_window_is_the_selectivity_the_index_is_for(self) -> None: share = cur.fetchone()[0] assert 0 < share < 0.10 - def test_the_prefilter_does_not_scan_the_executions_table(self) -> None: - """The regression the index is meant to remove.""" + def test_the_index_can_serve_the_prefilter(self) -> None: + """*Usable*, not *chosen*. + + Whether the planner picks the index on a synthetic table turns on + random_page_cost, effective_cache_size, the PG major version and how the + freshly-loaded visibility map looks — none of which this fixture pins, so + asserting the choice reds the build on a config change with no code change. + Disabling seqscan asks the question that is actually about the query: can this + shape be served from the index at all? A prefilter narrowed to lead with a + different column fails here whatever the cost model says. + """ + sql = self._prefilter_sql() with connection.cursor() as cur: - cur.execute("EXPLAIN " + self._prefilter_sql()) + cur.execute("SET LOCAL enable_seqscan = off") + cur.execute("EXPLAIN " + sql) plan = "\n".join(row[0] for row in cur.fetchall()) - assert "Seq Scan on workflow_execution" not in plan, plan + assert f"Index Scan using {INDEX_NAME}" in plan or f"Index Only Scan using {INDEX_NAME}" in plan, ( + f"expected {INDEX_NAME} to be usable for the prefilter:\n{plan}" + ) diff --git a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py index c971ce3059..5a03e212b2 100644 --- a/backend/dashboard_metrics/tests/test_aggregation_dispatch.py +++ b/backend/dashboard_metrics/tests/test_aggregation_dispatch.py @@ -67,12 +67,35 @@ def test_the_endpoint_forwards_the_tier_to_the_task(self, tier: str) -> None: def test_an_omitted_tier_leaves_the_task_default_in_place(self) -> None: """Not 'hourly', and not nothing: the task's own default is `all`, and passing - anything here would override it during the pre-0005 deploy window. + anything here would override it during the window before 0006 applies. """ status, called_with = _post({}) assert status == 200 assert called_with == {} + @pytest.mark.parametrize("body", [{}, {"tier": None}, "not-a-dict"]) + def test_an_absent_tier_leaves_the_task_default_in_place(self, body) -> None: + """An explicit null and a non-dict body both mean "omitted", not "no tiers".""" + status, called_with = _post(body) + assert status == 200 + assert called_with == {} + + def test_the_source_window_reaches_the_task(self) -> None: + """0005's reconciliation row dispatches this against the same task path.""" + status, called_with = _post({"source_window_days": 7}) + assert status == 200 + assert called_with == {"source_window_days": 7} + + def test_both_kwargs_survive_together(self) -> None: + status, called_with = _post({"tier": "hourly", "source_window_days": 7}) + assert status == 200 + assert called_with == {"tier": "hourly", "source_window_days": 7} + + def test_a_non_integer_window_is_rejected(self) -> None: + status, called_with = _post({"source_window_days": "seven"}) + assert status == 400 + assert called_with is None + def test_an_unrecognised_tier_is_rejected_rather_than_ignored(self) -> None: """A silent no-op would look like a successful run that wrote nothing. diff --git a/backend/dashboard_metrics/tests/test_aggregation_tier.py b/backend/dashboard_metrics/tests/test_aggregation_tier.py index 4cb3cfbc5e..f868e1395d 100644 --- a/backend/dashboard_metrics/tests/test_aggregation_tier.py +++ b/backend/dashboard_metrics/tests/test_aggregation_tier.py @@ -12,6 +12,7 @@ import inspect import os +import time import django import pytest @@ -21,15 +22,26 @@ if not apps.ready: django.setup() +from django.core.cache import cache # noqa: E402 + from dashboard_metrics.tasks import ( # noqa: E402 + AGGREGATION_LOCK_TIMEOUT, + DASHBOARD_SOURCE_WINDOW_DAYS, AggregationTier, - _aggregation_lock_key, + _acquire_aggregation_lock, + _acquire_aggregation_locks, + _aggregation_lock_keys, + _tiers_written, _writes_daily_monthly, _writes_hourly, aggregate_metrics_from_sources, ) +def _keys(tier, window: int = DASHBOARD_SOURCE_WINDOW_DAYS) -> list[str]: + return _aggregation_lock_keys(tier, window) + + class TestWhichTiersEachRunWrites: @pytest.mark.parametrize( "tier,hourly,daily_monthly", @@ -91,15 +103,81 @@ def test_an_unrecognised_tier_raises(self) -> None: AggregationTier("houry") -class TestTheLockIsPerTier: - def test_every_tier_gets_its_own_key(self) -> None: - """The two schedules collide at the top of every hour. One global key and - whichever fired first would hold it while the other returned lock_held — - so the slower tier could be starved indefinitely. - """ - keys = {_aggregation_lock_key(t) for t in AggregationTier} - assert len(keys) == len(list(AggregationTier)) +class TestTheTierTableIsExhaustive: + """A member with no entry must raise, not write nothing and report success.""" - def test_the_key_names_the_tier(self) -> None: + def test_every_declared_tier_has_an_entry(self) -> None: for tier in AggregationTier: - assert _aggregation_lock_key(tier).endswith(tier.value) + assert _tiers_written(tier) + + def test_an_unhandled_member_raises_rather_than_writing_nothing(self) -> None: + # Stands in for a member added to the enum without a _TIER_WRITES entry. + ghost = type("_Ghost", (), {"value": "weekly"})() + with pytest.raises(AssertionError, match="Unhandled AggregationTier"): + _tiers_written(ghost) + + +class TestTheLockCoversWhatIsWritten: + """Keyed by granularity written, not by enum member. + + Keying on the label alone gives ALL a third key that excludes nothing, so an ALL + run and the scheduled hourly run write EventMetricsHourly concurrently. These + exercise the lock rather than its key string: a version of + _acquire_aggregation_lock that ignored its argument would pass a key-shape test. + """ + + @pytest.fixture(autouse=True) + def _clear(self): + cache.clear() + yield + cache.clear() + + def test_the_two_scheduled_tiers_never_block_each_other(self) -> None: + assert _acquire_aggregation_locks(_keys(AggregationTier.HOURLY)) + assert _acquire_aggregation_locks(_keys(AggregationTier.DAILY_MONTHLY)) + + def test_a_tier_blocks_itself(self) -> None: + assert _acquire_aggregation_locks(_keys(AggregationTier.HOURLY)) + assert not _acquire_aggregation_locks(_keys(AggregationTier.HOURLY)) + + def test_all_is_blocked_by_either_half(self) -> None: + """The exclusion a per-member key silently dropped.""" + assert _acquire_aggregation_locks(_keys(AggregationTier.HOURLY)) + assert not _acquire_aggregation_locks(_keys(AggregationTier.ALL)) + + def test_a_blocked_run_releases_whatever_it_took(self) -> None: + """ALL takes hourly first; failing on daily_monthly must not strand hourly.""" + assert _acquire_aggregation_locks(_keys(AggregationTier.DAILY_MONTHLY)) + assert not _acquire_aggregation_locks(_keys(AggregationTier.ALL)) + assert _acquire_aggregation_locks(_keys(AggregationTier.HOURLY)) + + def test_a_wider_window_is_a_different_job(self) -> None: + """The reconciliation pass is never retried, so it must not be starved by the + 15-minute schedule it races against.""" + assert _acquire_aggregation_locks(_keys(AggregationTier.HOURLY, 2)) + assert _acquire_aggregation_locks(_keys(AggregationTier.ALL, 7)) + + +class TestTheLockSelfHeals: + """Both reclaim branches, neither of which was executed by any test.""" + + @pytest.fixture(autouse=True) + def _clear(self): + cache.clear() + yield + cache.clear() + + def test_a_lock_older_than_the_timeout_is_reclaimed(self) -> None: + key = _keys(AggregationTier.HOURLY)[0] + cache.set(key, str(time.time() - AGGREGATION_LOCK_TIMEOUT - 1), 3600) + assert _acquire_aggregation_lock(key) + + def test_a_fresh_lock_is_not_reclaimed(self) -> None: + key = _keys(AggregationTier.HOURLY)[0] + cache.set(key, str(time.time()), 3600) + assert not _acquire_aggregation_lock(key) + + def test_a_corrupted_lock_value_is_reclaimed(self) -> None: + key = _keys(AggregationTier.HOURLY)[0] + cache.set(key, "running", 3600) + assert _acquire_aggregation_lock(key) diff --git a/backend/dashboard_metrics/tests/test_migration_graph.py b/backend/dashboard_metrics/tests/test_migration_graph.py new file mode 100644 index 0000000000..f89c07708f --- /dev/null +++ b/backend/dashboard_metrics/tests/test_migration_graph.py @@ -0,0 +1,53 @@ +"""Guard: the migration graph builds (UN-3974). + +Django builds the **entire** graph before executing anything, so one migration +depending on a node that does not exist aborts `migrate`, `makemigrations` and +`showmigrations` for every app in the project — the deploy's migrate step fails, not +just this app's. + +Nothing else catches it. The backend suite runs with `--no-migrations`, so test-DB +creation never builds the graph, and every migration test in this app reaches its +module through `importlib.import_module`, which resolves a file path rather than a +graph node. GitHub also reports a stacked branch as mergeable, because a missing +dependency is not a textual conflict. + +DB-free: building the graph reads the migration files, not the database. +""" + +from __future__ import annotations + +import os + +import django +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from django.db.migrations.loader import MigrationLoader # noqa: E402 +from django.test import SimpleTestCase # noqa: E402 + + +class MigrationGraphTests(SimpleTestCase): + def test_the_graph_builds(self) -> None: + """A dependency on an absent migration raises NodeNotFoundError here.""" + loader = MigrationLoader(None, ignore_no_migrations=True) + loader.build_graph() + self.assertTrue(loader.graph.nodes, "no migrations loaded — the guard is inert") + + def test_every_app_has_exactly_one_leaf(self) -> None: + """Two leaves in one app block `migrate` for every app, not just that one. + + This is what a merge of two branches that each added a migration produces, and + it is invisible until deploy for the same `--no-migrations` reason. + """ + loader = MigrationLoader(None, ignore_no_migrations=True) + loader.build_graph() + + leaves: dict[str, list[str]] = {} + for app_label, name in loader.graph.leaf_nodes(): + leaves.setdefault(app_label, []).append(name) + + conflicts = {app: names for app, names in leaves.items() if len(names) > 1} + self.assertEqual(conflicts, {}, f"apps with multiple leaf migrations: {conflicts}") diff --git a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py index 728a1def82..47d9fedea3 100644 --- a/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py +++ b/backend/dashboard_metrics/tests/test_pg_periodic_task_declarations.py @@ -19,10 +19,19 @@ import importlib import json +import os from types import SimpleNamespace from typing import Any +import django import pytest +from django.apps import apps + +os.environ.setdefault("DJANGO_SETTINGS_MODULE", "backend.settings.test") +if not apps.ready: + django.setup() + +from dashboard_metrics.tasks import AggregationTier # noqa: E402 _BEAT_MIGRATION = "dashboard_metrics.migrations.0002_setup_periodic_tasks" _PG_MIGRATION = "dashboard_metrics.migrations.0004_pg_periodic_tasks" @@ -126,7 +135,7 @@ def test_no_spec_presets_a_run_time(self, pg_specs): class _SplitRecorder: - """Captures what 0005 does to one scheduler table, keeping creates and updates apart. + """Captures what 0006 does to one scheduler table, keeping creates and updates apart. The distinction is the point: creating a row writes every default, updating one writes only the named fields. Conflating them is exactly the bug this guards. @@ -135,12 +144,20 @@ class _SplitRecorder: def __init__(self, existing: Any = None) -> None: self.created: dict[str, dict[str, Any]] = {} self.updated: dict[str, dict[str, Any]] = {} + self.deleted: list[str] = [] self.bumps = 0 + # How many rows a filtered update matches. 0 models the row being absent, + # which is the case the migration now refuses to report as success. + self.rows_present: int | None = None self._existing = existing self._filtered_on: str = "" + self._filtered_in: list[str] = [] - def filter(self, name: str = "", **_kw: Any) -> _SplitRecorder: + def filter( + self, name: str = "", name__in: list[str] | None = None, **_kw: Any + ) -> _SplitRecorder: self._filtered_on = name + self._filtered_in = list(name__in or []) return self def first(self) -> Any: @@ -148,7 +165,7 @@ def first(self) -> Any: def update(self, **kwargs: Any) -> int: self.updated[self._filtered_on] = kwargs - return 1 + return 1 if self.rows_present is None else self.rows_present def update_or_create( self, name: str = "", defaults: dict[str, Any] | None = None, **_kw: Any @@ -163,7 +180,8 @@ def get_or_create(self, **kwargs: Any) -> tuple[dict[str, Any], bool]: return kwargs, True def delete(self) -> tuple[int, dict[str, Any]]: - return (0, {}) + self.deleted.extend(self._filtered_in or [self._filtered_on]) + return (len(self.deleted), {}) def _run_split(beat_row: Any = None, pg_row: Any = None) -> dict[str, _SplitRecorder]: @@ -305,6 +323,67 @@ def get_model(self, _app: str, model: str) -> type: assert tracker.bumps == 1 +class TestTheFrozenLiteralsMatchTheEnumToday: + """These are wire values, so a rename has to fail loudly rather than pass. + + The migration cannot import the enum, and it never re-runs — so renaming an + AggregationTier value and "keeping this in step" leaves live rows carrying the old + string while every test goes green. Comparing the two here turns that into a + failure at the moment of the rename. + """ + + def test_the_declared_tiers_are_exactly_the_schedulable_ones(self) -> None: + mod = importlib.import_module(_SPLIT_MIGRATION) + declared = {spec["tier"] for spec in mod.AGGREGATION_SCHEDULES} + # ALL is the signature default and the pre-migration row's meaning; no + # schedule row ever carries it. + schedulable = {t.value for t in AggregationTier} - {AggregationTier.ALL.value} + assert declared == schedulable + + +class TestTheRollbackRestoresOneRow: + """merge_schedules is this PR's stated safety story and had no coverage at all.""" + + def _run_merge(self, rows_present: int | None = None): + mod = importlib.import_module(_SPLIT_MIGRATION) + beat, pg, tracker = _SplitRecorder(), _SplitRecorder(), _SplitRecorder() + beat.rows_present = rows_present + pg.rows_present = rows_present + + class _Apps: + def get_model(self, _app: str, model: str) -> type: + table = { + "PeriodicTask": beat, + "PgPeriodicTask": pg, + "PeriodicTasks": tracker, + }.get(model, _SplitRecorder()) + return type("M", (), {"objects": table}) + + mod.merge_schedules(_Apps(), None) + return {"beat": beat, "pg": pg, "tracker": tracker} + + def test_the_added_row_is_deleted_from_both_tables(self) -> None: + merged = self._run_merge() + for table in ("beat", "pg"): + assert _NEW_ROW in merged[table].deleted + + def test_the_existing_row_gets_its_pre_split_payload_back(self) -> None: + merged = self._run_merge() + assert merged["beat"].updated[_EXISTING_ROW]["kwargs"] == "{}" + assert merged["pg"].updated[_EXISTING_ROW]["task_kwargs"] == {} + + def test_the_rollback_leaves_ownership_alone_like_the_forward_direction(self) -> None: + merged = self._run_merge() + assert "enabled" not in merged["beat"].updated[_EXISTING_ROW] + assert set(merged["pg"].updated[_EXISTING_ROW]) == {"task_kwargs"} + + def test_a_rollback_that_restores_nothing_raises(self) -> None: + """A bulk update matching no row would otherwise report a clean rollback while + leaving the daily and monthly tiers with no schedule at all.""" + with pytest.raises(RuntimeError, match=_EXISTING_ROW): + self._run_merge(rows_present=0) + + class TestTheRewriteLeavesSchedulerOwnershipAlone: """The existing row may already be owned by the PG scheduler, with its Beat twin disabled by converge_pg_scheduler. Rewriting `pg_owned` or `enabled` here would diff --git a/backend/dashboard_metrics/tests/test_tier_split_equivalence.py b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py index 13599deb2c..3f6fbe54b2 100644 --- a/backend/dashboard_metrics/tests/test_tier_split_equivalence.py +++ b/backend/dashboard_metrics/tests/test_tier_split_equivalence.py @@ -7,8 +7,11 @@ Two properties, and both matter: -- the `hourly` schedule reproduces what the single pre-split run wrote to - EventMetricsHourly, exactly — that is the "unchanged" half +- the `hourly` schedule reproduces what an all-tiers run writes to EventMetricsHourly, + exactly — that is the "unchanged" half. Note this is a **partition** property of the + post-change code, not a comparison against the pre-split implementation, which this + branch does not have: `test_the_hourly_tier_holds_the_figures_the_fixture_implies` + is what anchors it to an absolute number - `hourly` and `daily_monthly` together reproduce every row the pre-split run wrote to any table — that is the "nothing is lost" half, which the AC assumes rather than states @@ -20,6 +23,7 @@ import uuid from datetime import timedelta from typing import Any +from unittest.mock import patch import os @@ -41,7 +45,11 @@ EventMetricsHourly, EventMetricsMonthly, ) -from dashboard_metrics.tasks import AggregationTier, _run_aggregation # noqa: E402 +from dashboard_metrics.tasks import ( # noqa: E402 + AggregationTier, + _run_aggregation, + _truncate_to_month, +) # (model, the column naming its period) — the period field differs per tier. _TIERS = [ @@ -60,16 +68,21 @@ def setUp(self) -> None: self.workflow = Workflow.objects.create( workflow_name="tier-split-wf", organization=self.org ) - now = timezone.now() + self.now = now = timezone.now() # One row per window the aggregation reads — last 24h for the hourly tier, last # 7 days for daily, inside the previous month for monthly. The two recent ones # also make the org visible to the active-org prefilter, without which nothing # runs at all. + # + # The previous-month row is derived from the month boundary, not a fixed + # "25 days ago": for the last few days of any month that lands in the *current* + # month and the cross-boundary coverage silently disappears. + last_month_day = _truncate_to_month(now) - timedelta(days=1) windows = [ now - timedelta(hours=2), now - timedelta(hours=5), now - timedelta(days=3), - now - timedelta(days=25), + last_month_day, ] executions = self._add_executions(windows) # Both aggregation paths have to be exercised: the per-metric queries go through @@ -136,8 +149,15 @@ def _clear(self) -> None: model._base_manager.all().delete() def _run(self, tier: AggregationTier) -> dict[str, set[tuple[Any, ...]]]: + """One clock for every run in a test. + + _run_aggregation reads timezone.now() itself, so three unpatched invocations + compute three different window starts — and a run straddling an hour or a month + boundary would fail on a non-regression. + """ self._clear() - _run_aggregation(tier) + with patch("dashboard_metrics.tasks.timezone.now", return_value=self.now): + _run_aggregation(tier) return self._snapshot() def test_the_pre_split_run_writes_all_three_tiers(self) -> None: @@ -160,6 +180,24 @@ def test_the_fixture_exercises_both_aggregation_paths(self) -> None: assert "documents_processed" in names, f"{table}: no per-metric figure" assert "llm_calls" in names, f"{table}: no combined-LLM figure" + def test_the_hourly_tier_holds_the_figures_the_fixture_implies(self) -> None: + """An absolute expectation, not a comparison of the code against itself. + + Every other assertion in this file runs the same post-change function twice, so + a regression in the shared path — the window arithmetic, the org_identifier + handoff, an upsert that writes zeroes — moves both sides equally and stays + green. This one names a number the fixture determines. + """ + self._run(AggregationTier.HOURLY) + rows = EventMetricsHourly._base_manager.filter( + metric_name="documents_processed" + ) + assert sum(row.metric_value for row in rows) == 2, ( + "exactly the -2h and -5h file executions fall inside the 24h window; " + "the -3d and previous-month ones must not" + ) + assert {row.metric_value for row in rows} != {0} + def test_hourly_reproduces_the_pre_split_hourly_figures(self) -> None: """The "figures unchanged" half of AC-2, row for row rather than in aggregate.""" before = self._run(AggregationTier.ALL)["EventMetricsHourly"] diff --git a/workers/scheduler/dashboard_metrics_tasks.py b/workers/scheduler/dashboard_metrics_tasks.py index 18e002c6e3..fd3ed56300 100644 --- a/workers/scheduler/dashboard_metrics_tasks.py +++ b/workers/scheduler/dashboard_metrics_tasks.py @@ -105,13 +105,23 @@ def _log_if_skipped(name: str, result: dict[str, Any]) -> None: @worker_task(name="dashboard_metrics.aggregate_from_sources") -def dashboard_metrics_aggregate(tier: str | None = None) -> dict[str, Any]: +def dashboard_metrics_aggregate( + tier: str | None = None, source_window_days: int | None = None +) -> dict[str, Any]: """Aggregate source tables into the hourly/daily/monthly metrics tables. - ``tier`` comes from the schedule row's kwargs and selects which tiers to write; - omitted means all of them. + Both kwargs come from the schedule row and both are optional: ``tier`` selects + which tiers to write, ``source_window_days`` widens the daily lookback for the + reconciliation pass. Omitting either applies the backend task's own default. """ - body = {"tier": tier} if tier is not None else None + body = { + key: value + for key, value in ( + ("tier", tier), + ("source_window_days", source_window_days), + ) + if value is not None + } or None result = _call_internal(_AGGREGATE_PATH, body=body) _log_if_skipped("dashboard_metrics.aggregate_from_sources", result) return result diff --git a/workers/tests/test_dashboard_metrics_tasks.py b/workers/tests/test_dashboard_metrics_tasks.py index 055a7725d5..c795dc9305 100644 --- a/workers/tests/test_dashboard_metrics_tasks.py +++ b/workers/tests/test_dashboard_metrics_tasks.py @@ -8,6 +8,7 @@ from __future__ import annotations +import inspect import sys from pathlib import Path from unittest.mock import MagicMock, patch @@ -82,7 +83,7 @@ def test_aggregate_forwards_the_tier_from_the_schedule_row(self, tier): assert call.call_args.kwargs["body"] == {"tier": tier} def test_aggregate_omits_the_body_when_no_tier_is_given(self): - # Pre-0005 rows carry no tier kwarg; the backend's default then applies, which + # Rows written before 0006 carry no tier kwarg; the backend default then applies, # is every tier rather than none. with patch.object(dmt, "_call_internal", return_value={"success": True}) as call: dmt.dashboard_metrics_aggregate() @@ -119,6 +120,40 @@ def test_lock_held_result_is_surfaced_not_swallowed(self, caplog): assert result["skipped"] is True +class TestTheReconciliationKwargSurvives: + """0005 declares a row against this same task path carrying source_window_days. + + The PG scheduler copies task_kwargs verbatim into the payload, so a proxy that + does not accept it raises TypeError per tick — not covered by autoretry_for, and + dropped at MAX_ATTEMPTS=1. The gap-repair pass simply never runs. + """ + + _DECLARED = [{}, {"tier": "hourly"}, {"source_window_days": 7}] + + @pytest.mark.parametrize("kwargs", _DECLARED) + def test_every_scheduled_kwarg_set_binds(self, kwargs) -> None: + inspect.signature(dmt.dashboard_metrics_aggregate).bind(**kwargs) + + def test_the_source_window_reaches_the_endpoint(self) -> None: + with patch.object(dmt, "_call_internal", return_value={}) as call: + dmt.dashboard_metrics_aggregate(source_window_days=7) + assert call.call_args.kwargs["body"] == {"source_window_days": 7} + + def test_both_kwargs_travel_together(self) -> None: + with patch.object(dmt, "_call_internal", return_value={}) as call: + dmt.dashboard_metrics_aggregate(tier="hourly", source_window_days=7) + assert call.call_args.kwargs["body"] == { + "tier": "hourly", + "source_window_days": 7, + } + + def test_omitting_both_sends_no_body(self) -> None: + # The backend then applies its own defaults rather than ones invented here. + with patch.object(dmt, "_call_internal", return_value={}) as call: + dmt.dashboard_metrics_aggregate() + assert call.call_args.kwargs["body"] is None + + class TestInternalCall: def _response(self, status_code=200, payload=None): r = MagicMock()