fix(cohort): Add count field to cohort model (#8585)

* add count field to cohort model

* make count safe

* update version with count

ee/clickhouse/models/cohort.py

@@ -282,7 +282,7 @@ def insert_static_cohort(person_uuids: List[Optional[uuid.UUID]], cohort_id: int
     sync_execute(INSERT_PERSON_STATIC_COHORT, persons)
 
 
-def recalculate_cohortpeople(cohort: Cohort):
+def recalculate_cohortpeople(cohort: Cohort) -> Optional[int]:
     cohort_filter, cohort_params = format_person_query(cohort, 0, custom_match_field="id")
 
     before_count = sync_execute(GET_COHORT_SIZE_SQL, {"cohort_id": cohort.pk, "team_id": cohort.team_id})
@@ -307,14 +307,21 @@ def recalculate_cohortpeople(cohort: Cohort):
     remove_cohortpeople_sql = REMOVE_PEOPLE_NOT_MATCHING_COHORT_ID_SQL.format(cohort_filter=cohort_filter)
     sync_execute(remove_cohortpeople_sql, {**cohort_params, "cohort_id": cohort.pk, "team_id": cohort.team_id})
 
-    count = sync_execute(GET_COHORT_SIZE_SQL, {"cohort_id": cohort.pk, "team_id": cohort.team_id})
-    logger.info(
-        "Recalculating cohortpeople done",
-        team_id=cohort.team_id,
-        cohort_id=cohort.pk,
-        size_before=before_count[0][0],
-        size=count[0][0],
-    )
+    count_result = sync_execute(GET_COHORT_SIZE_SQL, {"cohort_id": cohort.pk, "team_id": cohort.team_id})
+
+    if count_result and len(count_result) and len(count_result[0]):
+        count = count_result[0][0]
+
+        logger.info(
+            "Recalculating cohortpeople done",
+            team_id=cohort.team_id,
+            cohort_id=cohort.pk,
+            size_before=before_count[0][0],
+            size=count,
+        )
+        return count
+
+    return None
 
 
 def simplified_cohort_filter_properties(cohort: Cohort, team: Team) -> List[Property]:

latest_migrations.manifest

@@ -4,7 +4,7 @@ axes: 0006_remove_accesslog_trusted
 contenttypes: 0002_remove_content_type_name
 database: 0002_auto_20190129_2304
 ee: 0007_dashboard_permissions
-posthog: 0206_global_tags_setup
+posthog: 0207_cohort_count
 rest_hooks: 0002_swappable_hook_model
 sessions: 0001_initial
 social_django: 0010_uid_db_index

posthog/api/cohort.py

@@ -42,7 +42,6 @@ from posthog.tasks.calculate_cohort import (
 
 class CohortSerializer(serializers.ModelSerializer):
     created_by = UserBasicSerializer(read_only=True)
-    count = serializers.SerializerMethodField()
     earliest_timestamp_func = get_earliest_timestamp
 
     class Meta:
@@ -144,11 +143,6 @@ class CohortSerializer(serializers.ModelSerializer):
 
         return cohort
 
-    def get_count(self, action: Cohort) -> Optional[int]:
-        if hasattr(action, "count"):
-            return action.count  # type: ignore
-        return None
-
 
 class CohortViewSet(StructuredViewSetMixin, viewsets.ModelViewSet):
     queryset = Cohort.objects.all()
@@ -160,14 +154,6 @@ class CohortViewSet(StructuredViewSetMixin, viewsets.ModelViewSet):
         if self.action == "list":
             queryset = queryset.filter(deleted=False)
 
-        cohort_people_count = (
-            CohortPeople.objects.filter(cohort_id=OuterRef("id"), version=OuterRef("version"))
-            .values("cohort_id")
-            .annotate(count=Count("person_id", distinct=True))
-            .values("count")
-        )
-
-        queryset = queryset.annotate(count=cohort_people_count)
         return queryset.prefetch_related("created_by").order_by("-created_at")
 
 

posthog/api/person.py

@@ -494,16 +494,7 @@ class PersonViewSet(StructuredViewSetMixin, viewsets.ModelViewSet):
 
         person = self.get_queryset().get(id=str(request.GET["person_id"]))
 
-        cohort_people_count = (
-            CohortPeople.objects.filter(cohort_id=OuterRef("id"), version=OuterRef("version"))
-            .values("cohort_id")
-            .annotate(count=Count("person_id", distinct=True))
-            .values("count")
-        )
-
-        cohorts = Cohort.objects.annotate(count=Subquery(cohort_people_count)).filter(
-            people__id=person.id, deleted=False
-        )
+        cohorts = Cohort.objects.filter(people__id=person.id, deleted=False)
         return response.Response({"results": CohortSerializer(cohorts, many=True).data})
 
 

posthog/migrations/0207_cohort_count.py

@@ -0,0 +1,14 @@
+# Generated by Django 3.2.5 on 2022-02-14 15:42
+
+from django.db import migrations, models
+
+
+class Migration(migrations.Migration):
+
+    dependencies = [
+        ("posthog", "0206_global_tags_setup"),
+    ]
+
+    operations = [
+        migrations.AddField(model_name="cohort", name="count", field=models.IntegerField(blank=True, null=True),),
+    ]

posthog/models/cohort.py

@@ -80,6 +80,7 @@ class Cohort(models.Model):
     people: models.ManyToManyField = models.ManyToManyField("Person", through="CohortPeople")
     version: models.IntegerField = models.IntegerField(blank=True, null=True)
     pending_version: models.IntegerField = models.IntegerField(blank=True, null=True)
+    count: models.IntegerField = models.IntegerField(blank=True, null=True)
 
     created_by: models.ForeignKey = models.ForeignKey("User", on_delete=models.SET_NULL, blank=True, null=True)
     created_at: models.DateTimeField = models.DateTimeField(default=timezone.now, blank=True, null=True)
@@ -139,12 +140,12 @@ class Cohort(models.Model):
                 "cohort_calculation_started", id=self.pk, current_version=self.version, new_version=pending_version
             )
 
-            recalculate_cohortpeople(self)
+            count = recalculate_cohortpeople(self)
             self.calculate_people(new_version=pending_version)
 
             # Update filter to match pending version if still valid
             Cohort.objects.filter(pk=self.pk).filter(Q(version__lt=pending_version) | Q(version__isnull=True)).update(
-                version=pending_version
+                version=pending_version, count=count
             )
             self.refresh_from_db()