feat: Set feature export response on initial API request (#3126)

api/features/import_export/constants.py

@@ -13,7 +13,7 @@
 PROCESSING = "PROCESSING"
 FAILED = "FAILED"
 
-FEATURE_IMPORT_STATUSES = (
+FEATURE_EXPORT_STATUSES = FEATURE_IMPORT_STATUSES = (
     (SUCCESS, "Success"),
     (PROCESSING, "Processing"),
     (FAILED, "Failed"),

api/features/import_export/migrations/0002_status_and_data_featureexport.py

@@ -0,0 +1,23 @@
+# Generated by Django 3.2.23 on 2023-12-08 16:45
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ('features_import_export', '0001_initial'),
+    ]
+
+    operations = [
+        migrations.AddField(
+            model_name='featureexport',
+            name='status',
+            field=models.CharField(choices=[('SUCCESS', 'Success'), ('PROCESSING', 'Processing'), ('FAILED', 'Failed')], default='PROCESSING', max_length=50),
+        ),
+        migrations.AlterField(
+            model_name='featureexport',
+            name='data',
+            field=models.CharField(blank=True, max_length=1000000, null=True),
+        ),
+    ]

api/features/import_export/models.py

@@ -1,6 +1,7 @@
 from django.db import models
 
 from features.import_export.constants import (
+    FEATURE_EXPORT_STATUSES,
     FEATURE_IMPORT_STATUSES,
     FEATURE_IMPORT_STRATEGIES,
     MAX_FEATURE_EXPORT_SIZE,
@@ -25,9 +26,21 @@ class FeatureExport(models.Model):
         swappable=True,
     )
 
+    status = models.CharField(
+        choices=FEATURE_EXPORT_STATUSES,
+        max_length=50,
+        blank=False,
+        null=False,
+        default=PROCESSING,
+    )
+
     # This is a JSON string of data used for file download
-    # once the task has completed assembly.
-    data = models.CharField(max_length=MAX_FEATURE_EXPORT_SIZE)
+    # once the task has completed assembly. It is null on upload.
+    data = models.CharField(
+        max_length=MAX_FEATURE_EXPORT_SIZE,
+        blank=True,
+        null=True,
+    )
     created_at = models.DateTimeField(auto_now_add=True)
 
 

api/features/import_export/serializers.py

@@ -2,7 +2,12 @@
 from rest_framework import serializers
 from rest_framework.exceptions import ValidationError
 
-from .constants import MAX_FEATURE_IMPORT_SIZE, OVERWRITE_DESTRUCTIVE, SKIP
+from .constants import (
+    MAX_FEATURE_IMPORT_SIZE,
+    OVERWRITE_DESTRUCTIVE,
+    PROCESSING,
+    SKIP,
+)
 from .models import FeatureExport, FeatureImport
 from .tasks import (
     export_features_for_environment,
@@ -14,14 +19,21 @@ class CreateFeatureExportSerializer(serializers.Serializer):
     environment_id = serializers.IntegerField(required=True)
     tag_ids = serializers.ListField(child=serializers.IntegerField())
 
-    def save(self) -> None:
+    def save(self) -> FeatureExport:
+        feature_export = FeatureExport.objects.create(
+            environment_id=self.validated_data["environment_id"],
+            status=PROCESSING,
+        )
+
         export_features_for_environment.delay(
             kwargs={
-                "environment_id": self.validated_data["environment_id"],
+                "feature_export_id": feature_export.id,
                 "tag_ids": self.validated_data["tag_ids"],
             }
         )
 
+        return feature_export
+
 
 class FeatureExportSerializer(serializers.ModelSerializer):
     name = serializers.SerializerMethodField()
@@ -32,6 +44,7 @@ class Meta:
             "id",
             "name",
             "environment_id",
+            "status",
             "created_at",
         )
 

api/features/import_export/tasks.py

@@ -16,7 +16,7 @@
     register_task_handler,
 )
 
-from .constants import OVERWRITE_DESTRUCTIVE, SKIP
+from .constants import FAILED, OVERWRITE_DESTRUCTIVE, SKIP, SUCCESS
 from .models import FeatureExport, FeatureImport
 
 
@@ -29,10 +29,12 @@ def clear_stale_feature_imports_and_exports() -> None:
     FeatureImport.objects.filter(created_at__lt=two_weeks_ago).delete()
 
 
-@register_task_handler()
-def export_features_for_environment(
-    environment_id: int, tag_ids: Optional[list[int]] = None
+def _export_features_for_environment(
+    feature_export: FeatureExport, tag_ids: Optional[list[int]]
 ) -> None:
+    """
+    Caller for the export_features_for_environment to handle fails.
+    """
     additional_filters = Q(
         identity__isnull=True,
         feature_segment__isnull=True,
@@ -43,7 +45,7 @@ def export_features_for_environment(
     if tag_ids:
         additional_filters &= Q(feature__tags__in=tag_ids)
 
-    environment = Environment.objects.get(id=environment_id)
+    environment = feature_export.environment
     feature_states = get_environment_flags_list(
         environment=environment,
         additional_filters=additional_filters,
@@ -76,9 +78,24 @@ def export_features_for_environment(
             }
         )
 
-    FeatureExport.objects.create(
-        environment_id=environment_id, data=json.dumps(payload)
-    )
+    feature_export.status = SUCCESS
+    feature_export.data = json.dumps(payload)
+    feature_export.save()
+
+
+@register_task_handler()
+def export_features_for_environment(
+    feature_export_id: int, tag_ids: Optional[list[int]] = None
+) -> None:
+    feature_export = FeatureExport.objects.get(id=feature_export_id)
+
+    try:
+        _export_features_for_environment(feature_export, tag_ids)
+        assert feature_export.status == SUCCESS
+    except Exception:
+        feature_export.status = FAILED
+        feature_export.save()
+        raise
 
 
 @register_task_handler()

api/features/import_export/views.py

@@ -27,16 +27,17 @@
 @swagger_auto_schema(
     method="POST",
     request_body=CreateFeatureExportSerializer(),
-    responses={201: CreateFeatureExportSerializer()},
+    responses={201: FeatureExportSerializer()},
 )
 @api_view(["POST"])
 @permission_classes([CreateFeatureExportPermissions])
 def create_feature_export(request: Request) -> Response:
     serializer = CreateFeatureExportSerializer(data=request.data)
     serializer.is_valid(raise_exception=True)
-    serializer.save()
+    feature_export = serializer.save()
+    response_serializer = FeatureExportSerializer(feature_export)
 
-    return Response(serializer.validated_data, status=201)
+    return Response(response_serializer.data, status=201)
 
 
 @swagger_auto_schema(

api/tests/unit/features/import_export/test_unit_features_import_export_tasks.py

@@ -7,7 +7,11 @@
 from environments.identities.models import Identity
 from environments.models import Environment
 from features.feature_types import MULTIVARIATE, STANDARD
-from features.import_export.constants import OVERWRITE_DESTRUCTIVE, SKIP
+from features.import_export.constants import (
+    OVERWRITE_DESTRUCTIVE,
+    PROCESSING,
+    SKIP,
+)
 from features.import_export.models import FeatureExport, FeatureImport
 from features.import_export.tasks import (
     clear_stale_feature_imports_and_exports,
@@ -101,12 +105,15 @@ def test_export_and_import_features_for_environment_with_skip(
         initial_value="keepme",
     )
 
-    # When
-    export_features_for_environment(environment.id)
-
-    feature_export = FeatureExport.objects.get(
+    feature_export = FeatureExport.objects.create(
         environment=environment,
+        status=PROCESSING,
     )
+
+    # When
+    export_features_for_environment(feature_export.id)
+
+    feature_export.refresh_from_db()
     assert len(feature_export.data) > 200
 
     feature_import = FeatureImport.objects.create(
@@ -229,13 +236,15 @@ def test_export_and_import_features_for_environment_with_overwrite_destructive(
         project=project2,
         initial_value="keepme",
     )
+    feature_export = FeatureExport.objects.create(
+        environment=environment,
+        status=PROCESSING,
+    )
 
     # When
-    export_features_for_environment(environment.id, [design_tag.id])
+    export_features_for_environment(feature_export.id, [design_tag.id])
 
-    feature_export = FeatureExport.objects.get(
-        environment=environment,
-    )
+    feature_export.refresh_from_db()
     assert len(feature_export.data) > 200
 
     feature_import = FeatureImport.objects.create(

api/tests/unit/features/import_export/test_unit_features_import_export_views.py

@@ -1,13 +1,15 @@
 import json
 from typing import Callable
 
+import pytest
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.urls import reverse
+from pytest_mock import MockerFixture
 from rest_framework.test import APIClient
 
 from environments.models import Environment
 from environments.permissions.models import UserEnvironmentPermission
-from features.import_export.constants import OVERWRITE_DESTRUCTIVE
+from features.import_export.constants import OVERWRITE_DESTRUCTIVE, PROCESSING
 from features.import_export.models import FeatureExport, FeatureImport
 from projects.models import Project
 from projects.permissions import VIEW_PROJECT
@@ -207,9 +209,11 @@ def test_feature_import_unauthorized(
     assert response.status_code == 403
 
 
+@pytest.mark.freeze_time("2023-12-08T06:05:47.320000+00:00")
 def test_create_feature_export(
     admin_client: APIClient,
     environment: Environment,
+    mocker: MockerFixture,
 ) -> None:
     # Given
     tag = Tag.objects.create(
@@ -218,6 +222,9 @@ def test_create_feature_export(
         color="#228B22",
     )
 
+    task_mock = mocker.patch(
+        "features.import_export.serializers.export_features_for_environment"
+    )
     url = reverse("api-v1:features:create-feature-export")
     data = {"environment_id": environment.id, "tag_ids": [tag.id]}
     assert FeatureExport.objects.count() == 0
@@ -229,16 +236,23 @@ def test_create_feature_export(
 
     # Then
     assert response.status_code == 201
+
+    assert FeatureExport.objects.count() == 1
+    feature_export = FeatureExport.objects.all().first()
+    # Picked up later by task for processing.
+
+    assert feature_export.data is None
     assert response.data == {
+        "created_at": "2023-12-08T06:05:47.320000Z",
         "environment_id": environment.id,
-        "tag_ids": [tag.id],
+        "id": feature_export.id,
+        "name": "Test Environment | 2023-12-08 06:05 UTC",
+        "status": PROCESSING,
     }
-    assert FeatureExport.objects.count() == 1
 
-    # Created by export_features_for_environment task.
-    feature_export = FeatureExport.objects.all().first()
-    assert feature_export.data
-    assert feature_export.environment == environment
+    task_mock.delay.assert_called_once_with(
+        kwargs={"feature_export_id": feature_export.id, "tag_ids": [tag.id]},
+    )
 
 
 def test_create_feature_export_unauthorized(