add override for functions with precision before tz

PostHog · Mar 19, 2024 · 78f3a41 · 78f3a41
1 parent 612bf95
commit 78f3a41
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 6 deletions.
diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py
@@ -823,12 +823,20 @@ def visit_call(self, node: ast.Call):
 
                 if func_meta.tz_aware:
                     has_tz_override = len(node.args) == func_meta.max_args
-                    if (relevant_clickhouse_name == "now64" and len(node.args) == 0) or (
-                        relevant_clickhouse_name == "parseDateTime64BestEffortOrNull" and len(node.args) == 1
-                    ):
-                        args.append("6")  # These two CH functions require the precision argument before timezone
+
                     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))
+                    ):
+                        # 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
@@ -1538,12 +1538,14 @@ def test_override_timezone(self):
                 """
                     SELECT
                         toDateTime(timestamp),
-                        toDateTime(timestamp, 'US/Pacific')
+                        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) FROM events WHERE equals(events.team_id, {self.team.pk}) LIMIT 10000",
+            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,
@@ -1552,5 +1554,7 @@ def test_override_timezone(self):
                 "hogql_val_1": "UTC",
                 "hogql_val_2": "UTC",
                 "hogql_val_3": "US/Pacific",
+                "hogql_val_4": "UTC",
+                "hogql_val_5": "US/Pacific",
             },
         )