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

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/utils.py

@@ -313,7 +313,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
         """
 

rust/property-defs-rs/src/types.rs

@@ -175,8 +175,8 @@ impl Event {
         let updates = self.into_updates_inner();
         if updates.len() > skip_threshold {
             warn!(
-                "Event {} for team {} has more than 10,000 properties, skipping",
-                event, team_id
+                "Event {} for team {} has more than {} properties, skipping",
+                event, team_id, skip_threshold
             );
             metrics::counter!(EVENTS_SKIPPED, &[("reason", "too_many_properties")]).increment(1);
             return vec![];
@@ -427,8 +427,8 @@ impl EventDefinition {
         sqlx::query!(
             r#"
             INSERT INTO posthog_eventdefinition (id, name, volume_30_day, query_usage_30_day, team_id, project_id, last_seen_at, created_at)
-            VALUES ($1, $2, NULL, NULL, $3, $4, $5, NOW()) ON CONFLICT
-            ON CONSTRAINT posthog_eventdefinition_team_id_name_80fa0b87_uniq
+            VALUES ($1, $2, NULL, NULL, $3, $4, $5, NOW())
+            ON CONFLICT (coalesce(project_id, team_id), name)
             DO UPDATE SET last_seen_at = $5
         "#,
             Uuid::now_v7(),
@@ -472,7 +472,7 @@ impl PropertyDefinition {
             r#"
             INSERT INTO posthog_propertydefinition (id, name, type, group_type_index, is_numerical, volume_30_day, query_usage_30_day, team_id, project_id, property_type)
             VALUES ($1, $2, $3, $4, $5, NULL, NULL, $6, $7, $8)
-            ON CONFLICT (team_id, name, type, coalesce(group_type_index, -1))
+            ON CONFLICT (coalesce(project_id, team_id), name, type, coalesce(group_type_index, -1))
             DO UPDATE SET property_type=EXCLUDED.property_type WHERE posthog_propertydefinition.property_type IS NULL
         "#,
             Uuid::now_v7(),

rust/property-defs-rs/tests/test_migrations/20240830124836_setup_propdefs_tables.sql

@@ -15,6 +15,7 @@ CREATE TABLE IF NOT EXISTS posthog_eventdefinition (
     CONSTRAINT posthog_eventdefinition_team_id_name_80fa0b87_uniq UNIQUE (team_id, name)
 );
 
+CREATE UNIQUE INDEX event_definition_proj_uniq ON posthog_eventdefinition (coalesce(project_id, team_id), name);
 
 CREATE TABLE IF NOT EXISTS posthog_propertydefinition (
     id UUID PRIMARY KEY,
@@ -31,6 +32,7 @@ CREATE TABLE IF NOT EXISTS posthog_propertydefinition (
 );
 
 CREATE UNIQUE INDEX posthog_propertydefinition_uniq ON posthog_propertydefinition (team_id, name, type, coalesce(group_type_index, -1));
+CREATE UNIQUE INDEX posthog_propdef_proj_uniq ON posthog_propertydefinition (coalesce(project_id, team_id), name, type, coalesce(group_type_index, -1));
 
 
 CREATE TABLE IF NOT EXISTS posthog_eventproperty (
@@ -42,3 +44,4 @@ CREATE TABLE IF NOT EXISTS posthog_eventproperty (
 );
 
 CREATE UNIQUE INDEX posthog_event_property_unique_team_event_property ON posthog_eventproperty (team_id, event, property);
+CREATE UNIQUE INDEX posthog_event_property_unique_proj_event_property ON posthog_eventproperty (coalesce(project_id, team_id), event, property);

rust/property-defs-rs/tests/updates.rs

@@ -1,7 +1,9 @@
 use chrono::{DateTime, Duration, Utc};
 use property_defs_rs::types::{Event, PropertyParentType, PropertyValueType};
 use serde_json::json;
-use sqlx::PgPool;
+use sqlx::postgres::PgArguments;
+use sqlx::{Arguments, PgPool};
+use uuid::Uuid;
 
 #[sqlx::test(migrations = "./tests/test_migrations")]
 async fn test_updates(db: PgPool) {
@@ -15,7 +17,6 @@ async fn test_updates(db: PgPool) {
             "some_bool_as_string": "true"
         }
     "#;
-
     let event_src = json!({
         "team_id": 1,
         "project_id": 1,
@@ -105,6 +106,51 @@ async fn test_updates(db: PgPool) {
     .unwrap();
 }
 
+#[sqlx::test(migrations = "./tests/test_migrations")]
+async fn test_update_on_project_id_conflict(db: PgPool) {
+    let definition_created_at: DateTime<Utc> = Utc::now() - Duration::days(1);
+    let mut args = PgArguments::default();
+    args.add(Uuid::now_v7()).unwrap();
+    args.add("foo").unwrap();
+    args.add(1).unwrap();
+    args.add(definition_created_at).unwrap();
+    sqlx::query_with(
+        r#"
+        INSERT INTO posthog_eventdefinition (id, name, volume_30_day, query_usage_30_day, team_id, project_id, last_seen_at, created_at)
+        VALUES ($1, $2, NULL, NULL, $3, NULL, $4, $4) -- project_id is NULL! This definition is from before environments
+    "#, args
+    ).execute(&db).await.unwrap();
+
+    assert_eventdefinition_exists(
+        &db,
+        "foo",
+        1,
+        definition_created_at,
+        Duration::milliseconds(0),
+    )
+    .await
+    .unwrap();
+
+    let before = Utc::now();
+    let event_src = json!({
+        "team_id": 3,
+        "project_id": 1,
+        "event": "foo",
+        "properties": "{}"
+    });
+
+    let event = serde_json::from_value::<Event>(event_src.clone()).unwrap();
+    for update in event.into_updates(10000) {
+        update.issue(&db).await.unwrap();
+    }
+
+    // The event def we created earlier got updated, even though it has a different `team_id`,
+    // because `coalesce(project_id, team_id)` matches
+    assert_eventdefinition_exists(&db, "foo", 1, before, Duration::seconds(1))
+        .await
+        .unwrap();
+}
+
 async fn assert_eventdefinition_exists(
     db: &PgPool,
     name: &str,