From 8569c85172a7b0f2b8584a3e0b36f754da60b0d9 Mon Sep 17 00:00:00 2001
From: Zach Aysan <zachaysan@gmail.com>
Date: Thu, 8 Aug 2024 19:19:38 +0000
Subject: [PATCH] Switch to live objects manager

---
 api/environments/models.py | 2 +-
 api/projects/models.py     | 5 ++++-
 api/segments/views.py      | 2 +-
 3 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/api/environments/models.py b/api/environments/models.py
index dd58ce9fa605..b2895ca8f644 100644
--- a/api/environments/models.py
+++ b/api/environments/models.py
@@ -344,7 +344,7 @@ def get_segments_from_cache(self) -> typing.List[Segment]:
         segments = environment_segments_cache.get(self.id)
         if not segments:
             segments = list(
-                Segment.objects.filter(
+                Segment.live_objects.filter(
                     feature_segments__feature_states__environment=self
                 ).prefetch_related(
                     "rules",
diff --git a/api/projects/models.py b/api/projects/models.py
index d20c26eea833..a9cae6e62ea1 100644
--- a/api/projects/models.py
+++ b/api/projects/models.py
@@ -142,19 +142,22 @@ def edge_v2_identity_overrides_migrated(self) -> bool:
         return self.edge_v2_migration_status == EdgeV2MigrationStatus.COMPLETE
 
     def get_segments_from_cache(self):
+        from segments.models import Segment
+
         segments = project_segments_cache.get(self.id)
 
         if not segments:
             # This is optimised to account for rules nested one levels deep (since we
             # don't support anything above that from the UI at the moment). Anything
             # past that will require additional queries / thought on how to optimise.
-            segments = self.segments.all().prefetch_related(
+            segments = Segment.live_objects.filter(project=self).prefetch_related(
                 "rules",
                 "rules__conditions",
                 "rules__rules",
                 "rules__rules__conditions",
                 "rules__rules__rules",
             )
+
             project_segments_cache.set(
                 self.id, segments, timeout=settings.CACHE_PROJECT_SEGMENTS_SECONDS
             )
diff --git a/api/segments/views.py b/api/segments/views.py
index 566dc9a6c1a1..c1426195aba0 100644
--- a/api/segments/views.py
+++ b/api/segments/views.py
@@ -39,7 +39,7 @@ def get_queryset(self):
         )
         project = get_object_or_404(permitted_projects, pk=self.kwargs["project_pk"])
 
-        queryset = project.segments.all()
+        queryset = Segment.live_objects.filter(project=project)
 
         if self.action == "list":
             # TODO: at the moment, the UI only shows the name and description of the segment in the list view.