fix: timings wiring up (PostHog#29010)

sloppyjuicy · Feb 20, 2025 · eeea5a6 · eeea5a6
1 parent 6bde178
commit eeea5a6
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 26 deletions.
diff --git a/posthog/hogql/autocomplete.py b/posthog/hogql/autocomplete.py
@@ -372,9 +372,9 @@ def get_hogql_autocomplete(
     if database_arg is not None:
         database = database_arg
     else:
-        database = create_hogql_database(team_id=team.pk, team_arg=team)
+        database = create_hogql_database(team_id=team.pk, team_arg=team, timings=timings)
 
-    context = HogQLContext(team_id=team.pk, team=team, database=database)
+    context = HogQLContext(team_id=team.pk, team=team, database=database, timings=timings)
     if query.sourceQuery:
         if query.sourceQuery.kind == "HogQLQuery" and (
             query.sourceQuery.query is None or query.sourceQuery.query == ""

diff --git a/posthog/hogql/hogql.py b/posthog/hogql/hogql.py
@@ -32,7 +32,7 @@ def translate_hogql(
         if context.database is None:
             if context.team_id is None:
                 raise ValueError("Cannot translate HogQL for a filter with no team specified")
-            context.database = create_hogql_database(context.team_id, context.modifiers)
+            context.database = create_hogql_database(context.team_id, context.modifiers, timings=context.timings)
         node = parse_expr(query, placeholders=placeholders)
         select_query = ast.SelectQuery(select=[node], select_from=ast.JoinExpr(table=ast.Field(chain=["events"])))
 

diff --git a/posthog/hogql/printer.py b/posthog/hogql/printer.py
@@ -117,7 +117,9 @@ def prepare_ast_for_printing(
 ) -> _T_AST | None:
     if context.database is None:
         with context.timings.measure("create_hogql_database"):
-            context.database = create_hogql_database(context.team_id, context.modifiers, context.team)
+            context.database = create_hogql_database(
+                context.team_id, context.modifiers, context.team, timings=context.timings
+            )
 
     context.modifiers = set_default_in_cohort_via(context.modifiers)
 

diff --git a/posthog/hogql/query.py b/posthog/hogql/query.py
@@ -123,26 +123,25 @@ def _apply_limit(self):
                     )
 
     def _generate_hogql(self):
-        with self.timings.measure("hogql"):
-            with self.timings.measure("prepare_ast"):
-                hogql_query_context = dataclasses.replace(
-                    self.context,
-                    team_id=self.team.pk,
-                    team=self.team,
-                    enable_select_queries=True,
-                    timings=self.timings,
-                    modifiers=self.query_modifiers,
-                )
+        hogql_query_context = dataclasses.replace(
+            self.context,
+            team_id=self.team.pk,
+            team=self.team,
+            enable_select_queries=True,
+            timings=self.timings,
+            modifiers=self.query_modifiers,
+        )
 
-            with self.timings.measure("clone"):
-                cloned_query = clone_expr(self.select_query, True)
+        with self.timings.measure("clone"):
+            cloned_query = clone_expr(self.select_query, True)
 
+        with self.timings.measure("prepare_ast_for_printing"):
             select_query_hogql = cast(
                 ast.SelectQuery,
                 prepare_ast_for_printing(node=cloned_query, context=hogql_query_context, dialect="hogql"),
             )
 
-        with self.timings.measure("print_ast"):
+        with self.timings.measure("print_prepared_ast"):
             self.hogql = print_prepared_ast(
                 select_query_hogql,
                 hogql_query_context,
@@ -186,13 +185,14 @@ def _generate_clickhouse_sql(self):
                 timings=self.timings,
                 modifiers=self.query_modifiers,
             )
-            self.clickhouse_sql = print_ast(
-                self.select_query,
-                context=self.clickhouse_context,
-                dialect="clickhouse",
-                settings=settings,
-                pretty=self.pretty if self.pretty is not None else True,
-            )
+            with self.timings.measure("print_ast"):
+                self.clickhouse_sql = print_ast(
+                    self.select_query,
+                    context=self.clickhouse_context,
+                    dialect="clickhouse",
+                    settings=settings,
+                    pretty=self.pretty if self.pretty is not None else True,
+                )
         except Exception as e:
             if self.debug:
                 self.clickhouse_sql = ""
@@ -259,8 +259,10 @@ def generate_clickhouse_sql(self) -> tuple[str, HogQLContext]:
         self._process_variables()
         self._process_placeholders()
         self._apply_limit()
-        self._generate_hogql()
-        self._generate_clickhouse_sql()
+        with self.timings.measure("_generate_hogql"):
+            self._generate_hogql()
+        with self.timings.measure("_generate_clickhouse_sql"):
+            self._generate_clickhouse_sql()
         return self.clickhouse_sql, self.clickhouse_context
 
     def execute(self) -> HogQLQueryResponse: