feat(hogql): allow timezone overrides (#21004)

PostHog · Mar 27, 2024 · 8d55545 · 8d55545
1 parent e027202
commit 8d55545
Show file tree

Hide file tree

Showing 3 changed files with 54 additions and 10 deletions.
diff --git a/posthog/hogql/functions/mapping.py b/posthog/hogql/functions/mapping.py
@@ -161,22 +161,22 @@ class HogQLFunctionMeta:
     "toDate": HogQLFunctionMeta(
         "toDateOrNull",
         1,
-        1,
+        2,
         overloads=[((ast.DateTimeType, ast.DateType), "toDate")],
         tz_aware=True,
     ),
     "toDateTime": HogQLFunctionMeta(
         "parseDateTime64BestEffortOrNull",
         1,
-        1,
+        2,
         overloads=[((ast.DateTimeType, ast.DateType, ast.IntegerType), "toDateTime")],
         tz_aware=True,
     ),
     "toUUID": HogQLFunctionMeta("toUUIDOrNull", 1, 1),
     "toString": HogQLFunctionMeta("toString", 1, 1),
     "toJSONString": HogQLFunctionMeta("toJSONString", 1, 1),
-    "parseDateTime": HogQLFunctionMeta("parseDateTimeOrNull", 2, 2, tz_aware=True),
-    "parseDateTimeBestEffort": HogQLFunctionMeta("parseDateTime64BestEffortOrNull", 1, 1, tz_aware=True),
+    "parseDateTime": HogQLFunctionMeta("parseDateTimeOrNull", 2, 3, tz_aware=True),
+    "parseDateTimeBestEffort": HogQLFunctionMeta("parseDateTime64BestEffortOrNull", 1, 2, tz_aware=True),
     # dates and times
     "toTimeZone": HogQLFunctionMeta("toTimeZone", 2, 2),
     "timeZoneOf": HogQLFunctionMeta("timeZoneOf", 1, 1),
@@ -219,8 +219,8 @@ class HogQLFunctionMeta:
     "dateSub": HogQLFunctionMeta("dateSub", 3, 3),
     "timeStampAdd": HogQLFunctionMeta("timeStampAdd", 2, 2),
     "timeStampSub": HogQLFunctionMeta("timeStampSub", 2, 2),
-    "now": HogQLFunctionMeta("now64", tz_aware=True),
-    "NOW": HogQLFunctionMeta("now64", tz_aware=True),
+    "now": HogQLFunctionMeta("now64", 0, 1, tz_aware=True),
+    "NOW": HogQLFunctionMeta("now64", 0, 1, tz_aware=True),
     "nowInBlock": HogQLFunctionMeta("nowInBlock", 1, 1),
     "today": HogQLFunctionMeta("today"),
     "yesterday": HogQLFunctionMeta("yesterday"),

diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py
@@ -822,11 +822,21 @@ def visit_call(self, node: ast.Call):
                                 break  # Found an overload matching the first function org
 
                 if func_meta.tz_aware:
-                    if (relevant_clickhouse_name == "now64" and len(node.args) == 0) or (
-                        relevant_clickhouse_name == "parseDateTime64BestEffortOrNull" and len(node.args) == 1
+                    has_tz_override = len(node.args) == func_meta.max_args
+
+                    if not has_tz_override:
+                        args.append(self.visit(ast.Constant(value=self._get_timezone())))
+
+                    if (
+                        relevant_clickhouse_name == "now64"
+                        and (len(node.args) == 0 or (has_tz_override and len(node.args) == 1))
+                    ) or (
+                        relevant_clickhouse_name == "parseDateTime64BestEffortOrNull"
+                        and (len(node.args) == 1 or (has_tz_override and len(node.args) == 2))
                     ):
-                        args.append("6")  # These two CH functions require the precision argument before timezone
-                    args.append(self.visit(ast.Constant(value=self._get_timezone())))
+                        # These two CH functions require a precision argument before timezone
+                        args = args[:-1] + ["6"] + args[-1:]
+
                 if node.name == "toStartOfWeek" and len(node.args) == 1:
                     # If week mode hasn't been specified, use the project's default.
                     # For Monday-based weeks mode 3 is used (which is ISO 8601), for Sunday-based mode 0 (CH default)

diff --git a/posthog/hogql/test/test_printer.py b/posthog/hogql/test/test_printer.py
@@ -1524,3 +1524,37 @@ def test_lookup_organic_medium_type(self):
             ),
             printed,
         )
+
+    def test_override_timezone(self):
+        context = HogQLContext(
+            team_id=self.team.pk,
+            enable_select_queries=True,
+            database=Database(None, WeekStartDay.SUNDAY),
+        )
+        context.database.events.fields["test_date"] = DateDatabaseField(name="test_date")  # type: ignore
+
+        self.assertEqual(
+            self._select(
+                """
+                    SELECT
+                        toDateTime(timestamp),
+                        toDateTime(timestamp, 'US/Pacific'),
+                        now(),
+                        now('US/Pacific')
+                    FROM events
+                """,
+                context,
+            ),
+            f"SELECT toDateTime(toTimeZone(events.timestamp, %(hogql_val_0)s), %(hogql_val_1)s), toDateTime(toTimeZone(events.timestamp, %(hogql_val_2)s), %(hogql_val_3)s), now64(6, %(hogql_val_4)s), now64(6, %(hogql_val_5)s) FROM events WHERE equals(events.team_id, {self.team.pk}) LIMIT 10000",
+        )
+        self.assertEqual(
+            context.values,
+            {
+                "hogql_val_0": "UTC",
+                "hogql_val_1": "UTC",
+                "hogql_val_2": "UTC",
+                "hogql_val_3": "US/Pacific",
+                "hogql_val_4": "UTC",
+                "hogql_val_5": "US/Pacific",
+            },
+        )