Add retention metric type for cohort retention analysis

Author: nicosuave

sidemantic/adapters/sidemantic.py

@@ -297,6 +297,13 @@ def _parse_model(self, model_def: dict) -> Model | None:
                 conversion_event=measure_def.get("conversion_event"),
                 conversion_window=measure_def.get("conversion_window"),
                 offset_window=measure_def.get("offset_window"),
+                # Retention parameters
+                cohort_event=measure_def.get("cohort_event"),
+                activity_event=measure_def.get("activity_event"),
+                periods=measure_def.get("periods"),
+                retention_granularity=measure_def.get("granularity")
+                if measure_def.get("type") == "retention"
+                else None,
                 # Cumulative/window parameters
                 window=measure_def.get("window"),
                 grain_to_date=measure_def.get("grain_to_date"),
@@ -413,6 +420,10 @@ def _parse_metric(self, metric_def: dict) -> Metric | None:
             conversion_event=metric_def.get("conversion_event"),
             conversion_window=metric_def.get("conversion_window"),
             offset_window=metric_def.get("offset_window"),
+            cohort_event=metric_def.get("cohort_event"),
+            activity_event=metric_def.get("activity_event"),
+            periods=metric_def.get("periods"),
+            retention_granularity=metric_def.get("granularity") if metric_type == "retention" else None,
             window=metric_def.get("window"),
             grain_to_date=metric_def.get("grain_to_date"),
             window_expression=metric_def.get("window_expression"),
@@ -574,6 +585,15 @@ def _export_model(self, model: Model) -> dict:
                     measure_def["conversion_window"] = measure.conversion_window
                 if measure.offset_window:
                     measure_def["offset_window"] = measure.offset_window
+                # Retention parameters
+                if measure.cohort_event:
+                    measure_def["cohort_event"] = measure.cohort_event
+                if measure.activity_event:
+                    measure_def["activity_event"] = measure.activity_event
+                if measure.periods is not None:
+                    measure_def["periods"] = measure.periods
+                if measure.retention_granularity:
+                    measure_def["granularity"] = measure.retention_granularity
                 # Cumulative/window parameters
                 if measure.window:
                     measure_def["window"] = measure.window
@@ -655,6 +675,14 @@ def _export_metric(self, measure: Metric, graph) -> dict:
             result["conversion_window"] = measure.conversion_window
         if measure.offset_window:
             result["offset_window"] = measure.offset_window
+        if measure.cohort_event:
+            result["cohort_event"] = measure.cohort_event
+        if measure.activity_event:
+            result["activity_event"] = measure.activity_event
+        if measure.periods is not None:
+            result["periods"] = measure.periods
+        if measure.retention_granularity:
+            result["granularity"] = measure.retention_granularity
         if measure.sql:
             result["sql"] = measure.sql
             # Auto-detect and export dependencies for derived measures

sidemantic/core/metric.py

@@ -204,10 +204,15 @@ def validate_type_specific_fields(self):
                 raise ValueError("conversion metric requires 'base_event' field")
             if not self.conversion_event:
                 raise ValueError("conversion metric requires 'conversion_event' field")
+        if self.type == "retention":
+            if not self.entity:
+                raise ValueError("retention metric requires 'entity' field")
+            if not self.cohort_event:
+                raise ValueError("retention metric requires 'cohort_event' field")
         return self
 
     # Metric type (if this is a complex metric, not just a simple aggregation)
-    type: Literal["ratio", "derived", "cumulative", "time_comparison", "conversion"] | None = Field(
+    type: Literal["ratio", "derived", "cumulative", "time_comparison", "conversion", "retention"] | None = Field(
         None, description="Metric type for complex calculations"
     )
 
@@ -252,6 +257,18 @@ def validate_type_specific_fields(self):
     conversion_event: str | None = Field(None, description="Target event filter")
     conversion_window: str | None = Field(None, description="Conversion time window")
 
+    # Retention parameters
+    cohort_event: str | None = Field(
+        None, description="SQL filter for cohort-defining event (e.g., \"event = 'install'\")"
+    )
+    activity_event: str | None = Field(
+        None, description='SQL filter for activity event (default: any event, e.g., "event IS NOT NULL")'
+    )
+    periods: int | None = Field(None, description="Number of retention periods to compute (e.g., 28 for 28-day)")
+    retention_granularity: Literal["day", "week", "month"] | None = Field(
+        None, description="Time granularity for retention periods (day, week, month)"
+    )
+
     # Common parameters
     filters: list[str] | None = Field(None, description="Optional WHERE clause filters")
     fill_nulls_with: int | float | str | None = Field(None, description="Default value when result is NULL")

sidemantic/sql/generator.py

@@ -323,6 +323,9 @@ def metric_needs_window(m):
             # Conversion metrics need special handling
             if metric.type == "conversion":
                 return True
+            # Retention metrics need special handling
+            if metric.type == "retention":
+                return True
             return False
 
         needs_window_functions = any(metric_needs_window(m) for m in metrics)
@@ -2197,6 +2200,170 @@ def resolve_ratio_ref(ref: str) -> str:
 
         raise NotImplementedError(f"Metric type {metric.type} not yet implemented")
 
+    def _generate_retention_query(
+        self,
+        metric_name: str,
+        dimensions: list[str],
+        filters: list[str] | None = None,
+        order_by: list[str] | None = None,
+        limit: int | None = None,
+    ) -> str:
+        """Generate SQL for cohort retention metrics.
+
+        Uses a multi-CTE pattern:
+        1. cohorts: identify each entity's first qualifying event (cohort_date)
+        2. activity: distinct entity activity dates
+        3. retention: join cohorts to activity, compute periods_since and active_users
+        4. cohort_sizes: count of entities per cohort_date
+        5. Final SELECT: retention percentage per cohort_date and period
+
+        Args:
+            metric_name: Name of the retention metric
+            dimensions: List of dimension references (unused for retention, reserved)
+            filters: List of filter expressions
+            order_by: List of fields to order by
+            limit: Maximum number of rows to return
+
+        Returns:
+            SQL query string
+        """
+        import re as _re
+
+        # Resolve metric and model
+        metric = None
+        model = None
+
+        if "." in metric_name:
+            model_name, measure_name = metric_name.split(".", 1)
+            model = self.graph.get_model(model_name)
+            if model:
+                metric = model.get_metric(measure_name)
+        else:
+            try:
+                metric = self.graph.get_metric(metric_name)
+            except KeyError:
+                pass
+
+        if not metric or not metric.entity or not metric.cohort_event:
+            raise ValueError(f"Retention metric {metric_name} missing required fields (entity, cohort_event)")
+
+        # Find the model that owns this metric if not already found
+        if not model:
+            for m_name, m in self.graph.models.items():
+                if m.get_metric(metric_name):
+                    model = m
+                    break
+            if not model:
+                for m_name, m in self.graph.models.items():
+                    for dim in m.dimensions:
+                        if dim.name == metric.entity:
+                            model = m
+                            break
+                    if model:
+                        break
+
+        if not model:
+            raise ValueError(f"No model found for retention metric {metric_name}")
+
+        # Defaults
+        periods = metric.periods or 28
+        granularity = metric.retention_granularity or "day"
+        activity_event = metric.activity_event or "TRUE"
+
+        # Validate entity identifier
+        if not _re.match(r"^[a-zA-Z_][a-zA-Z0-9_.]*$", metric.entity):
+            raise ValueError(f"Invalid entity identifier: {metric.entity}")
+        if not isinstance(periods, int) or periods < 1:
+            raise ValueError(f"Invalid periods value: {periods}")
+
+        # Find timestamp dimension
+        timestamp_dim = None
+        for dim in model.dimensions:
+            if dim.type == "time":
+                timestamp_dim = dim.name
+                break
+
+        if not timestamp_dim:
+            raise ValueError("Retention metrics require a time dimension on the model")
+
+        # Build FROM clause
+        if model.sql:
+            from_clause = f"({model.sql}) AS t"
+        else:
+            from_clause = model.table
+
+        # Build granularity-specific date truncation and interval
+        if granularity == "day":
+            trunc_expr = f"{timestamp_dim}::date"
+            diff_expr = "a.active_date - c.cohort_date"
+            periods_label = "days_since"
+        elif granularity == "week":
+            trunc_expr = f"{self._date_trunc('week', timestamp_dim)}::date"
+            diff_expr = "(a.active_date - c.cohort_date) / 7"
+            periods_label = "weeks_since"
+        elif granularity == "month":
+            trunc_expr = f"{self._date_trunc('month', timestamp_dim)}::date"
+            diff_expr = (
+                "(EXTRACT(YEAR FROM a.active_date) - EXTRACT(YEAR FROM c.cohort_date)) * 12"
+                " + (EXTRACT(MONTH FROM a.active_date) - EXTRACT(MONTH FROM c.cohort_date))"
+            )
+            periods_label = "months_since"
+        else:
+            raise ValueError(f"Unsupported retention granularity: {granularity}")
+
+        # Build optional WHERE filters for the source data
+        filter_clause = ""
+        if filters:
+            filter_clause = " AND " + " AND ".join(filters)
+
+        order_clause = "\nORDER BY r.cohort_date, r.periods_since"
+        if order_by:
+            order_fields = []
+            for field in order_by:
+                field_name = field.split(".", 1)[1] if "." in field else field
+                order_fields.append(field_name)
+            order_clause = f"\nORDER BY {', '.join(order_fields)}"
+
+        limit_clause = ""
+        if limit is not None:
+            limit_clause = f"\nLIMIT {limit}"
+
+        sql = f"""WITH cohorts AS (
+  SELECT {metric.entity}, MIN({trunc_expr}) AS cohort_date
+  FROM {from_clause}
+  WHERE {metric.cohort_event}{filter_clause}
+  GROUP BY {metric.entity}
+),
+activity AS (
+  SELECT DISTINCT {metric.entity}, {trunc_expr} AS active_date
+  FROM {from_clause}
+  WHERE {activity_event}{filter_clause}
+),
+retention AS (
+  SELECT
+    c.cohort_date,
+    CAST({diff_expr} AS INTEGER) AS periods_since,
+    COUNT(DISTINCT c.{metric.entity}) AS active_users
+  FROM cohorts c
+  JOIN activity a ON c.{metric.entity} = a.{metric.entity} AND a.active_date >= c.cohort_date
+  WHERE CAST({diff_expr} AS INTEGER) <= {periods}
+  GROUP BY 1, 2
+),
+cohort_sizes AS (
+  SELECT cohort_date, COUNT(DISTINCT {metric.entity}) AS cohort_size
+  FROM cohorts GROUP BY 1
+)
+SELECT
+  r.cohort_date,
+  r.periods_since AS {periods_label},
+  r.active_users,
+  c.cohort_size,
+  ROUND(r.active_users * 100.0 / c.cohort_size, 1) AS retention_pct
+FROM retention r
+JOIN cohort_sizes c ON r.cohort_date = c.cohort_date{order_clause}{limit_clause}"""
+
+        return sql.strip()
+
     def _generate_conversion_query(
         self,
         metric_name: str,
@@ -2427,6 +2594,7 @@ def _generate_with_window_functions(
         time_comparison_metrics = []
         offset_ratio_metrics = []
         conversion_metrics = []
+        retention_metrics = []
         base_metrics = []
         time_comparison_base_plans = {}
         regular_expression_metric_plans = {}
@@ -2570,6 +2738,9 @@ def collect_leaf_base_metrics(
             elif metric and metric.type == "conversion":
                 add_unique(conversion_metrics, m)
                 # Conversion metrics need special handling - don't add to base_metrics
+            elif metric and metric.type == "retention":
+                add_unique(retention_metrics, m)
+                # Retention metrics need special handling - don't add to base_metrics
             else:
                 # Regular metric or measure
                 metric_ref = canonical_ref(m, resolved_context)
@@ -2596,6 +2767,10 @@ def collect_leaf_base_metrics(
         if precomputed_cumulative_metrics:
             cumulative_metrics = [m for m in cumulative_metrics if m not in precomputed_cumulative_metrics]
 
+        # Handle retention metrics separately - they need a completely different pattern
+        if retention_metrics:
+            return self._generate_retention_query(retention_metrics[0], dimensions, filters, order_by, limit)
+
         # Handle conversion metrics separately - they need a completely different pattern
         if conversion_metrics:
             return self._generate_conversion_query(conversion_metrics[0], dimensions, filters, order_by, limit)