fix(web-analytics): handle multi-set partition correctly (#35231)

lricoy · web-flow · commit a4c6ea181148 · 2025-07-18T13:33:17.000-04:00
diff --git a/dags/tests/test_web_preaggregated_partitions.py b/dags/tests/test_web_preaggregated_partitions.py
@@ -0,0 +1,147 @@
+import pytest
+from datetime import datetime, UTC
+from unittest.mock import Mock, patch, call
+
+import dagster
+from dagster import TimeWindow
+
+from dags.web_preaggregated_daily import (
+    pre_aggregate_web_analytics_data,
+)
+from posthog.models.web_preaggregated.sql import DROP_PARTITION_SQL
+
+
+class TestPartitionHandling:
+    def setup_method(self):
+        self.mock_context = Mock(spec=dagster.AssetExecutionContext)
+        self.mock_context.log = Mock()
+        self.mock_context.op_config = {"team_ids": [1, 2], "extra_clickhouse_settings": ""}
+
+    @pytest.mark.parametrize(
+        "start_date_str,end_date_str,expected_partitions",
+        [
+            # Single day partition
+            ("2024-01-01", "2024-01-02", ["2024-01-01"]),
+            # Two day partition
+            ("2024-01-01", "2024-01-03", ["2024-01-01", "2024-01-02"]),
+            # Week-long partition
+            (
+                "2024-01-01",
+                "2024-01-08",
+                ["2024-01-01", "2024-01-02", "2024-01-03", "2024-01-04", "2024-01-05", "2024-01-06", "2024-01-07"],
+            ),
+            # Month-long partition (testing edge case)
+            ("2024-01-01", "2024-02-01", [f"2024-01-{day:02d}" for day in range(1, 32)]),
+        ],
+    )
+    @patch("dags.web_preaggregated_daily.sync_execute")
+    def test_partition_dropping_for_different_time_windows(
+        self, mock_sync_execute, start_date_str, end_date_str, expected_partitions
+    ):
+        start_datetime = datetime.fromisoformat(start_date_str).replace(tzinfo=UTC)
+        end_datetime = datetime.fromisoformat(end_date_str).replace(tzinfo=UTC)
+        self.mock_context.partition_time_window = TimeWindow(start_datetime, end_datetime)
+
+        mock_sql_generator = Mock(return_value="INSERT INTO test_table ...")
+        pre_aggregate_web_analytics_data(
+            context=self.mock_context,
+            table_name="web_stats_daily",
+            sql_generator=mock_sql_generator,
+        )
+
+        actual_drop_calls = [
+            call_args for call_args in mock_sync_execute.call_args_list if "DROP PARTITION" in str(call_args)
+        ]
+        assert len(actual_drop_calls) == len(expected_partitions)
+
+    @patch("dags.web_preaggregated_daily.sync_execute")
+    def test_partition_drop_error_handling(self, mock_sync_execute):
+        start_datetime = datetime(2024, 1, 1, tzinfo=UTC)
+        end_datetime = datetime(2024, 1, 3, tzinfo=UTC)
+        self.mock_context.partition_time_window = TimeWindow(start_datetime, end_datetime)
+
+        def side_effect(sql):
+            if "DROP PARTITION" in sql:
+                raise Exception("Partition doesn't exist")
+            return None
+
+        mock_sync_execute.side_effect = side_effect
+        mock_sql_generator = Mock(return_value="INSERT INTO test_table ...")
+
+        pre_aggregate_web_analytics_data(
+            context=self.mock_context,
+            table_name="web_stats_daily",
+            sql_generator=mock_sql_generator,
+        )
+
+        assert self.mock_context.log.info.call_count >= 4
+        insert_calls = [call_args for call_args in mock_sync_execute.call_args_list if "INSERT INTO" in str(call_args)]
+        assert len(insert_calls) == 1
+
+    @patch("dags.web_preaggregated_daily.sync_execute")
+    def test_granularity_parameter_usage(self, mock_sync_execute):
+        start_datetime = datetime(2024, 1, 1, tzinfo=UTC)
+        end_datetime = datetime(2024, 1, 2, tzinfo=UTC)
+        self.mock_context.partition_time_window = TimeWindow(start_datetime, end_datetime)
+
+        mock_sql_generator = Mock(return_value="INSERT INTO test_table ...")
+        pre_aggregate_web_analytics_data(
+            context=self.mock_context,
+            table_name="web_stats_daily",
+            sql_generator=mock_sql_generator,
+        )
+
+        expected_sql = DROP_PARTITION_SQL("web_stats_daily", "2024-01-01", granularity="daily")
+        assert call(expected_sql) in mock_sync_execute.call_args_list
+        assert "'20240101'" in expected_sql
+
+    def test_missing_partition_time_window_raises_error(self):
+        self.mock_context.partition_time_window = None
+        mock_sql_generator = Mock()
+
+        with pytest.raises(dagster.Failure, match="This asset should only be run with a partition_time_window"):
+            pre_aggregate_web_analytics_data(
+                context=self.mock_context,
+                table_name="web_stats_daily",
+                sql_generator=mock_sql_generator,
+            )
+
+    @patch("dags.web_preaggregated_daily.sync_execute")
+    def test_insert_query_failure_raises_dagster_failure(self, mock_sync_execute):
+        start_datetime = datetime(2024, 1, 1, tzinfo=UTC)
+        end_datetime = datetime(2024, 1, 2, tzinfo=UTC)
+        self.mock_context.partition_time_window = TimeWindow(start_datetime, end_datetime)
+
+        def side_effect(sql):
+            if "INSERT INTO" in sql:
+                raise Exception("Insert failed")
+            return None
+
+        mock_sync_execute.side_effect = side_effect
+        mock_sql_generator = Mock(return_value="INSERT INTO test_table ...")
+
+        with pytest.raises(dagster.Failure, match="Failed to pre-aggregate web_stats_daily"):
+            pre_aggregate_web_analytics_data(
+                context=self.mock_context,
+                table_name="web_stats_daily",
+                sql_generator=mock_sql_generator,
+            )
+
+    @patch("dags.web_preaggregated_daily.sync_execute")
+    def test_same_start_and_end_date_drops_partition(self, mock_sync_execute):
+        start_datetime = datetime(2024, 1, 1, tzinfo=UTC)
+        end_datetime = datetime(2024, 1, 1, tzinfo=UTC)
+        self.mock_context.partition_time_window = TimeWindow(start_datetime, end_datetime)
+
+        mock_sql_generator = Mock(return_value="INSERT INTO test_table ...")
+        pre_aggregate_web_analytics_data(
+            context=self.mock_context,
+            table_name="web_stats_daily",
+            sql_generator=mock_sql_generator,
+        )
+
+        drop_calls = [call_args for call_args in mock_sync_execute.call_args_list if "DROP PARTITION" in str(call_args)]
+        assert len(drop_calls) == 1
+
+        insert_calls = [call_args for call_args in mock_sync_execute.call_args_list if "INSERT INTO" in str(call_args)]
+        assert len(insert_calls) == 1
diff --git a/dags/web_preaggregated_daily.py b/dags/web_preaggregated_daily.py
@@ -71,17 +71,27 @@ def pre_aggregate_web_analytics_data(
     date_end = end_datetime.strftime("%Y-%m-%d")
 
     try:
-        # Drop the partition first, ensuring a clean state before insertion
+        # Drop all partitions in the time window, ensuring a clean state before insertion
         # Note: No ON CLUSTER needed since tables are replicated (not sharded) and replication handles distribution
-        drop_partition_query = DROP_PARTITION_SQL(table_name, date_start)
-        context.log.info(f"Dropping partition for {date_start}: {drop_partition_query}")
+        current_date = start_datetime.date()
+        end_date = end_datetime.date()
 
-        try:
-            sync_execute(drop_partition_query)
-            context.log.info(f"Successfully dropped partition for {date_start}")
-        except Exception as drop_error:
-            # Partition might not exist when running for the first time or when running in a empty backfill, which is fine
-            context.log.info(f"Partition for {date_start} doesn't exist or couldn't be dropped: {drop_error}")
+        # For time windows: start is inclusive, end is exclusive (except for single-day partitions)
+        while current_date < end_date or (current_date == start_datetime.date() == end_date):
+            partition_date_str = current_date.strftime("%Y-%m-%d")
+            drop_partition_query = DROP_PARTITION_SQL(table_name, partition_date_str, granularity="daily")
+            context.log.info(f"Dropping partition for {partition_date_str}: {drop_partition_query}")
+
+            try:
+                sync_execute(drop_partition_query)
+                context.log.info(f"Successfully dropped partition for {partition_date_str}")
+            except Exception as drop_error:
+                # Partition might not exist when running for the first time or when running in a empty backfill, which is fine
+                context.log.info(
+                    f"Partition for {partition_date_str} doesn't exist or couldn't be dropped: {drop_error}"
+                )
+
+            current_date += timedelta(days=1)
 
         insert_query = sql_generator(
             date_start=date_start,
diff --git a/posthog/models/web_preaggregated/test_web_preaggregated_sql.py b/posthog/models/web_preaggregated/test_web_preaggregated_sql.py
@@ -145,3 +145,50 @@ def test_hourly_vs_daily_partition_difference(self):
         assert "'20240115'" in daily_sql
         # Hourly should be YYYYMMDDHH format (defaulting to 00)
         assert "'2024011500'" in hourly_sql
+
+    def test_multi_day_partition_scenario_daily(self):
+        dates = ["2024-01-01", "2024-01-02", "2024-01-03"]
+
+        for date in dates:
+            sql = DROP_PARTITION_SQL("web_stats_daily", date, granularity="daily")
+            expected_partition = date.replace("-", "")
+            assert f"'{expected_partition}'" in sql
+            assert "ALTER TABLE web_stats_daily" in sql
+            assert "DROP PARTITION" in sql
+
+    def test_month_boundary_partitions(self):
+        test_cases = [
+            ("2024-01-31", "daily", "20240131"),
+            ("2024-02-01", "daily", "20240201"),
+            ("2024-02-29", "daily", "20240229"),  # Leap year
+            ("2024-03-01", "daily", "20240301"),
+        ]
+
+        for date, granularity, expected_partition in test_cases:
+            sql = DROP_PARTITION_SQL("web_stats_daily", date, granularity=granularity)
+            assert f"'{expected_partition}'" in sql
+
+    def test_year_boundary_partitions(self):
+        test_cases = [
+            ("2023-12-31", "daily", "20231231"),
+            ("2024-01-01", "daily", "20240101"),
+        ]
+
+        for date, granularity, expected_partition in test_cases:
+            sql = DROP_PARTITION_SQL("web_stats_daily", date, granularity=granularity)
+            assert f"'{expected_partition}'" in sql
+
+    def test_granularity_parameter_is_case_sensitive_and_invalid_defaults_to_daily(self):
+        date = "2024-01-15"
+
+        # Test valid granularities
+        daily_sql = DROP_PARTITION_SQL("test_table", date, granularity="daily")
+        hourly_sql = DROP_PARTITION_SQL("test_table", date, granularity="hourly")
+
+        assert "'20240115'" in daily_sql
+        assert "'2024011500'" in hourly_sql
+
+        # Test that invalid granularity defaults to daily behavior
+        # (This tests the else clause in the function)
+        invalid_sql = DROP_PARTITION_SQL("test_table", date, granularity="invalid")
+        assert "'20240115'" in invalid_sql
