Improve data model for live segments

Co-authored-by: Claude <[email protected]>

Author: emyller

api/import_export/export.py

@@ -8,7 +8,7 @@
 import boto3
 from django.core import serializers
 from django.core.serializers.json import DjangoJSONEncoder
-from django.db.models import F, Model, Q
+from django.db.models import Model, Q
 
 from edge_api.identities.export import export_edge_identity_and_overrides
 from environments.identities.models import Identity
@@ -130,28 +130,28 @@ def export_projects(
         *_export_entities(
             _EntityExportConfig(
                 Segment,
-                Q(project__organisation__id=organisation_id, id=F("version_of")),
+                Q(project__organisation__id=organisation_id, version_of__isnull=True),
             ),
             _EntityExportConfig(
                 SegmentRule,
                 Q(
                     segment__project__organisation__id=organisation_id,
-                    segment_id=F("segment__version_of"),
+                    segment__version_of__isnull=True,
                 )
                 | Q(
                     rule__segment__project__organisation__id=organisation_id,
-                    rule__segment_id=F("rule__segment__version_of"),
+                    rule__segment__version_of__isnull=True,
                 ),
             ),
             _EntityExportConfig(
                 Condition,
                 Q(
                     rule__segment__project__organisation__id=organisation_id,
-                    rule__segment_id=F("rule__segment__version_of"),
+                    rule__segment__version_of__isnull=True,
                 )
                 | Q(
                     rule__rule__segment__project__organisation__id=organisation_id,
-                    rule__rule__segment_id=F("rule__rule__segment__version_of"),
+                    rule__rule__segment__version_of__isnull=True,
                 ),
             ),
             _EntityExportConfig(Tag, default_filter),

api/segments/migrations/0031_set_version_of_to_null_for_canonical_segments.py

@@ -0,0 +1,21 @@
+# Generated by Django 4.2.22 on 2025-11-11 03:43
+
+from django.db import migrations
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("segments", "0030_add_default_to_segment_version"),
+    ]
+
+    operations = [
+        # Set version_of to NULL for canonical segments (where version_of_id = id).
+        # This follows the same pattern as migration 0023 which originally set
+        # version_of_id = id. Like that migration, this may block during deployment
+        # depending on the number of canonical segments in the database.
+        migrations.RunSQL(
+            sql="UPDATE segments_segment SET version_of_id = NULL WHERE version_of_id = id;",
+            reverse_sql="UPDATE segments_segment SET version_of_id = id WHERE version_of_id IS NULL;",
+        ),
+    ]

api/segments/models.py

@@ -8,9 +8,8 @@
 from django.core.exceptions import ValidationError
 from django.db import models, transaction
 from django_lifecycle import (  # type: ignore[import-untyped]
-    AFTER_CREATE,
+    BEFORE_CREATE,
     LifecycleModelMixin,
-    hook,
 )
 from flag_engine.segments import constants
 
@@ -44,7 +43,7 @@ def get_queryset(self):  # type: ignore[no-untyped-def]
         Returns only the canonical segments, which will always be
         the highest version.
         """
-        return super().get_queryset().filter(id=models.F("version_of"))
+        return super().get_queryset().filter(version_of__isnull=True)
 
 
 class ConfiguredOrderManager(SoftDeleteExportableManager, models.Manager[ModelT]):
@@ -135,22 +134,22 @@ def __str__(self):  # type: ignore[no-untyped-def]
     def get_skip_create_audit_log(self) -> bool:
         if self.is_system_segment:
             return True
-        is_revision = bool(self.version_of_id and self.version_of_id != self.id)
+        is_revision = self.version_of_id is not None
         return is_revision
 
-    @hook(AFTER_CREATE, when="version_of", is_now=None)
-    def set_version_of_to_self_if_none(self):  # type: ignore[no-untyped-def]
-        """
-        This allows the segment model to reference all versions of
-        itself including itself.
-        """
-        self.version_of = self
-        self.save_without_historical_record()
+    @BEFORE_CREATE(when="version_of", is_now=None)
+    def set_default_version_to_one_if_new_segment(self):  # type: ignore[no-untyped-def]
+        if self.version is None:
+            self.version = 1
 
     @transaction.atomic
     def clone(self, is_revision: bool = False, **extra_attrs: typing.Any) -> "Segment":
         """
-        Create a revision of the segment
+        Create a revision of the segment.
+
+        Args:
+            is_revision (bool): If True, creates a revision of the segment; if False, promotes the clone as canonical.
+            **extra_attrs: Additional attributes to set on the cloned segment.
         """
         cloned_segment = deepcopy(self)
         cloned_segment.pk = None
@@ -164,7 +163,7 @@ def clone(self, is_revision: bool = False, **extra_attrs: typing.Any) -> "Segmen
         cloned_segment.copy_rules_and_conditions_from(self)
 
         # Handle versioning
-        version_of = self if is_revision else cloned_segment
+        version_of = self if is_revision else None
         cloned_segment.version_of = extra_attrs.get("version_of", version_of)
         cloned_segment.version = self.version if is_revision else 1
         Segment.objects.filter(pk=cloned_segment.pk).update(

api/tests/unit/segments/test_unit_segments_migrations.py

@@ -155,7 +155,7 @@ def test_add_versioning_to_segments_forwards(migrator: Migrator) -> None:
     # Then the version_of attribute is correctly set.
     NewSegment = new_state.apps.get_model("segments", "Segment")
     new_segment = NewSegment.objects.get(id=segment.id)
-    assert new_segment.version_of == new_segment
+    assert new_segment.version_of_id == new_segment.id
 
 
 @pytest.mark.skipif(

api/tests/unit/segments/test_unit_segments_models.py

@@ -182,7 +182,7 @@ def test_Segment_clone__can_create_standalone_segment_clone(
     # Then
     assert cloned_segment != segment
     assert cloned_segment.name == "another-segment"
-    assert cloned_segment.version_of == cloned_segment
+    assert cloned_segment.version_of is None
     assert cloned_segment.version == 1
 
 

api/tests/unit/segments/test_unit_segments_views.py

@@ -207,7 +207,9 @@ def test_segments_limit_ignores_old_segment_versions(
 
     # and create some older versions for the segment fixture
     segment.clone(is_revision=True)
-    assert Segment.objects.filter(version_of=segment).count() == 2
+    # With the NULL-based canonical representation, only revisions are counted here
+    # (the canonical segment has version_of=None).
+    assert Segment.objects.filter(version_of=segment).count() == 1
     assert Segment.live_objects.count() == 1
 
     url = reverse("api-v1:projects:project-segments-list", args=[project.id])
@@ -1058,7 +1060,9 @@ def test_update_segment_versioned_segment(
     assert response.status_code == status.HTTP_200_OK
 
     # Now verify that a new versioned segment has been set.
-    assert Segment.objects.filter(version_of=segment).count() == 2
+    # With the NULL-based canonical representation, canonical segments have
+    # version_of=None, so only revisions are counted here.
+    assert Segment.objects.filter(version_of=segment).count() == 1
 
     # Now check the previously versioned segment to match former count of conditions.
 
@@ -1093,7 +1097,10 @@ def test_update_segment_versioned_segment_with_thrown_exception(
         rule=nested_rule, property="foo", operator=EQUAL, value="bar"
     )
 
-    assert segment.version == 1 == Segment.objects.filter(version_of=segment).count()
+    # With the NULL-based canonical representation, canonical segments have
+    # version_of=None, so they are not included in this count.
+    assert segment.version == 1
+    assert Segment.objects.filter(version_of=segment).count() == 0
 
     new_condition_property = "foo2"
     new_condition_value = "bar"
@@ -1144,7 +1151,9 @@ def test_update_segment_versioned_segment_with_thrown_exception(
     segment.refresh_from_db()
 
     # Now verify that the version of the segment has not been changed.
-    assert segment.version == 1 == Segment.objects.filter(version_of=segment).count()
+    # The transaction should have rolled back, so no revisions should exist.
+    assert segment.version == 1
+    assert Segment.objects.filter(version_of=segment).count() == 0
 
 
 @pytest.mark.parametrize(

api/util/mappers/engine.py

@@ -201,10 +201,9 @@ def map_environment_to_engine(
     organisation: "Organisation" = project.organisation
 
     # Read relationships - grab all the data needed from the ORM here.
-
-    project_segments = [
-        ps for ps in project.segments.all() if ps.id == ps.version_of_id
-    ]
+    # Note: project.segments is prefetched with Segment.live_objects which already
+    # filters by version_of__isnull=True, so no need to filter again here.
+    project_segments = list(project.segments.all())
 
     project_segment_rules_by_segment_id: Dict[
         int,