feat(environments): Make taxonomy reads + writes project–based

posthog/api/property_definition.py

@@ -1,8 +1,9 @@
 import dataclasses
 import json
 from typing import Any, Optional, cast
+from django.db.models.functions import Coalesce
 
-from django.db import connection
+from django.db import connection, models
 from loginas.utils import is_impersonated_session
 from rest_framework import mixins, request, response, serializers, status, viewsets
 from posthog.api.utils import action
@@ -99,7 +100,7 @@ class QueryContext:
     The raw query is used to both query and count these results
     """
 
-    team_id: int
+    project_id: int
     table: str
     property_definition_fields: str
     property_definition_table: str
@@ -232,7 +233,7 @@ def with_search(self, search_query: str, search_kwargs: dict) -> "QueryContext":
         return dataclasses.replace(
             self,
             search_query=search_query,
-            params={**self.params, "team_id": self.team_id, **search_kwargs},
+            params={**self.params, "project_id": self.project_id, **search_kwargs},
         )
 
     def with_excluded_properties(self, excluded_properties: Optional[str], type: str) -> "QueryContext":
@@ -264,7 +265,7 @@ def as_sql(self, order_by_verified: bool):
             SELECT {self.property_definition_fields}, {self.event_property_field} AS is_seen_on_filtered_events
             FROM {self.table}
             {self._join_on_event_property()}
-            WHERE {self.property_definition_table}.team_id = %(team_id)s
+            WHERE coalesce({self.property_definition_table}.project_id, {self.property_definition_table}.team_id) = %(project_id)s
               AND type = %(type)s
               AND coalesce(group_type_index, -1) = %(group_type_index)s
               {self.excluded_properties_filter}
@@ -281,7 +282,7 @@ def as_count_sql(self):
             SELECT count(*) as full_count
             FROM {self.table}
             {self._join_on_event_property()}
-            WHERE {self.property_definition_table}.team_id = %(team_id)s
+            WHERE coalesce({self.property_definition_table}.project_id, {self.property_definition_table}.team_id) = %(project_id)s
               AND type = %(type)s
               AND coalesce(group_type_index, -1) = %(group_type_index)s
              {self.excluded_properties_filter} {self.name_filter} {self.numerical_filter} {self.search_query} {self.event_property_filter} {self.is_feature_flag_filter}
@@ -296,7 +297,7 @@ def _join_on_event_property(self):
             {self.event_property_join_type} (
                 SELECT DISTINCT property
                 FROM posthog_eventproperty
-                WHERE team_id = %(team_id)s {self.event_name_join_filter}
+                WHERE coalesce(project_id, team_id) = %(project_id)s {self.event_name_join_filter}
             ) {self.posthog_eventproperty_table_join_alias}
             ON {self.posthog_eventproperty_table_join_alias}.property = name
             """
@@ -537,7 +538,7 @@ def dangerously_get_queryset(self):
 
         query_context = (
             QueryContext(
-                team_id=self.team_id,
+                project_id=self.project_id,
                 table=(
                     "ee_enterprisepropertydefinition FULL OUTER JOIN posthog_propertydefinition ON posthog_propertydefinition.id=ee_enterprisepropertydefinition.propertydefinition_ptr_id"
                     if use_enterprise_taxonomy
@@ -621,8 +622,10 @@ def seen_together(self, request: request.Request, *args: Any, **kwargs: Any) ->
         serializer = SeenTogetherQuerySerializer(data=request.GET)
         serializer.is_valid(raise_exception=True)
 
-        matches = EventProperty.objects.filter(
-            team_id=self.team_id,
+        matches = EventProperty.objects.alias(
+            effective_project_id=Coalesce("project_id", "team_id", output_field=models.BigIntegerField())
+        ).filter(
+            effective_project_id=self.project_id,  # type: ignore
             event__in=serializer.validated_data["event_names"],
             property=serializer.validated_data["property_name"],
         )

posthog/api/test/__snapshots__/test_decide.ambr

@@ -3890,6 +3890,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id",
          "posthog_team"."id",
@@ -4260,6 +4261,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -4693,6 +4695,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -4986,6 +4989,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -5171,6 +5175,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id",
          "posthog_team"."id",
@@ -5602,6 +5607,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -5834,6 +5840,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -6165,6 +6172,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -6344,6 +6352,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -6475,6 +6484,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id",
          "posthog_team"."id",
@@ -6906,6 +6916,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -7138,6 +7149,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -7465,6 +7477,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"
@@ -7640,6 +7653,7 @@
          "posthog_hogfunction"."inputs",
          "posthog_hogfunction"."encrypted_inputs",
          "posthog_hogfunction"."filters",
+         "posthog_hogfunction"."mappings",
          "posthog_hogfunction"."masking",
          "posthog_hogfunction"."template_id"
   FROM "posthog_hogfunction"

posthog/api/utils.py

@@ -312,7 +312,7 @@ def create_event_definitions_sql(
             SELECT {",".join(event_definition_fields)}
             FROM posthog_eventdefinition
             {enterprise_join}
-            WHERE team_id = %(project_id)s {conditions}
+            WHERE coalesce(project_id, team_id) = %(project_id)s {conditions}
             ORDER BY {additional_ordering}name ASC
         """
 

posthog/migrations/0530_taxonomy_unique_on_project.py

@@ -0,0 +1,119 @@
+# Generated by Django 4.2.15 on 2024-12-09 15:51
+
+from django.db import migrations
+from django.db import models
+import django.db.models.functions.comparison
+import posthog.models.utils
+from django.contrib.postgres.operations import AddIndexConcurrently, RemoveIndexConcurrently
+
+
+class Migration(migrations.Migration):
+    atomic = False  # Added to support concurrent index creation
+    dependencies = [("posthog", "0529_hog_function_mappings")]
+
+    operations = [
+        # First clean up rows that would fail the project-based unique constraints we're adding
+        migrations.RunSQL(
+            sql="""
+            DELETE FROM posthog_propertydefinition
+            WHERE team_id IN (
+                SELECT id FROM posthog_team WHERE id != project_id
+            );""",
+            reverse_sql=migrations.RunSQL.noop,
+            elidable=True,
+        ),
+        migrations.RunSQL(
+            sql="""
+            DELETE FROM posthog_eventdefinition
+            WHERE team_id IN (
+                SELECT id FROM posthog_team WHERE id != project_id
+            );""",
+            reverse_sql=migrations.RunSQL.noop,
+            elidable=True,
+        ),
+        migrations.RunSQL(
+            sql="""
+            DELETE FROM posthog_eventproperty
+            WHERE team_id IN (
+                SELECT id FROM posthog_team WHERE id != project_id
+            );""",
+            reverse_sql=migrations.RunSQL.noop,
+            elidable=True,
+        ),
+        # Remove misguided `project_id`-only indexes from the previous migration
+        RemoveIndexConcurrently(
+            model_name="eventproperty",
+            name="posthog_eve_proj_id_22de03_idx",
+        ),
+        RemoveIndexConcurrently(
+            model_name="eventproperty",
+            name="posthog_eve_proj_id_26dbfb_idx",
+        ),
+        RemoveIndexConcurrently(
+            model_name="propertydefinition",
+            name="index_property_def_query_proj",
+        ),
+        RemoveIndexConcurrently(
+            model_name="propertydefinition",
+            name="posthog_pro_project_3583d2_idx",
+        ),
+        # Add new useful indexes using `coalesce(project_id, team_id)`
+        AddIndexConcurrently(
+            model_name="eventproperty",
+            index=models.Index(
+                django.db.models.functions.comparison.Coalesce(models.F("project_id"), models.F("team_id")),
+                models.F("event"),
+                name="posthog_eve_proj_id_22de03_idx",
+            ),
+        ),
+        AddIndexConcurrently(
+            model_name="eventproperty",
+            index=models.Index(
+                django.db.models.functions.comparison.Coalesce(models.F("project_id"), models.F("team_id")),
+                models.F("property"),
+                name="posthog_eve_proj_id_26dbfb_idx",
+            ),
+        ),
+        AddIndexConcurrently(
+            model_name="propertydefinition",
+            index=models.Index(
+                django.db.models.functions.comparison.Coalesce(models.F("project_id"), models.F("team_id")),
+                models.F("type"),
+                django.db.models.functions.comparison.Coalesce(models.F("group_type_index"), -1),
+                models.OrderBy(models.F("query_usage_30_day"), descending=True, nulls_last=True),
+                models.OrderBy(models.F("name")),
+                name="index_property_def_query_proj",
+            ),
+        ),
+        AddIndexConcurrently(
+            model_name="propertydefinition",
+            index=models.Index(
+                django.db.models.functions.comparison.Coalesce(models.F("project_id"), models.F("team_id")),
+                models.F("type"),
+                models.F("is_numerical"),
+                name="posthog_pro_project_3583d2_idx",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="eventdefinition",
+            constraint=posthog.models.utils.UniqueConstraintByExpression(
+                concurrently=True, expression="(coalesce(project_id, team_id), name)", name="event_definition_proj_uniq"
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="eventproperty",
+            constraint=posthog.models.utils.UniqueConstraintByExpression(
+                concurrently=True,
+                expression="(coalesce(project_id, team_id), event, property)",
+                name="posthog_event_property_unique_proj_event_property",
+            ),
+        ),
+        migrations.AddConstraint(
+            model_name="propertydefinition",
+            constraint=posthog.models.utils.UniqueConstraintByExpression(
+                concurrently=True,
+                expression="(coalesce(project_id, team_id), name, type, coalesce(group_type_index, -1))",
+                name="posthog_propdef_proj_uniq",
+            ),
+        ),
+    ]

posthog/migrations/max_migration.txt

@@ -1 +1 @@
-0529_hog_function_mappings
+0530_taxonomy_unique_on_project

posthog/models/event_definition.py

@@ -3,7 +3,7 @@
 from django.utils import timezone
 
 from posthog.models.team import Team
-from posthog.models.utils import UUIDModel
+from posthog.models.utils import UUIDModel, UniqueConstraintByExpression
 
 
 class EventDefinition(UUIDModel):
@@ -27,7 +27,6 @@ class EventDefinition(UUIDModel):
     volume_30_day = models.IntegerField(default=None, null=True)
 
     class Meta:
-        unique_together = ("team", "name")
         indexes = [
             # Index on project_id foreign key
             models.Index(fields=["project"], name="posthog_eve_proj_id_f93fcbb0"),
@@ -37,6 +36,14 @@ class Meta:
                 opclasses=["gin_trgm_ops"],
             ),  # To speed up DB-based fuzzy searching
         ]
+        unique_together = ("team", "name")
+        constraints = [
+            UniqueConstraintByExpression(
+                concurrently=True,
+                name="event_definition_proj_uniq",
+                expression="(coalesce(project_id, team_id), name)",
+            )
+        ]
 
     def __str__(self) -> str:
         return f"{self.name} / {self.team.name}"

posthog/models/event_property.py

@@ -1,7 +1,9 @@
 from django.db import models
 
 from posthog.models.team import Team
-from posthog.models.utils import sane_repr
+from posthog.models.utils import UniqueConstraintByExpression, sane_repr
+from django.db.models.expressions import F
+from django.db.models.functions import Coalesce
 
 
 class EventProperty(models.Model):
@@ -15,15 +17,20 @@ class Meta:
             models.UniqueConstraint(
                 fields=["team", "event", "property"],
                 name="posthog_event_property_unique_team_event_property",
-            )
+            ),
+            UniqueConstraintByExpression(
+                concurrently=True,
+                name="posthog_event_property_unique_proj_event_property",
+                expression="(coalesce(project_id, team_id), event, property)",
+            ),
         ]
         indexes = [
             # Index on project_id foreign key
             models.Index(fields=["project"], name="posthog_eve_proj_id_dd2337d2"),
             models.Index(fields=["team", "event"]),
-            models.Index(fields=["project", "event"], name="posthog_eve_proj_id_22de03_idx"),
+            models.Index(Coalesce(F("project_id"), F("team_id")), F("event"), name="posthog_eve_proj_id_22de03_idx"),
             models.Index(fields=["team", "property"]),
-            models.Index(fields=["project", "property"], name="posthog_eve_proj_id_26dbfb_idx"),
+            models.Index(Coalesce(F("project_id"), F("team_id")), F("property"), name="posthog_eve_proj_id_26dbfb_idx"),
         ]
 
     __repr__ = sane_repr("event", "property", "team_id")