From 612bf9569da95351b7838e19ffa5add903e1a10c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Obermu=CC=88ller?= Date: Tue, 19 Mar 2024 16:07:46 +0100 Subject: [PATCH 1/2] add override for toDateTime --- posthog/hogql/functions/mapping.py | 12 ++++++------ posthog/hogql/printer.py | 4 +++- posthog/hogql/test/test_printer.py | 30 ++++++++++++++++++++++++++++++ 3 files changed, 39 insertions(+), 7 deletions(-) diff --git a/posthog/hogql/functions/mapping.py b/posthog/hogql/functions/mapping.py index 5edf1a68a826a..cd908d725341d 100644 --- 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 index db6523cfa4078..2731897748c99 100644 --- a/posthog/hogql/printer.py +++ b/posthog/hogql/printer.py @@ -822,11 +822,13 @@ def visit_call(self, node: ast.Call): break # Found an overload matching the first function org 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 - args.append(self.visit(ast.Constant(value=self._get_timezone()))) + if not has_tz_override: + args.append(self.visit(ast.Constant(value=self._get_timezone()))) 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 index 9bbf4c40aff98..ad74a0ca87648 100644 --- a/posthog/hogql/test/test_printer.py +++ b/posthog/hogql/test/test_printer.py @@ -1524,3 +1524,33 @@ 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') + 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", + ) + self.assertEqual( + context.values, + { + "hogql_val_0": "UTC", + "hogql_val_1": "UTC", + "hogql_val_2": "UTC", + "hogql_val_3": "US/Pacific", + }, + ) From 78f3a4195834aec7ccdd66bdb5deb961d0c6f6cc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Thomas=20Obermu=CC=88ller?= Date: Tue, 19 Mar 2024 16:28:57 +0100 Subject: [PATCH 2/2] add override for functions with precision before tz --- posthog/hogql/printer.py | 16 ++++++++++++---- posthog/hogql/test/test_printer.py | 8 ++++++-- 2 files changed, 18 insertions(+), 6 deletions(-) diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py index 2731897748c99..98d3bdc4bc8a5 100644 --- 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 index ad74a0ca87648..1f3b2db78e7ed 100644 --- 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", }, )