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..09667e867f --- /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.""" + 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, _ = crontab_model.objects.get_or_create( + minute="0", + hour="4", + day_of_week="*", + day_of_month="*", + month_of_year="*", + 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" + ), + }, + ) + + +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() + + +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..c0b1b1a9b8 100644 --- a/backend/dashboard_metrics/tasks.py +++ b/backend/dashboard_metrics/tasks.py @@ -8,12 +8,15 @@ 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 import transaction +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 +36,17 @@ DASHBOARD_HOURLY_METRICS_RETENTION_DAYS = 30 DASHBOARD_DAILY_METRICS_RETENTION_DAYS = 365 +# Daily-tier source lookback, sized against the worst observed +# created_at -> terminal-status lag. +DASHBOARD_SOURCE_WINDOW_DAYS = 2 + +# Wider lookback for the once-daily reconciliation pass. +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. +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,42 +179,83 @@ 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 _delete_orphan_monthly(month_start: date, fresh_keys: set[tuple]) -> int: + """Drop monthly rows from month_start that the rollup no longer produces. - Uses _base_manager to bypass DefaultOrganizationManagerMixin. + 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"] + 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 - Args: - aggregations: Dict keyed by (org_id, month_str, metric_name, project, tag) + deleted, _ = EventMetricsMonthly._base_manager.filter(pk__in=stale_pks).delete() + return deleted - Returns: - Number of rows upserted + +def _rollup_monthly_from_daily(month_start: date) -> 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. """ - 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 + ] + # An empty daily tier means the source is gone, not that every month is + # zero — leave 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) @@ -260,22 +315,18 @@ 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. - - 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. +def aggregate_metrics_from_sources( + source_window_days: int = DASHBOARD_SOURCE_WINDOW_DAYS, +) -> dict[str, Any]: + """Aggregate source tables into the hourly, daily and monthly tiers. - 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. + 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. - 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) + Args: + source_window_days: Daily-tier source lookback. The once-daily + reconciliation pass reruns this task at + DASHBOARD_RECONCILE_WINDOW_DAYS to repair gaps after downtime. Returns: Dict with aggregation summary for all three tiers @@ -285,7 +336,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,19 +348,12 @@ 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. - - 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. - """ + """Run a single metric query at hourly and daily granularity.""" extra_kwargs = extra_kwargs or {} # === HOURLY (last 24h) === @@ -324,43 +368,31 @@ 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. + Two queries covering four metrics. """ # === HOURLY (last 24h) === for row in MetricsQueryService.get_llm_metrics_combined( @@ -374,82 +406,175 @@ 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) + + +# 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 org into hourly/daily aggregates. - if day_ts >= daily_start: - key = (org_id, day_ts.date().isoformat(), metric_name, "default", "") - _upsert_agg(daily_agg, key, metric_type, value) + A failing metric is logged and counted, leaving the rest to proceed. - key = (org_id, month_key, metric_name, "default", "") - _upsert_agg(monthly_agg, key, metric_type, value) + 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 a 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 -def _run_aggregation() -> dict[str, Any]: - """Execute the actual aggregation logic. + return hourly_agg, daily_agg, errors - Separated from the task function to keep the lock management clean. - """ - end_date = timezone.now() - # 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) - 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 +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 + + 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 execution activity in the prefilter lookback.""" + 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() + ) - # 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 _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]: + """Execute the aggregation, separately from the task's lock handling.""" + end_date = timezone.now() + + # 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( + _truncate_to_month(end_date) - timedelta(days=1) + ).date() stats = { "hourly": {"upserted": 0}, @@ -460,109 +585,40 @@ def _run_aggregation() -> dict[str, Any]: } # 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 - active_org_ids = set( - WorkflowExecution.objects.filter( - created_at__gte=daily_start, - ) - .values_list("workflow__organization_id", flat=True) - .distinct() - ) - total_orgs = Organization.objects.count() + 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] = {} - 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, - 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, - ) - 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 - + _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 + 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']}, " @@ -571,19 +627,7 @@ def _run_aggregation() -> dict[str, Any]: 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 03ef136508..6cb9316d1c 100644 --- a/backend/dashboard_metrics/tests/test_tasks.py +++ b/backend/dashboard_metrics/tests/test_tasks.py @@ -1,20 +1,37 @@ """Unit tests for Dashboard Metrics Celery tasks.""" -from datetime import datetime, timedelta +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 ( EventMetricsDaily, EventMetricsHourly, + 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, + _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 +215,318 @@ 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() + + 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 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.""" + + 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) + + 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()