fix: optimise segment deletion and cloning with bulk operations (#6401)

Author: gagantrivedi

api/core/dataclasses.py

@@ -0,0 +1,25 @@
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from rest_framework.request import Request
+
+    from api_keys.models import MasterAPIKey
+    from users.models import FFAdminUser
+
+
+@dataclass
+class AuthorData:
+    user: "FFAdminUser | None" = None
+    api_key: "MasterAPIKey | None" = None
+
+    @classmethod
+    def from_request(cls, request: "Request") -> "AuthorData":
+        from users.models import FFAdminUser
+
+        if type(request.user) is FFAdminUser:
+            return cls(user=request.user)
+        elif hasattr(request.user, "key"):
+            return cls(api_key=request.user.key)
+        else:
+            raise ValueError("Request user must be FFAdminUser or have an API key")

api/features/feature_states/serializers.py

@@ -1,9 +1,9 @@
 from rest_framework import serializers
 
+from core.dataclasses import AuthorData
 from environments.models import Environment
 from features.models import Feature, FeatureState
 from features.versioning.dataclasses import (
-    AuthorData,
     FlagChangeSet,
     FlagChangeSetV2,
     SegmentOverrideChangeSet,

api/features/versioning/dataclasses.py

@@ -2,35 +2,12 @@
 
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import TYPE_CHECKING
 
 from pydantic import BaseModel, computed_field
 
+from core.dataclasses import AuthorData
 from features.feature_states.models import FeatureValueType
 
-if TYPE_CHECKING:
-    from rest_framework.request import Request
-
-    from api_keys.models import MasterAPIKey
-    from users.models import FFAdminUser
-
-
-@dataclass
-class AuthorData:
-    user: FFAdminUser | None = None
-    api_key: MasterAPIKey | None = None
-
-    @classmethod
-    def from_request(cls, request: Request) -> AuthorData:
-        from users.models import FFAdminUser
-
-        if type(request.user) is FFAdminUser:
-            return cls(user=request.user)
-        elif hasattr(request.user, "key"):
-            return cls(api_key=request.user.key)
-        else:
-            raise ValueError("Request user must be FFAdminUser or have an API key")
-
 
 class Conflict(BaseModel):
     segment_id: int | None = None

api/features/versioning/versioning_service.py

@@ -5,11 +5,11 @@
 from django.utils import timezone
 from rest_framework.exceptions import NotFound
 
+from core.dataclasses import AuthorData
 from environments.models import Environment
 from features.feature_states.models import FeatureValueType
 from features.models import Feature, FeatureSegment, FeatureState, FeatureStateValue
 from features.versioning.dataclasses import (
-    AuthorData,
     FlagChangeSet,
     FlagChangeSetV2,
 )

api/segments/models.py

@@ -28,6 +28,7 @@
 from features.models import Feature
 from metadata.models import Metadata
 from projects.models import Project
+from segments.services import copy_segment_rules_and_conditions
 
 ModelT = typing.TypeVar("ModelT", bound=models.Model)
 
@@ -180,40 +181,10 @@ def clone(self, is_revision: bool = False, **extra_attrs: typing.Any) -> "Segmen
         return cloned_segment
 
     def copy_rules_and_conditions_from(self, source_segment: "Segment") -> None:
-        """
-        Recursively copy rules and conditions from another segment
-        """
-        assert transaction.get_connection().in_atomic_block, "Must run in a transaction"
-
-        # Delete existing rules
-        SegmentRule.objects.filter(segment=self).delete()
-
-        source_rules = SegmentRule.objects.filter(
-            models.Q(segment=source_segment) | models.Q(rule__segment=source_segment)
+        copy_segment_rules_and_conditions(
+            target_segment=self, source_segment=source_segment
         )
 
-        # Ensure top-level rules are cloned first (for dependencies)
-        source_rules = source_rules.order_by(models.F("rule").asc(nulls_first=True))
-
-        rule_to_cloned_rule_map: dict[SegmentRule, SegmentRule] = {}
-        for rule in source_rules:
-            cloned_rule = deepcopy(rule)
-            cloned_rule.pk = None
-            cloned_rule.uuid = uuid.uuid4()
-            cloned_rule.segment = self if rule.segment else None
-            if rule.rule in rule_to_cloned_rule_map:
-                cloned_rule.rule = rule_to_cloned_rule_map[rule.rule]
-            cloned_rule.save()
-            rule_to_cloned_rule_map[rule] = cloned_rule
-
-        source_conditions = Condition.objects.filter(rule__in=rule_to_cloned_rule_map)
-        for condition in source_conditions:
-            cloned_condition = deepcopy(condition)
-            cloned_condition.pk = None
-            cloned_condition.uuid = uuid.uuid4()
-            cloned_condition.rule = rule_to_cloned_rule_map[condition.rule]
-            cloned_condition.save()
-
     def get_create_log_message(self, history_instance) -> typing.Optional[str]:  # type: ignore[no-untyped-def]
         return SEGMENT_CREATED_MESSAGE % self.name
 

api/segments/services.py

@@ -0,0 +1,160 @@
+import typing
+import uuid
+
+from django.db import models, transaction
+from django.utils import timezone
+
+from core.dataclasses import AuthorData
+
+if typing.TYPE_CHECKING:
+    from segments.models import Segment
+
+
+def delete_segment(
+    segment: "Segment",
+    author: AuthorData,
+) -> None:
+    """
+    Delete a segment using optimized bulk operations.
+
+    Uses bulk UPDATE/DELETE operations instead of individual soft-deletes,
+    reducing the number of database queries from O(n) to O(1) where n is
+    the number of rules and conditions.
+
+    Note: This is a temporary solution until we redesign the segment data model.
+    """
+    from features.models import FeatureSegment
+    from segments.models import Condition, Segment, SegmentRule
+    from segments.tasks import create_segment_deleted_audit_log
+
+    now = timezone.now()
+
+    segment_name = segment.name
+    segment_uuid = str(segment.uuid)
+    segment_id = segment.id
+    project_id = segment.project_id
+
+    segment_ids = list(
+        Segment.objects.filter(
+            models.Q(id=segment.id) | models.Q(version_of_id=segment.id)
+        ).values_list("id", flat=True)
+    )
+
+    top_level_rule_ids = list(
+        SegmentRule.objects.filter(segment_id__in=segment_ids).values_list(
+            "id", flat=True
+        )
+    )
+
+    all_rule_ids = set(top_level_rule_ids)
+    current_level_ids = top_level_rule_ids
+
+    while current_level_ids:
+        nested_rule_ids = list(
+            SegmentRule.objects.filter(rule_id__in=current_level_ids).values_list(
+                "id", flat=True
+            )
+        )
+        all_rule_ids.update(nested_rule_ids)
+        current_level_ids = nested_rule_ids
+
+    all_rule_ids_list = list(all_rule_ids)
+
+    with transaction.atomic():
+        FeatureSegment.objects.filter(segment_id__in=segment_ids).delete()
+        Condition.objects.filter(rule_id__in=all_rule_ids_list).update(deleted_at=now)
+        SegmentRule.objects.filter(id__in=all_rule_ids_list).update(deleted_at=now)
+        Segment.objects.filter(id__in=segment_ids).update(deleted_at=now)
+
+    create_segment_deleted_audit_log.delay(
+        args=(
+            project_id,
+            segment_name,
+            segment_id,
+            segment_uuid,
+            author.user.id if author.user else None,
+            author.api_key.id if author.api_key else None,
+            now.isoformat(),
+        )
+    )
+
+
+def copy_segment_rules_and_conditions(
+    target_segment: "Segment",
+    source_segment: "Segment",
+) -> None:
+    """
+    Copy rules and conditions from source to target segment using bulk operations.
+
+    If target has existing rules, they are hard-deleted first.
+
+    """
+    from segments.models import Condition, SegmentRule
+
+    assert transaction.get_connection().in_atomic_block, "Must run in a transaction"
+
+    SegmentRule.objects.filter(segment=target_segment).delete()
+
+    top_level_rules = list(SegmentRule.objects.filter(segment=source_segment))
+    if not top_level_rules:
+        return
+
+    rule_id_to_cloned_rule: dict[int, SegmentRule] = {}
+
+    cloned_top_rules = [
+        SegmentRule(
+            uuid=uuid.uuid4(),
+            segment=target_segment,
+            rule=None,
+            type=rule.type,
+        )
+        for rule in top_level_rules
+    ]
+    SegmentRule.objects.bulk_create(cloned_top_rules)
+
+    for rule, cloned in zip(top_level_rules, cloned_top_rules):
+        rule_id_to_cloned_rule[rule.id] = cloned
+
+    current_level_rule_ids = [r.id for r in top_level_rules]
+    while current_level_rule_ids:
+        nested_rules = list(
+            SegmentRule.objects.filter(rule_id__in=current_level_rule_ids)
+        )
+        cloned_nested = []
+        for rule in nested_rules:
+            assert rule.rule_id is not None
+            cloned_nested.append(
+                SegmentRule(
+                    uuid=uuid.uuid4(),
+                    segment=None,
+                    rule=rule_id_to_cloned_rule[rule.rule_id],
+                    type=rule.type,
+                )
+            )
+
+        if cloned_nested:
+            SegmentRule.objects.bulk_create(cloned_nested)
+
+        for rule, cloned in zip(nested_rules, cloned_nested):
+            rule_id_to_cloned_rule[rule.id] = cloned
+
+        current_level_rule_ids = [r.id for r in nested_rules]
+
+    source_conditions = list(
+        Condition.objects.filter(rule_id__in=rule_id_to_cloned_rule.keys())
+    )
+    if source_conditions:
+        Condition.objects.bulk_create(
+            [
+                Condition(
+                    uuid=uuid.uuid4(),
+                    rule=rule_id_to_cloned_rule[c.rule_id],
+                    operator=c.operator,
+                    property=c.property,
+                    value=c.value,
+                    description=c.description,
+                    created_with_segment=True,
+                )
+                for c in source_conditions
+            ]
+        )

api/segments/tasks.py

@@ -1,10 +1,37 @@
+from django.utils.dateparse import parse_datetime
 from task_processor.decorators import (
     register_task_handler,
 )
 
+from audit.constants import SEGMENT_DELETED_MESSAGE
+from audit.models import AuditLog
+from audit.related_object_type import RelatedObjectType
 from segments.models import Segment
 
 
 @register_task_handler()
 def delete_segment(segment_id: int) -> None:
     Segment.objects.get(pk=segment_id).delete()
+
+
+@register_task_handler()
+def create_segment_deleted_audit_log(
+    project_id: int,
+    segment_name: str,
+    segment_id: int,
+    segment_uuid: str,
+    user_id: int | None,
+    master_api_key_id: int | None,
+    created_date: str,
+) -> None:
+    AuditLog.objects.create(
+        project_id=project_id,
+        environment=None,
+        log=SEGMENT_DELETED_MESSAGE % segment_name,
+        author_id=user_id,
+        master_api_key_id=master_api_key_id,
+        related_object_id=segment_id,
+        related_object_type=RelatedObjectType.SEGMENT.name,
+        related_object_uuid=segment_uuid,
+        created_date=parse_datetime(created_date),
+    )

api/segments/views.py

@@ -11,6 +11,7 @@
 from rest_framework.response import Response
 
 from app.pagination import CustomPagination
+from core.dataclasses import AuthorData
 from edge_api.identities.models import EdgeIdentity
 from environments.identities.models import Identity
 from environments.models import Environment
@@ -28,6 +29,7 @@
     SegmentListQuerySerializer,
     SegmentSerializer,
 )
+from .services import delete_segment
 
 logger = logging.getLogger()
 
@@ -122,6 +124,12 @@ def associated_features(self, request, *args, **kwargs):  # type: ignore[no-unty
         serializer = self.get_serializer(queryset, many=True)
         return Response(serializer.data)
 
+    def destroy(self, request: Request, *args: Any, **kwargs: Any) -> Response:
+        segment = self.get_object()
+        author = AuthorData.from_request(request)
+        delete_segment(segment, author=author)
+        return Response(status=status.HTTP_204_NO_CONTENT)
+
     @swagger_auto_schema(
         request_body=CloneSegmentSerializer,
         responses={201: SegmentSerializer()},

api/tests/unit/features/versioning/test_unit_versioning_dataclasses.py

@@ -2,7 +2,8 @@
 
 import pytest
 
-from features.versioning.dataclasses import AuthorData, Conflict
+from core.dataclasses import AuthorData
+from features.versioning.dataclasses import Conflict
 
 
 @pytest.mark.parametrize("segment_id, expected_result", ((None, True), (1, False)))