fix(exports): format csv with formula and breakdown (#40607)

Author: andyzzhao

posthog/tasks/exports/csv_exporter.py

@@ -22,6 +22,12 @@
 from ...exceptions import QuerySizeExceeded
 from ...hogql.constants import CSV_EXPORT_BREAKDOWN_LIMIT_INITIAL, CSV_EXPORT_BREAKDOWN_LIMIT_LOW, CSV_EXPORT_LIMIT
 from ...hogql.query import LimitContext
+from ...hogql_queries.insights.trends.breakdown import (
+    BREAKDOWN_NULL_DISPLAY,
+    BREAKDOWN_NULL_STRING_LABEL,
+    BREAKDOWN_OTHER_DISPLAY,
+    BREAKDOWN_OTHER_STRING_LABEL,
+)
 from ..exporter import EXPORT_ASSET_UNKNOWN_COUNTER, EXPORT_FAILED_COUNTER, EXPORT_SUCCEEDED_COUNTER, EXPORT_TIMER
 from .ordered_csv_renderer import OrderedCsvRenderer
 
@@ -78,7 +84,7 @@ def add_query_params(url: str, params: dict[str, str]) -> str:
     return urlunparse(parsed)
 
 
-def _convert_response_to_csv_data(data: Any) -> Generator[Any, None, None]:
+def _convert_response_to_csv_data(data: Any, breakdown_filter: Optional[dict] = None) -> Generator[Any, None, None]:
     if isinstance(data.get("results"), list):
         results = data.get("results")
         if len(results) > 0 and (isinstance(results[0], list) or isinstance(results[0], tuple)) and data.get("types"):
@@ -166,7 +172,9 @@ def _convert_response_to_csv_data(data: Any) -> Generator[Any, None, None]:
 
                 yield line
             return
-        elif isinstance(first_result.get("data"), list):
+        elif isinstance(first_result.get("data"), list) or (
+            first_result.get("data") is None and "aggregated_value" in first_result
+        ):
             is_comparison = first_result.get("compare_label")
 
             # take date labels from current results, when comparing against previous
@@ -186,9 +194,33 @@ def _convert_response_to_csv_data(data: Any) -> Generator[Any, None, None]:
 
                 if isinstance(action, dict) and action.get("custom_name"):
                     line["custom name"] = action.get("custom_name")
-                if item.get("aggregated_value"):
-                    line["total count"] = item.get("aggregated_value")
-                else:
+
+                if "breakdown_value" in item:
+                    breakdown_value = item.get("breakdown_value")
+                    breakdown_values = breakdown_value if isinstance(breakdown_value, list) else [breakdown_value]
+
+                    # Get breakdown property names from filter
+                    breakdowns = breakdown_filter.get("breakdowns", []) if breakdown_filter else []
+                    # For single breakdown, check legacy "breakdown" field
+                    if not breakdowns and breakdown_filter and "breakdown" in breakdown_filter:
+                        breakdowns = [{"property": breakdown_filter.get("breakdown")}]
+
+                    for idx, val in enumerate(breakdown_values):
+                        # Get the property name from the breakdown filter
+                        prop_name = breakdowns[idx].get("property") if idx < len(breakdowns) else None
+                        if not prop_name:
+                            continue
+                        # Format special breakdown values for display
+                        formatted_val = str(val) if val is not None else ""
+                        if formatted_val == BREAKDOWN_OTHER_STRING_LABEL:
+                            formatted_val = BREAKDOWN_OTHER_DISPLAY
+                        elif formatted_val == BREAKDOWN_NULL_STRING_LABEL:
+                            formatted_val = BREAKDOWN_NULL_DISPLAY
+                        line[prop_name] = formatted_val
+
+                if item.get("aggregated_value") is not None:
+                    line["Total Sum"] = item.get("aggregated_value")
+                elif item.get("data"):
                     for index, data in enumerate(item["data"]):
                         line[label_item["labels"][index]] = data
 
@@ -289,7 +321,9 @@ def get_from_hogql_query(exported_asset: ExportedAsset, limit: int, resource: di
 
         if isinstance(query_response, BaseModel):
             query_response = query_response.model_dump(by_alias=True)
-        yield from _convert_response_to_csv_data(query_response)
+
+        breakdown_filter = query.get("breakdownFilter") if query else None
+        yield from _convert_response_to_csv_data(query_response, breakdown_filter=breakdown_filter)
         return
 
 

posthog/tasks/exports/test/test_csv_exporter.py

@@ -807,13 +807,13 @@ def test_csv_exporter_trends_query_with_compare_previous_option(
             lines = (content or "").strip().splitlines()
 
             expected_lines = [
-                "series,21-Mar-2023,22-Mar-2023,23-Mar-2023,24-Mar-2023,25-Mar-2023,26-Mar-2023,27-Mar-2023,28-Mar-2023",
-                "Chrome - current,2.0,0.0,0.0,0.0,0.0,0.0,0.0,3.0",
-                "Firefox - current,1.0,0.0,0.0,0.0,0.0,0.0,0.0,2.0",
-                "Safari - current,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0",
-                "Chrome - previous,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2.0",
-                "Firefox - previous,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0",
-                "Safari - previous,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0",
+                "series,$browser,21-Mar-2023,22-Mar-2023,23-Mar-2023,24-Mar-2023,25-Mar-2023,26-Mar-2023,27-Mar-2023,28-Mar-2023",
+                "Chrome - current,Chrome,2.0,0.0,0.0,0.0,0.0,0.0,0.0,3.0",
+                "Firefox - current,Firefox,1.0,0.0,0.0,0.0,0.0,0.0,0.0,2.0",
+                "Safari - current,Safari,1.0,0.0,0.0,0.0,0.0,0.0,0.0,0.0",
+                "Chrome - previous,Chrome,0.0,0.0,0.0,0.0,0.0,0.0,0.0,2.0",
+                "Firefox - previous,Firefox,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0",
+                "Safari - previous,Safari,0.0,0.0,0.0,0.0,0.0,0.0,0.0,1.0",
             ]
 
             self.assertEqual(lines, expected_lines)
@@ -885,3 +885,231 @@ def test_csv_exporter_trends_actors(
                     "d0780d6b-ccd0-44fa-a227-47efe4f3f30d,,,user_2,1,user_2",
                 ],
             )
+
+    @patch("posthog.models.exported_asset.UUIDT")
+    def test_csv_exporter_trends_query_with_formula(
+        self, mocked_uuidt: Any, MAX_SELECT_RETURNED_ROWS: int = 10
+    ) -> None:
+        with freeze_time("2024-05-15T12:00:00.000Z"):
+            _create_person(distinct_ids=["formula_test_user_xyz"], team=self.team)
+
+        events_by_person = {
+            "formula_test_user_xyz": [
+                {"event": "formula_test_event_a", "timestamp": datetime(2024, 5, 15, 13, 46)},
+                {"event": "formula_test_event_b", "timestamp": datetime(2024, 5, 15, 13, 47)},
+            ],
+        }
+        journeys_for(events_by_person, self.team)
+        flush_persons_and_events()
+
+        exported_asset = ExportedAsset(
+            team=self.team,
+            export_format=ExportedAsset.ExportFormat.CSV,
+            export_context={
+                "source": {
+                    "kind": "TrendsQuery",
+                    "dateRange": {"date_to": "2024-05-15", "date_from": "2024-05-15"},
+                    "series": [
+                        {"kind": "EventsNode", "event": "formula_test_event_a", "name": "Event A", "math": "total"},
+                        {"kind": "EventsNode", "event": "formula_test_event_b", "name": "Event B", "math": "total"},
+                    ],
+                    "interval": "day",
+                    "trendsFilter": {
+                        "showLegend": True,
+                        "display": "ActionsTable",
+                        "formula": "A+B",
+                    },
+                }
+            },
+        )
+        exported_asset.save()
+        mocked_uuidt.return_value = "a-guid"
+
+        with self.settings(OBJECT_STORAGE_ENABLED=True, OBJECT_STORAGE_EXPORTS_FOLDER="Test-Exports"):
+            csv_exporter.export_tabular(exported_asset)
+
+            assert (
+                exported_asset.content_location
+                == f"{TEST_PREFIX}/csv/team-{self.team.id}/task-{exported_asset.id}/a-guid"
+            )
+
+            content = object_storage.read(exported_asset.content_location)
+            lines = (content or "").strip().split("\r\n")
+            self.assertEqual(
+                lines,
+                [
+                    "series,Total Sum",
+                    "Formula (A+B),2.0",
+                ],
+            )
+
+    @patch("posthog.models.exported_asset.UUIDT")
+    def test_csv_exporter_trends_query_with_formula_and_single_breakdown(
+        self, mocked_uuidt: Any, MAX_SELECT_RETURNED_ROWS: int = 10
+    ) -> None:
+        with freeze_time("2024-06-10T12:00:00.000Z"):
+            _create_person(distinct_ids=["breakdown_user_single"], team=self.team)
+
+        _create_event(
+            event="breakdown_single_event_a",
+            distinct_id="breakdown_user_single",
+            team=self.team,
+            timestamp=datetime(2024, 6, 10, 13, 46),
+            properties={"country": "USA"},
+        )
+        _create_event(
+            event="breakdown_single_event_b",
+            distinct_id="breakdown_user_single",
+            team=self.team,
+            timestamp=datetime(2024, 6, 10, 13, 47),
+            properties={"country": "USA"},
+        )
+        flush_persons_and_events()
+
+        exported_asset = ExportedAsset(
+            team=self.team,
+            export_format=ExportedAsset.ExportFormat.CSV,
+            export_context={
+                "source": {
+                    "kind": "TrendsQuery",
+                    "dateRange": {"date_to": "2024-06-10", "date_from": "2024-06-10"},
+                    "series": [
+                        {
+                            "kind": "EventsNode",
+                            "event": "breakdown_single_event_a",
+                            "name": "Event A",
+                            "math": "total",
+                        },
+                        {
+                            "kind": "EventsNode",
+                            "event": "breakdown_single_event_b",
+                            "name": "Event B",
+                            "math": "total",
+                        },
+                    ],
+                    "interval": "day",
+                    "trendsFilter": {
+                        "showLegend": True,
+                        "display": "ActionsTable",
+                        "formula": "A+B",
+                    },
+                    "breakdownFilter": {
+                        "breakdown": "country",
+                        "breakdown_type": "event",
+                    },
+                }
+            },
+        )
+        exported_asset.save()
+        mocked_uuidt.return_value = "a-guid"
+
+        with self.settings(OBJECT_STORAGE_ENABLED=True, OBJECT_STORAGE_EXPORTS_FOLDER="Test-Exports"):
+            csv_exporter.export_tabular(exported_asset)
+
+            assert (
+                exported_asset.content_location
+                == f"{TEST_PREFIX}/csv/team-{self.team.id}/task-{exported_asset.id}/a-guid"
+            )
+
+            content = object_storage.read(exported_asset.content_location)
+            lines = (content or "").strip().split("\r\n")
+            self.assertEqual(
+                lines,
+                [
+                    "series,country,Total Sum",
+                    "Formula (A+B),USA,2.0",
+                ],
+            )
+
+    @patch("posthog.models.exported_asset.UUIDT")
+    def test_csv_exporter_trends_query_with_formula_and_multiple_breakdowns(
+        self, mocked_uuidt: Any, MAX_SELECT_RETURNED_ROWS: int = 10
+    ) -> None:
+        with freeze_time("2024-07-20T12:00:00.000Z"):
+            _create_person(distinct_ids=["multi_breakdown_user_1"], team=self.team)
+            _create_person(distinct_ids=["multi_breakdown_user_2"], team=self.team)
+
+        _create_event(
+            event="multi_breakdown_event_a",
+            distinct_id="multi_breakdown_user_1",
+            team=self.team,
+            timestamp=datetime(2024, 7, 20, 13, 46),
+            properties={"$browser": "Chrome"},
+        )
+        _create_event(
+            event="multi_breakdown_event_b",
+            distinct_id="multi_breakdown_user_1",
+            team=self.team,
+            timestamp=datetime(2024, 7, 20, 13, 47),
+            properties={"$browser": "Chrome"},
+        )
+        _create_event(
+            event="multi_breakdown_event_a",
+            distinct_id="multi_breakdown_user_2",
+            team=self.team,
+            timestamp=datetime(2024, 7, 20, 13, 48),
+            properties={"$browser": "Firefox"},
+        )
+        flush_persons_and_events()
+
+        exported_asset = ExportedAsset(
+            team=self.team,
+            export_format=ExportedAsset.ExportFormat.CSV,
+            export_context={
+                "source": {
+                    "kind": "TrendsQuery",
+                    "dateRange": {"date_to": "2024-07-20", "date_from": "2024-07-20"},
+                    "series": [
+                        {
+                            "kind": "EventsNode",
+                            "event": "multi_breakdown_event_a",
+                            "name": "Event A",
+                            "math": "total",
+                        },
+                        {
+                            "kind": "EventsNode",
+                            "event": "multi_breakdown_event_b",
+                            "name": "Event B",
+                            "math": "total",
+                        },
+                    ],
+                    "interval": "day",
+                    "trendsFilter": {
+                        "showLegend": True,
+                        "display": "ActionsTable",
+                        "formula": "A+B",
+                    },
+                    "breakdownFilter": {
+                        "breakdowns": [
+                            {"property": "distinct_id", "type": "event_metadata"},
+                            {"property": "$browser", "type": "event"},
+                        ]
+                    },
+                }
+            },
+        )
+        exported_asset.save()
+        mocked_uuidt.return_value = "a-guid"
+
+        with self.settings(OBJECT_STORAGE_ENABLED=True, OBJECT_STORAGE_EXPORTS_FOLDER="Test-Exports"):
+            csv_exporter.export_tabular(exported_asset)
+
+            assert (
+                exported_asset.content_location
+                == f"{TEST_PREFIX}/csv/team-{self.team.id}/task-{exported_asset.id}/a-guid"
+            )
+
+            content = object_storage.read(exported_asset.content_location)
+            lines = (content or "").strip().split("\r\n")
+
+            # Sort data lines for consistent comparison (order may vary)
+            data_lines = sorted(lines[1:])
+
+            self.assertEqual(
+                lines[0:1] + data_lines,
+                [
+                    "series,distinct_id,$browser,Total Sum",
+                    "Formula (A+B),multi_breakdown_user_1,Chrome,2.0",
+                    "Formula (A+B),multi_breakdown_user_2,Firefox,1.0",
+                ],
+            )