Add missing logic for action days in cohort calculating

ee/clickhouse/models/cohort.py

@@ -1,7 +1,9 @@
 import uuid
-from datetime import datetime
+from datetime import datetime, timedelta
 from typing import Any, Dict, List, Optional, Tuple
 
+from django.utils import timezone
+
 from ee.clickhouse.client import sync_execute
 from ee.clickhouse.models.action import format_action_filter
 from ee.clickhouse.sql.cohort import CALCULATE_COHORT_PEOPLE_SQL
@@ -30,10 +32,16 @@ def format_person_query(cohort: Cohort) -> Tuple[str, Dict[str, Any]]:
         if group.get("action_id"):
             action = Action.objects.get(pk=group["action_id"], team_id=cohort.team.pk)
             action_filter_query, action_params = format_action_filter(action, prepend="_{}_action".format(group_idx))
-            extract_person = "SELECT distinct_id FROM events WHERE team_id = %(team_id)s AND {query}".format(
-                query=action_filter_query
+
+            date_query: str = ""
+            date_params: Dict[str, str] = {}
+            if group.get("days"):
+                date_query, date_params = parse_action_timestamps(int(group.get("days")))
+
+            extract_person = "SELECT distinct_id FROM events WHERE team_id = %(team_id)s {date_query} AND {query}".format(
+                query=action_filter_query, date_query=date_query
             )
-            params = {**params, **action_params}
+            params = {**params, **action_params, **date_params}
             filters.append("distinct_id IN (" + extract_person + ")")
 
         elif group.get("properties"):
@@ -53,6 +61,16 @@ def format_person_query(cohort: Cohort) -> Tuple[str, Dict[str, Any]]:
     return joined_filter, params
 
 
+def parse_action_timestamps(days: int) -> Tuple[str, Dict[str, str]]:
+    curr_time = timezone.now()
+    start_time = curr_time - timedelta(days=days)
+
+    return (
+        "and timestamp >= %(date_from)s AND timestamp <= %(date_to)s",
+        {"date_from": start_time.strftime("%Y-%m-%d %H:%M:%S"), "date_to": curr_time.strftime("%Y-%m-%d %H:%M:%S")},
+    )
+
+
 def format_filter_query(cohort: Cohort) -> Tuple[str, Dict[str, Any]]:
     person_query, params = format_person_query(cohort)
     person_id_query = CALCULATE_COHORT_PEOPLE_SQL.format(query=person_query)

ee/clickhouse/models/test/test_cohort.py

@@ -1,6 +1,8 @@
 from datetime import datetime
 from uuid import uuid4
 
+from freezegun import freeze_time
+
 from ee.clickhouse.client import sync_execute
 from ee.clickhouse.models.cohort import format_filter_query, get_person_ids_by_cohort_id
 from ee.clickhouse.models.event import create_event
@@ -106,6 +108,54 @@ def test_prop_cohort_basic_action(self):
         result = sync_execute(final_query, {**params, "team_id": self.team.pk})
         self.assertEqual(len(result), 1)
 
+    def test_prop_cohort_basic_action_days(self):
+
+        _create_person(distinct_ids=["some_other_id"], team_id=self.team.pk, properties={"$some_prop": "something"})
+
+        _create_person(
+            distinct_ids=["some_id"],
+            team_id=self.team.pk,
+            properties={"$some_prop": "something", "$another_prop": "something"},
+        )
+
+        action = _create_action(team=self.team, name="$pageview")
+        _create_event(
+            event="$pageview",
+            team=self.team,
+            distinct_id="some_id",
+            properties={"attr": "some_val"},
+            timestamp=datetime(2020, 1, 9, 12, 0, 1),
+        )
+
+        _create_event(
+            event="$pageview",
+            team=self.team,
+            distinct_id="some_other_id",
+            properties={"attr": "some_val"},
+            timestamp=datetime(2020, 1, 5, 12, 0, 1),
+        )
+
+        with freeze_time("2020-01-10"):
+            cohort1 = Cohort.objects.create(
+                team=self.team, groups=[{"action_id": action.pk, "days": 1}], name="cohort1",
+            )
+
+            filter = Filter(data={"properties": [{"key": "id", "value": cohort1.pk, "type": "cohort"}],})
+            query, params = parse_prop_clauses(filter.properties, self.team.pk)
+            final_query = "SELECT uuid FROM events WHERE team_id = %(team_id)s {}".format(query)
+            result = sync_execute(final_query, {**params, "team_id": self.team.pk})
+            self.assertEqual(len(result), 1)
+
+            cohort2 = Cohort.objects.create(
+                team=self.team, groups=[{"action_id": action.pk, "days": 7}], name="cohort2",
+            )
+
+            filter = Filter(data={"properties": [{"key": "id", "value": cohort2.pk, "type": "cohort"}],})
+            query, params = parse_prop_clauses(filter.properties, self.team.pk)
+            final_query = "SELECT uuid FROM events WHERE team_id = %(team_id)s {}".format(query)
+            result = sync_execute(final_query, {**params, "team_id": self.team.pk})
+            self.assertEqual(len(result), 2)
+
     def test_prop_cohort_multiple_groups(self):
 
         _create_person(distinct_ids=["some_other_id"], team_id=self.team.pk, properties={"$some_prop": "something"})