Add RetentionEntity type

There was no applicable type to reuse

frontend/src/queries/nodes/InsightQuery/utils/filtersToQueryNode.test.ts

@@ -404,8 +404,8 @@ describe('filtersToQueryNode', () => {
                 retention_type: 'retention_first_time',
                 retention_reference: 'total',
                 total_intervals: 2,
-                returning_entity: [{ a: 1 }],
-                target_entity: [{ b: 1 }],
+                returning_entity: { id: '1' },
+                target_entity: { id: '1' },
                 period: RetentionPeriod.Day,
             }
 
@@ -417,8 +417,8 @@ describe('filtersToQueryNode', () => {
                     retention_type: 'retention_first_time',
                     retention_reference: 'total',
                     total_intervals: 2,
-                    returning_entity: [{ a: 1 }],
-                    target_entity: [{ b: 1 }],
+                    returning_entity: { id: '1' },
+                    target_entity: { id: '1' },
                     period: RetentionPeriod.Day,
                 },
             }

frontend/src/queries/schema.json

@@ -2536,6 +2536,31 @@
             "required": ["key", "operator", "type", "value"],
             "type": "object"
         },
+        "RetentionEntity": {
+            "additionalProperties": false,
+            "properties": {
+                "id": {
+                    "type": ["string", "number"]
+                },
+                "kind": {
+                    "enum": ["ActionsNode", "EventsNode"],
+                    "type": "string"
+                },
+                "name": {
+                    "type": "string"
+                },
+                "order": {
+                    "type": "number"
+                },
+                "type": {
+                    "$ref": "#/definitions/EntityType"
+                },
+                "uuid": {
+                    "type": "string"
+                }
+            },
+            "type": "object"
+        },
         "RetentionFilter": {
             "additionalProperties": false,
             "description": "`RetentionFilterType` minus everything inherited from `FilterType`",
@@ -2551,10 +2576,10 @@
                     "$ref": "#/definitions/RetentionType"
                 },
                 "returning_entity": {
-                    "type": "object"
+                    "$ref": "#/definitions/RetentionEntity"
                 },
                 "target_entity": {
-                    "type": "object"
+                    "$ref": "#/definitions/RetentionEntity"
                 },
                 "total_intervals": {
                     "type": "integer"

frontend/src/scenes/insights/summarizeInsight.test.ts

@@ -308,12 +308,12 @@ describe('summarizing insights', () => {
                         target_entity: {
                             id: '$autocapture',
                             name: '$autocapture',
-                            type: 'event',
+                            type: 'events',
                         },
                         returning_entity: {
                             id: '$autocapture',
                             name: '$autocapture',
-                            type: 'event',
+                            type: 'events',
                         },
                         retention_type: RETENTION_FIRST_TIME,
                     } as RetentionFilterType,
@@ -333,12 +333,12 @@ describe('summarizing insights', () => {
                         target_entity: {
                             id: 'purchase',
                             name: 'purchase',
-                            type: 'event',
+                            type: 'events',
                         },
                         returning_entity: {
                             id: '$pageview',
                             name: '$pageview',
-                            type: 'event',
+                            type: 'events',
                         },
                         retention_type: RETENTION_RECURRING,
                         aggregation_group_type_index: 0,
@@ -731,12 +731,12 @@ describe('summarizing insights', () => {
                     target_entity: {
                         id: '$autocapture',
                         name: '$autocapture',
-                        type: 'event',
+                        type: 'events',
                     },
                     returning_entity: {
                         id: '$autocapture',
                         name: '$autocapture',
-                        type: 'event',
+                        type: 'events',
                     },
                     retention_type: RETENTION_FIRST_TIME,
                 },
@@ -760,12 +760,12 @@ describe('summarizing insights', () => {
                     target_entity: {
                         id: 'purchase',
                         name: 'purchase',
-                        type: 'event',
+                        type: 'events',
                     },
                     returning_entity: {
                         id: '$pageview',
                         name: '$pageview',
-                        type: 'event',
+                        type: 'events',
                     },
                     retention_type: RETENTION_RECURRING,
                 },

frontend/src/types.ts

@@ -37,6 +37,7 @@ import type {
     InsightVizNode,
     Node,
 } from './queries/schema'
+import { NodeKind } from './queries/schema'
 
 export type Optional<T, K extends string | number | symbol> = Omit<T, K> & { [K in keyof T]?: T[K] }
 
@@ -1874,13 +1875,24 @@ export interface PathsFilterType extends FilterType {
     path_end_key?: string // Paths People End Key
     path_dropoff_key?: string // Paths People Dropoff Key
 }
+
+export interface RetentionEntity {
+    id?: string | number // TODO: Fix weird typing issues
+    kind?: NodeKind.ActionsNode | NodeKind.EventsNode
+    name?: string
+    type?: EntityType
+    // @asType integer
+    order?: number
+    uuid?: string
+}
+
 export interface RetentionFilterType extends FilterType {
     retention_type?: RetentionType
     retention_reference?: 'total' | 'previous' // retention wrt cohort size or previous period
     /** @asType integer */
     total_intervals?: number // retention total intervals
-    returning_entity?: Record<string, any>
-    target_entity?: Record<string, any>
+    returning_entity?: RetentionEntity
+    target_entity?: RetentionEntity
     period?: RetentionPeriod
 }
 export interface LifecycleFilterType extends FilterType {

posthog/hogql/property.py

@@ -33,6 +33,7 @@
     PropertyGroupFilter,
     PropertyGroupFilterValue,
     FilterLogicalOperator,
+    RetentionEntity,
 )
 
 
@@ -387,18 +388,18 @@ def action_to_expr(action: Action) -> ast.Expr:
         return ast.Or(exprs=or_queries)
 
 
-def entity_to_expr(entity: dict, default_event=PAGEVIEW_EVENT) -> ast.Expr:
-    if entity["type"] == TREND_FILTER_TYPE_ACTIONS and entity["id"] is not None:
-        action = Action.objects.get(pk=entity["id"])
+def entity_to_expr(entity: RetentionEntity, default_event=PAGEVIEW_EVENT) -> ast.Expr:
+    if entity.type == TREND_FILTER_TYPE_ACTIONS and entity.id is not None:
+        action = Action.objects.get(pk=entity.id)
         return action_to_expr(action)
-    elif entity["type"] == TREND_FILTER_TYPE_EVENTS:
-        if entity["id"] is None:
+    elif entity.type == TREND_FILTER_TYPE_EVENTS:
+        if entity.id is None:
             return ast.Constant(value=True)
 
         return ast.CompareOperation(
             op=ast.CompareOperationOp.Eq,
             left=ast.Field(chain=["events", "event"]),
-            right=ast.Constant(value=entity["id"]),
+            right=ast.Constant(value=entity.id),
         )
 
     return ast.CompareOperation(

posthog/hogql_queries/insights/retention_query_runner.py

@@ -24,6 +24,7 @@
     HogQLQueryModifiers,
     RetentionQueryResponse,
     IntervalType,
+    RetentionEntity,
 )
 from posthog.schema import RetentionQuery, RetentionType
 
@@ -46,10 +47,12 @@ def __init__(
         super().__init__(query, team=team, timings=timings, modifiers=modifiers, in_export_context=in_export_context)
 
     def get_applicable_entity(self, event_query_type):
-        default_entity = {
-            "id": "$pageview",
-            "type": TREND_FILTER_TYPE_EVENTS,
-        }
+        default_entity = RetentionEntity(
+            **{
+                "id": "$pageview",
+                "type": TREND_FILTER_TYPE_EVENTS,
+            }
+        )
         target_entity = self.query.retentionFilter.target_entity or default_entity
         if event_query_type in [RetentionQueryType.TARGET, RetentionQueryType.TARGET_FIRST_TIME]:
             return target_entity

posthog/schema.py

@@ -485,6 +485,23 @@ class RecordingDurationFilter(BaseModel):
     value: float
 
 
+class Kind(str, Enum):
+    ActionsNode = "ActionsNode"
+    EventsNode = "EventsNode"
+
+
+class RetentionEntity(BaseModel):
+    model_config = ConfigDict(
+        extra="forbid",
+    )
+    id: Optional[Union[str, float]] = None
+    kind: Optional[Kind] = None
+    name: Optional[str] = None
+    order: Optional[float] = None
+    type: Optional[EntityType] = None
+    uuid: Optional[str] = None
+
+
 class RetentionReference(str, Enum):
     total = "total"
     previous = "previous"
@@ -631,7 +648,7 @@ class VizSpecificOptions(BaseModel):
     RETENTION: Optional[RETENTION] = None
 
 
-class Kind(str, Enum):
+class Kind1(str, Enum):
     unit = "unit"
     duration_s = "duration_s"
     percentage = "percentage"
@@ -644,7 +661,7 @@ class WebOverviewItem(BaseModel):
     changeFromPreviousPct: Optional[float] = None
     isIncreaseBad: Optional[bool] = None
     key: str
-    kind: Kind
+    kind: Kind1
     previous: Optional[float] = None
     value: Optional[float] = None
 
@@ -934,8 +951,8 @@ class RetentionFilter(BaseModel):
     period: Optional[RetentionPeriod] = None
     retention_reference: Optional[RetentionReference] = None
     retention_type: Optional[RetentionType] = None
-    returning_entity: Optional[Dict[str, Any]] = None
-    target_entity: Optional[Dict[str, Any]] = None
+    returning_entity: Optional[RetentionEntity] = None
+    target_entity: Optional[RetentionEntity] = None
     total_intervals: Optional[int] = None