diff --git a/posthog/hogql/autocomplete.py b/posthog/hogql/autocomplete.py index 304488d764426..aa8babe13741e 100644 --- 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 index 00553d3d65a94..f40c22b82ed95 100644 --- 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 index a80c2f36328e0..a090b5585db45 100644 --- 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 index 011b844da2446..a502f537215cd 100644 --- 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: