Worker task for test analytics export

apps/codecov-api/api/sentry/urls.py

@@ -1,9 +1,19 @@
 from django.urls import path
 
-from .views import account_link, account_unlink, test_analytics_eu
+from .views import (
+    account_link,
+    account_unlink,
+    create_ta_export,
+    get_ta_export,
+)
 
 urlpatterns = [
     path("internal/account/link/", account_link, name="account-link"),
     path("internal/account/unlink/", account_unlink, name="account-unlink"),
-    path("internal/test-analytics/eu/", test_analytics_eu, name="test-analytics-eu"),
+    path("internal/test-analytics/exports/", create_ta_export, name="create-ta-export"),
+    path(
+        "internal/test-analytics/exports/<str:task_id>/",
+        get_ta_export,
+        name="get-ta-export",
+    ),
 ]

apps/codecov-api/api/sentry/views.py

@@ -1,9 +1,8 @@
 import logging
-from itertools import groupby
 
 import sentry_sdk
+from celery.result import AsyncResult
 from django.conf import settings
-from django.views.decorators.gzip import gzip_page
 from rest_framework import serializers, status
 from rest_framework.decorators import (
     api_view,
@@ -12,17 +11,16 @@
 )
 from rest_framework.response import Response
 
-from api.public.v2.test_results.serializers import TestrunSerializer
 from codecov_auth.models import Account
 from codecov_auth.permissions import JWTAuthenticationPermission
+from services.task import TaskService, celery_app
+from shared.celery_config import export_test_analytics_data_task_name
 from shared.django_apps.codecov_auth.models import (
     GithubAppInstallation,
     Owner,
     Plan,
     Service,
 )
-from shared.django_apps.core.models import Repository
-from shared.django_apps.ta_timeseries.models import Testrun
 from shared.plan.constants import PlanName
 
 log = logging.getLogger(__name__)
@@ -250,67 +248,187 @@
     )
 
 
-class SentryTestAnalyticsEuSerializer(serializers.Serializer):
-    """Serializer for test analytics EU endpoint"""
+class CreateTestAnalyticsExportSerializer(serializers.Serializer):
+    """Serializer for create test analytics export endpoint"""
 
     integration_names = serializers.ListField(
         child=serializers.CharField(),
         help_text="The Sentry integration names",
         min_length=1,
     )
+    gcp_project_id = serializers.CharField(
+        help_text="The GCP project ID",
+        required=True,
+    )
+    destination_bucket = serializers.CharField(
+        help_text="The destination bucket",
+        required=True,
+    )
+    destination_prefix = serializers.CharField(
+        help_text="The destination prefix",
+        required=True,
+    )
 
 
-@gzip_page
 @api_view(["POST"])
 @authentication_classes([])
 @permission_classes([JWTAuthenticationPermission])
-def test_analytics_eu(request, *args, **kwargs):
-    serializer = SentryTestAnalyticsEuSerializer(data=request.data)
+def create_ta_export(request, *args, **kwargs):
+    """
+    Starts a Celery task to export test analytics data.
+
+    POST /sentry/internal/test-analytics/exports/
+
+    Returns:
+        {
+            "tasks": [
+                {
+                    "integration_name": "sample-integration-name",
+                    "task_id": "sample-celery-task-id",
+                    "status": "PENDING"
+                }
+            ],
+            "total_tasks": 1,
+            "successfully_scheduled": 1,
+            "status": "PENDING"
+        }
+    """
+
+    serializer = CreateTestAnalyticsExportSerializer(data=request.data)
     serializer.is_valid(raise_exception=True)
 
     integration_names = serializer.validated_data["integration_names"]
+    gcp_project_id = serializer.validated_data["gcp_project_id"]
+    destination_bucket = serializer.validated_data["destination_bucket"]
+    destination_prefix = serializer.validated_data["destination_prefix"]
+
+    log.info(
+        "Starting data export for the following integrations",
+        extra={
+            "integrations": integration_names,
+            "integration_count": len(integration_names),
+        },
+    )
+
+    task_service = TaskService()
+    task_results = []
 
-    # For every integration name, determine if an Owner record exist by filtering by name and service=github
-    test_runs_per_integration = {}
-    for name in integration_names:
+    for integration_name in integration_names:
         try:
-            owner = Owner.objects.get(name=name, service=Service.GITHUB)
-        except Owner.DoesNotExist:
-            log.warning(
-                f"Owner with name {name} and service {Service.GITHUB} not found"
+            result = task_service.schedule_task(
+                task_name=export_test_analytics_data_task_name,
+                kwargs={
+                    "integration_name": integration_name,
+                    "gcp_project_id": gcp_project_id,
+                    "destination_bucket": destination_bucket,
+                    "destination_prefix": destination_prefix,
+                },
+                apply_async_kwargs={},
+            )
+            task_id = result.id
+            task_status = result.status
+            task_results.append(
+                {
+                    "integration_name": integration_name,
+                    "task_id": task_id,
+                    "status": task_status,
+                }
+            )
+        except Exception as e:
+            sentry_sdk.capture_exception(
+                e,
+                tags={
+                    "event": "test_analytics_export_scheduling_failed",
+                    "integration_name": integration_name,
+                },
+            )
+            task_results.append(
+                {
+                    "integration_name": integration_name,
+                    "error": str(e),
+                    "status": "FAILED_TO_SCHEDULE",
+                }
             )
-            continue
 
-        # Only fetch name and repoid fields
-        repo_id_to_name = dict(
-            Repository.objects.filter(
-                author=owner, test_analytics_enabled=True
-            ).values_list("repoid", "name")
-        )
+    successful_tasks = [task for task in task_results if "task_id" in task]
+    failed_tasks = [task for task in task_results if "error" in task]
+
+    log.info(
+        "Completed data export scheduling for test analytics",
+        extra={
+            "successful_tasks": [
+                {
+                    "integration_name": task["integration_name"],
+                    "task_id": task["task_id"],
+                }
+                for task in successful_tasks
+            ],
+            "failed_tasks": [
+                {
+                    "integration_name": task["integration_name"],
+                    "error": task["error"],
+                    "status": task.get("status", "UNKNOWN"),
+                }
+                for task in failed_tasks
+            ],
+        },
+    )
 
-        if not repo_id_to_name:
-            test_runs_per_integration[name] = {}
-            continue
+    return Response(
+        {
+            "tasks": task_results,
+            "total_tasks": len(task_results),
+            "successfully_scheduled": len(successful_tasks),
+        },
+        status=status.HTTP_202_ACCEPTED,
+    )
 
-        # Fetch all test runs for all repositories in a single query
-        test_runs = Testrun.objects.filter(repo_id__in=repo_id_to_name.keys()).order_by(
-            "repo_id", "-timestamp"
-        )
 
-        # Group by repo_id (data is already ordered by repo_id) and serialize each group
-        test_runs_per_repository = {}
-        for repo_id, group in groupby(test_runs, key=lambda tr: tr.repo_id):
-            repo_name = repo_id_to_name[repo_id]  # Safe: we only fetch these repo_ids
-            test_runs_list = list(group)
-            test_runs_per_repository[repo_name] = TestrunSerializer(
-                test_runs_list, many=True
-            ).data
+@api_view(["GET"])
+@authentication_classes([])
+@permission_classes([JWTAuthenticationPermission])
+def get_ta_export(request, task_id, *args, **kwargs):
+    """
+    Check the status of a test analytics export task.
 
-        # Store each test_runs_per_repository in a dictionary
-        test_runs_per_integration[name] = test_runs_per_repository
+    GET /sentry/internal/test-analytics/exports/<task_id>
 
-    return Response(
+    Returns:
         {
-            "test_runs_per_integration": test_runs_per_integration,
+            "task_id": "celery-task-id-here",
+            "status": "SUCCESS|PENDING|FAILURE|RETRY|STARTED",
+            "result": {...}  # Only present if status is SUCCESS
+        }
+    """
+    log.info(
+        "Checking status for test analytics export task", extra={"task_id": task_id}
+    )
+
+    result = AsyncResult(task_id, app=celery_app)
+
+    response_data = {
+        "task_id": task_id,
+        "status": result.status,
+    }
+
+    if result.successful():
+        task_result = result.result
+        response_data["result"] = task_result
+        if isinstance(task_result, dict):
+            response_data["integration_name"] = task_result.get("integration_name")
+            if not task_result.get("successful", True):
+                response_data["task_reported_failure"] = True
+                response_data["error"] = task_result.get("error")
+    elif result.failed():
+        error_info = result.info
+        response_data["error"] = {
+            "message": str(error_info),
+            "type": type(error_info).__name__ if error_info else "Unknown",
         }
+
+    log.info(
+        "Completed status check for test analytics export task",
+        extra=response_data,
     )
+
+    return Response(response_data, status=status.HTTP_200_OK)

apps/worker/tasks/__init__.py

@@ -21,6 +21,7 @@
 from tasks.compute_component_comparison import compute_component_comparison_task
 from tasks.delete_owner import delete_owner_task
 from tasks.detect_flakes import detect_flakes_task
+from tasks.export_test_analytics_data import export_test_analytics_data_task
 from tasks.flush_repo import flush_repo
 from tasks.github_app_webhooks_check import gh_webhook_check_task
 from tasks.github_marketplace import ghm_sync_plans_task