From 5548c47e798d687ea0332e50dbded1cd6186e646 Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Thu, 5 Feb 2026 20:49:20 +0000 Subject: [PATCH 1/6] Move extend clause after summarize for calculated measures This ensures that extend operations (which may reference summarized columns) are executed after the summarize clause in the generated KQL query. --- sqlalchemy_kusto/dialect_kql.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 5be6be3..0cc0bd0 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -142,9 +142,15 @@ def visit_select( ) compiled_query_lines.append(f"| where {converted_where_clause}") + # Add summarize first if it exists + if "summarize" in projections_parts_dict: + compiled_query_lines.append(projections_parts_dict.pop("summarize")) + + # Then add extend after summarize if "extend" in projections_parts_dict: compiled_query_lines.append(projections_parts_dict.pop("extend")) + # Add remaining parts (project, sort) for statement_part in projections_parts_dict.values(): if statement_part: compiled_query_lines.append(statement_part) From c72ee6166c088aa0635d13fc5a0b3b74311e4fe5 Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Fri, 6 Feb 2026 00:35:30 +0000 Subject: [PATCH 2/6] Add test for calculated measure with ad hoc measure and constant and update test expectations for summarize-before-extend ordering --- sqlalchemy_kusto/dialect_kql.py | 2 +- tests/unit/test_dialect_kql.py | 46 ++++++++++++++++++++++++--------- 2 files changed, 35 insertions(+), 13 deletions(-) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 0cc0bd0..dac4054 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -107,7 +107,7 @@ def visit_select( from_object = select_stmt.get_final_froms()[0] if hasattr(from_object, "element"): query = self._get_most_inner_element(from_object.element) - (main, lets) = self._extract_let_statements(query.text) + main, lets = self._extract_let_statements(query.text) compiled_query_lines.extend(lets) compiled_query_lines.append( f"let {from_object.name} = ({self._convert_schema_in_statement(main)});" diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py index 04b8fc0..8791486 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -175,9 +175,10 @@ def test_group_by_text(): ).replace("\n", "") # raw query text from query query_expected = ( - '["ActiveUsersLastMonth"]| extend ["ActiveUserMetric"] = ["ActiveUsers"], ' - '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' + '["ActiveUsersLastMonth"]' '| summarize by ["EventInfo_Time"] / time(1d)' + '| extend ["ActiveUserMetric"] = ["ActiveUsers"], ' + '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["ActiveUserMetric"]' '| order by ["ActiveUserMetric"] desc' ) @@ -224,20 +225,19 @@ def test_group_by_text_vaccine_dataset(): query.compile(engine, compile_kwargs={"literal_binds": True}) ).replace("\n", "") query_expected = ( - 'database("superset").["CovidVaccineData"]| ' - 'extend ["country_name"] = ["country_name"]| ' - 'summarize by ["country_name"]| ' - 'project ["country_name"]| order by ["country_name"] asc' + 'database("superset").["CovidVaccineData"]' + '| summarize by ["country_name"]' + '| extend ["country_name"] = ["country_name"]' + '| project ["country_name"]' + '| order by ["country_name"] asc' ) assert query_compiled == query_expected def test_is_kql_function(): - assert KustoKqlCompiler._is_kql_function( - """case(Size <= 3, "Small", + assert KustoKqlCompiler._is_kql_function("""case(Size <= 3, "Small", Size <= 10, "Medium", - "Large")""" - ) + "Large")""") assert KustoKqlCompiler._is_kql_function("""bin(time(16d), 7d)""") assert KustoKqlCompiler._is_kql_function( """iff((EventType in ("Heavy Rain", "Flash Flood", "Flood")), "Rain event", "Not rain event")""" @@ -328,8 +328,8 @@ def test_distinct_count_by_text(): # raw query text from query query_expected = ( '["ActiveUsersLastMonth"]' - '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' + '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["DistinctUsers"]' '| order by ["ActiveUserMetric"] desc' ) @@ -354,8 +354,8 @@ def test_distinct_count_alt_by_text(): # raw query text from query query_expected = ( '["ActiveUsersLastMonth"]' - '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' + '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["DistinctUsers"]' '| order by ["ActiveUserMetric"] desc' ) @@ -549,6 +549,28 @@ def test_match_aggregates(column_name: str, expected_aggregate: str): assert kql_agg is None +def test_calculated_measure_with_adhoc_measure_and_constant(): + """Test calculated measure with an ad hoc measure and a constant. + + Measure 1 = count(*), Measure 2 = "Measure 1" * 2 + Measure 2 should compile to ["Measure 1"] * 2 + The extend clause must come after summarize for this to work. + """ + measure_1 = literal_column("count(*)").label("Measure 1") + measure_2 = literal_column('"Measure 1" * 2').label("Measure 2") + query = select([measure_1, measure_2]).select_from(text("SalesData")) + query_compiled = str( + query.compile(engine, compile_kwargs={"literal_binds": True}) + ).replace("\n", "") + query_expected = ( + '["SalesData"]' + '| summarize ["Measure 1"] = count() ' + '| extend ["Measure 2"] = ["Measure 1"] * 2' + '| project ["Measure 1"], ["Measure 2"]' + ) + assert query_compiled == query_expected + + @pytest.mark.parametrize( ("query_table_name", "expected_table_name"), [ From 63599f0fa571648146c84f308c3806ad98d72cb0 Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Fri, 13 Mar 2026 22:54:04 +0000 Subject: [PATCH 3/6] Split extend into pre/post-summarize for calculated columns vs measures Calculated columns (non-aggregate expressions like strcat, iff) now get their extend clause placed BEFORE summarize so they exist for group-by. Calculated measures (expressions referencing aggregate aliases) remain AFTER summarize where they can reference aggregated values. Tracks all aggregate aliases in a separate set to avoid overwrites when multiple measures use the same aggregate expression (e.g. two count(*)). Adds test_pre_aggregated_calculations from PR review comment to verify calculated columns used as group-by dimensions work correctly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sqlalchemy_kusto/dialect_kql.py | 44 ++++++++++++++++++++--------- tests/unit/test_dialect_kql.py | 50 +++++++++++++++++++++++++++------ 2 files changed, 71 insertions(+), 23 deletions(-) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index dac4054..9245cba 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -142,13 +142,10 @@ def visit_select( ) compiled_query_lines.append(f"| where {converted_where_clause}") - # Add summarize first if it exists - if "summarize" in projections_parts_dict: - compiled_query_lines.append(projections_parts_dict.pop("summarize")) - - # Then add extend after summarize - if "extend" in projections_parts_dict: - compiled_query_lines.append(projections_parts_dict.pop("extend")) + # Add clauses in correct order: pre-extend, summarize, post-extend + for key in ("pre_extend", "summarize", "post_extend"): + if key in projections_parts_dict: + compiled_query_lines.append(projections_parts_dict.pop(key)) # Add remaining parts (project, sort) for statement_part in projections_parts_dict.values(): @@ -204,7 +201,8 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s group_by_cols = select._group_by_clauses order_by_cols = select._order_by_clauses summarize_statement = "" - extend_statement = "" + pre_extend_statement = "" + post_extend_statement = "" project_statement = "" has_aggregates = False # The following is the logic @@ -219,8 +217,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s # N---> Add to projection if columns is not None: summarize_columns = set() - extend_columns = set() + pre_extend_columns: list[str] = [] + post_extend_columns: list[str] = [] projection_columns = [] + all_agg_aliases: set[str] = set() for column in [c for c in columns if c.name != "*"]: column_name, column_alias = self._extract_column_name_and_alias(column) column_alias = self._escape_and_quote_columns(column_alias, True) @@ -232,13 +232,26 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s summarize_columns.add( self._build_column_projection(kql_agg, column_alias) ) + if column_alias: + all_agg_aliases.add(column_alias) # No group by clause # Do the columns have aliases ? # Add additional and to handle case where : SELECT column_name as column_name elif column_alias and column_alias != column_name: - extend_columns.add( - self._build_column_projection(column_name, column_alias, True) + extend_entry = self._build_column_projection( + column_name, column_alias, True + ) + escaped_expr = self._escape_and_quote_columns(column_name) + + # Calculated measures (aggregate-dependent) go after summarize; + # calculated columns (no aggregate dependency) go before summarize + target = ( + post_extend_columns + if any(alias in escaped_expr for alias in all_agg_aliases) + else pre_extend_columns ) + target.append(extend_entry) + if column_alias: projection_columns.append( self._escape_and_quote_columns(column_alias, True) @@ -257,8 +270,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s summarize_statement = ( f"{summarize_statement} by {', '.join(by_columns)}" ) - if extend_columns: - extend_statement = f"| extend {', '.join(sorted(extend_columns))}" + if pre_extend_columns: + pre_extend_statement = f"| extend {', '.join(pre_extend_columns)}" + if post_extend_columns: + post_extend_statement = f"| extend {', '.join(post_extend_columns)}" project_statement = ( f"| project {', '.join(projection_columns)}" if projection_columns @@ -269,7 +284,8 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s f"| order by {', '.join(unwrapped_order_by)}" if unwrapped_order_by else "" ) return { - "extend": extend_statement, + "pre_extend": pre_extend_statement, + "post_extend": post_extend_statement, "summarize": summarize_statement, "project": project_statement, "sort": sort_statement, diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py index 8791486..efdd149 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -173,12 +173,12 @@ def test_group_by_text(): query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) ).replace("\n", "") - # raw query text from query + # raw query text from query - order matches column appearance query_expected = ( '["ActiveUsersLastMonth"]' + '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d), ' + '["ActiveUserMetric"] = ["ActiveUsers"]' '| summarize by ["EventInfo_Time"] / time(1d)' - '| extend ["ActiveUserMetric"] = ["ActiveUsers"], ' - '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["ActiveUserMetric"]' '| order by ["ActiveUserMetric"] desc' ) @@ -202,12 +202,13 @@ def test_function_text(f: str, expected: str): query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) ).replace("\n", "") + # Order matches column appearance in select query_expected = ( '["ActiveUsersLastMonth"]' - '| extend ["ActiveUserMetric"] = ["ActiveUsers"], ' - '["EventInfo_Time"] = ' + '| extend ["EventInfo_Time"] = ' + expected - + '| project ["EventInfo_Time"], ["ActiveUserMetric"]' + + ', ["ActiveUserMetric"] = ["ActiveUsers"]' + '| project ["EventInfo_Time"], ["ActiveUserMetric"]' ) assert query_compiled == query_expected @@ -226,8 +227,8 @@ def test_group_by_text_vaccine_dataset(): ).replace("\n", "") query_expected = ( 'database("superset").["CovidVaccineData"]' - '| summarize by ["country_name"]' '| extend ["country_name"] = ["country_name"]' + '| summarize by ["country_name"]' '| project ["country_name"]' '| order by ["country_name"] asc' ) @@ -328,8 +329,8 @@ def test_distinct_count_by_text(): # raw query text from query query_expected = ( '["ActiveUsersLastMonth"]' - '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' + '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["DistinctUsers"]' '| order by ["ActiveUserMetric"] desc' ) @@ -354,8 +355,8 @@ def test_distinct_count_alt_by_text(): # raw query text from query query_expected = ( '["ActiveUsersLastMonth"]' - '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' + '| summarize ["DistinctUsers"] = dcount(["ActiveUsers"]) by ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["DistinctUsers"]' '| order by ["ActiveUserMetric"] desc' ) @@ -571,6 +572,37 @@ def test_calculated_measure_with_adhoc_measure_and_constant(): assert query_compiled == query_expected +def test_pre_aggregated_calculations(): + """Test from PR #48 review: calculated column before aggregation. + + A calculated column (strcat) is created and used as a group-by dimension. + The extend must appear before summarize so the column exists for grouping. + """ + app_col = literal_column("App") + ns_col = literal_column("Namespace") + id_col = literal_column("_id") + + app_namespace = sa.func.strcat(app_col, ns_col).label("App_Namespace") + + query = ( + select(app_namespace, sa.func.count(id_col).label("TotalLogs")) + .select_from(text("Logs")) + .group_by(app_namespace) + ) + + query_compiled = str( + query.compile(engine, compile_kwargs={"literal_binds": True}) + ).replace("\n", "") + + query_expected = ( + '["Logs"]' + '| extend ["App_Namespace"] = strcat(App, Namespace)' + '| summarize ["TotalLogs"] = count(["_id"]) by ["App_Namespace"]' + '| project ["App_Namespace"], ["TotalLogs"]' + ) + assert query_compiled == query_expected + + @pytest.mark.parametrize( ("query_table_name", "expected_table_name"), [ From fecc047a29b54acf9015e6a6114b44e5b8099d5f Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Tue, 17 Mar 2026 00:25:02 +0000 Subject: [PATCH 4/6] Use robust regex-based dependency check for pre/post extend split Replace fragile substring matching on _escape_and_quote_columns output with a dedicated _expression_references_aliases method that uses regex to check for aggregate alias references in all forms (double-quoted, KQL-escaped, and bare identifiers). This correctly handles: - Aliases inside function calls (e.g., round("Measure 1", 2)) - Aliases on the RHS of operators (e.g., 2 * "Measure 1") Also fix missing parametrize argnames in test_schema_from_query and add tests for the above edge cases. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sqlalchemy_kusto/dialect_kql.py | 36 ++++++++++++++++++++++++--- tests/unit/test_dialect_kql.py | 43 ++++++++++++++++++++++++++++++++- 2 files changed, 74 insertions(+), 5 deletions(-) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 9245cba..9bf6327 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -220,7 +220,7 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s pre_extend_columns: list[str] = [] post_extend_columns: list[str] = [] projection_columns = [] - all_agg_aliases: set[str] = set() + all_agg_aliases_raw: set[str] = set() for column in [c for c in columns if c.name != "*"]: column_name, column_alias = self._extract_column_name_and_alias(column) column_alias = self._escape_and_quote_columns(column_alias, True) @@ -233,7 +233,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s self._build_column_projection(kql_agg, column_alias) ) if column_alias: - all_agg_aliases.add(column_alias) + raw = column_alias + if raw.startswith('["') and raw.endswith('"]'): + raw = raw[2:-2] + all_agg_aliases_raw.add(raw) # No group by clause # Do the columns have aliases ? # Add additional and to handle case where : SELECT column_name as column_name @@ -241,13 +244,14 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s extend_entry = self._build_column_projection( column_name, column_alias, True ) - escaped_expr = self._escape_and_quote_columns(column_name) # Calculated measures (aggregate-dependent) go after summarize; # calculated columns (no aggregate dependency) go before summarize target = ( post_extend_columns - if any(alias in escaped_expr for alias in all_agg_aliases) + if self._expression_references_aliases( + column_name, all_agg_aliases_raw + ) else pre_extend_columns ) target.append(extend_entry) @@ -369,6 +373,30 @@ def replacer(match): return modified_expression + @staticmethod + def _expression_references_aliases( + expression: str, raw_aliases: set[str] + ) -> bool: + """Check if an expression references any aggregate alias. + + Checks for both double-quoted ("alias") and KQL-escaped (["alias"]) + forms across the entire expression, including inside function calls + and on either side of operators. + """ + for alias in raw_aliases: + escaped_alias = re.escape(alias) + # Match ["alias"] or "alias" anywhere in the expression + if re.search( + rf'\["{escaped_alias}"\]|"{escaped_alias}"', expression + ): + return True + # For simple identifiers, also match unquoted bare references + if re.fullmatch(r"[A-Za-z_]\w*", alias) and re.search( + rf"\b{escaped_alias}\b", expression + ): + return True + return False + @staticmethod def _escape_and_quote_columns(name: str | None, is_alias=False) -> str: if name is None: diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py index efdd149..27e7f42 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -603,10 +603,51 @@ def test_pre_aggregated_calculations(): assert query_compiled == query_expected +def test_calculated_measure_alias_in_function(): + """Test that an aggregate alias inside a function is detected as post-extend. + + round("Measure 1", 2) references the aggregate alias inside a function call, + which _escape_and_quote_columns returns unchanged (it's a KQL function). + The robust check must still detect the dependency. + """ + measure_1 = literal_column("count(*)").label("Measure 1") + measure_2 = literal_column('round("Measure 1", 2)').label("Rounded") + query = select([measure_1, measure_2]).select_from(text("SalesData")) + query_compiled = str( + query.compile(engine, compile_kwargs={"literal_binds": True}) + ).replace("\n", "") + query_expected = ( + '["SalesData"]' + '| summarize ["Measure 1"] = count() ' + '| extend ["Rounded"] = round(["Measure 1"], 2)' + '| project ["Measure 1"], ["Rounded"]' + ) + assert query_compiled == query_expected + + +def test_calculated_measure_alias_on_rhs_of_operator(): + """Test that an aggregate alias on the RHS of an operator is detected. + + 2 * "Measure 1" has the alias on the right side of the operator. + """ + measure_1 = literal_column("count(*)").label("Measure 1") + measure_2 = literal_column('2 * "Measure 1"').label("Doubled") + query = select([measure_1, measure_2]).select_from(text("SalesData")) + query_compiled = str( + query.compile(engine, compile_kwargs={"literal_binds": True}) + ).replace("\n", "") + query_expected = ( + '["SalesData"]' + '| summarize ["Measure 1"] = count() ' + '| extend ["Doubled"] = ["2"] * "Measure 1"' + '| project ["Measure 1"], ["Doubled"]' + ) + assert query_compiled == query_expected + + @pytest.mark.parametrize( ("query_table_name", "expected_table_name"), [ - ("schema.table", 'database("schema").["table"]'), ('schema."table.name"', 'database("schema").["table.name"]'), ('"schema.name".table', 'database("schema.name").["table"]'), ('"schema.name"."table.name"', 'database("schema.name").["table.name"]'), From f032c854daf58b5c77153c2a88bcc3cf251cda8d Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Thu, 19 Mar 2026 21:48:18 +0000 Subject: [PATCH 5/6] fix: apply black formatting to dialect_kql.py Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sqlalchemy_kusto/dialect_kql.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 9bf6327..b68a2b3 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -374,9 +374,7 @@ def replacer(match): return modified_expression @staticmethod - def _expression_references_aliases( - expression: str, raw_aliases: set[str] - ) -> bool: + def _expression_references_aliases(expression: str, raw_aliases: set[str]) -> bool: """Check if an expression references any aggregate alias. Checks for both double-quoted ("alias") and KQL-escaped (["alias"]) @@ -386,9 +384,7 @@ def _expression_references_aliases( for alias in raw_aliases: escaped_alias = re.escape(alias) # Match ["alias"] or "alias" anywhere in the expression - if re.search( - rf'\["{escaped_alias}"\]|"{escaped_alias}"', expression - ): + if re.search(rf'\["{escaped_alias}"\]|"{escaped_alias}"', expression): return True # For simple identifiers, also match unquoted bare references if re.fullmatch(r"[A-Za-z_]\w*", alias) and re.search( From ef0cf70653412325599bfca6902adb4062f7954c Mon Sep 17 00:00:00 2001 From: Alison Gim Date: Thu, 19 Mar 2026 22:23:43 +0000 Subject: [PATCH 6/6] fix: address PR review comments - Two-pass alias collection: extract _collect_aggregate_aliases() to pre-scan all columns before classifying extend entries, fixing the forward-reference bug where a calculated measure listed before its aggregate dependency was misclassified as pre_extend. - Fix numeric literal escaping: when the LHS of an operator is a number (e.g. '2 * "Measure 1"'), keep it as a literal instead of wrapping it in ["2"], and properly escape the RHS through the standard path. - Add test_calculated_measure_forward_reference to verify the two-pass fix works when Measure 2 (referencing Measure 1) appears first in the select list. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- sqlalchemy_kusto/dialect_kql.py | 41 ++++++++++++++++++++++++++------- tests/unit/test_dialect_kql.py | 24 ++++++++++++++++++- 2 files changed, 56 insertions(+), 9 deletions(-) diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index b68a2b3..ce58fa7 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -220,8 +220,14 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s pre_extend_columns: list[str] = [] post_extend_columns: list[str] = [] projection_columns = [] - all_agg_aliases_raw: set[str] = set() - for column in [c for c in columns if c.name != "*"]: + + # Two-pass approach: first collect all aggregate aliases so that + # forward references (a calculated measure listed before the + # aggregate it depends on) are classified correctly. + columns_list = [c for c in columns if c.name != "*"] + all_agg_aliases_raw = self._collect_aggregate_aliases(columns_list) + + for column in columns_list: column_name, column_alias = self._extract_column_name_and_alias(column) column_alias = self._escape_and_quote_columns(column_alias, True) # Do we have a group by clause ? @@ -232,11 +238,6 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s summarize_columns.add( self._build_column_projection(kql_agg, column_alias) ) - if column_alias: - raw = column_alias - if raw.startswith('["') and raw.endswith('"]'): - raw = raw[2:-2] - all_agg_aliases_raw.add(raw) # No group by clause # Do the columns have aliases ? # Add additional and to handle case where : SELECT column_name as column_name @@ -373,6 +374,25 @@ def replacer(match): return modified_expression + def _collect_aggregate_aliases(self, columns_list) -> set[str]: + """Pre-scan columns to collect all aggregate aliases. + + This ensures calculated measures that forward-reference an aggregate + defined later in the select list are still classified correctly. + """ + aliases: set[str] = set() + for column in columns_list: + column_name, _ = self._extract_column_name_and_alias(column) + if self._extract_maybe_agg_column_parts(column_name): + _, col_alias = self._extract_column_name_and_alias(column) + col_alias = self._escape_and_quote_columns(col_alias, True) + if col_alias: + raw = col_alias + if raw.startswith('["') and raw.endswith('"]'): + raw = raw[2:-2] + aliases.add(raw) + return aliases + @staticmethod def _expression_references_aliases(expression: str, raw_aliases: set[str]) -> bool: """Check if an expression references any aggregate alias. @@ -417,10 +437,15 @@ def _escape_and_quote_columns(name: str | None, is_alias=False) -> str: parts = name.split(operator, 1) # Remove quotes if they exist at the edges col_part = parts[0].strip() + rhs = parts[1].strip() + # If LHS is a numeric literal, keep it as-is and escape the RHS + if KustoKqlCompiler._is_number_literal(col_part): + escaped_rhs = KustoKqlCompiler._escape_and_quote_columns(rhs) + return f"{col_part} {operator} {escaped_rhs}" if col_part.startswith('"') and col_part.endswith('"'): col_part = col_part[1:-1].strip() col_part = col_part.replace('"', '\\"') - return f'["{col_part}"] {operator} {parts[1].strip()}' # Wrap the column part + return f'["{col_part}"] {operator} {rhs}' # Wrap the column part # No operators found, just wrap the entire name name = name.replace('"', '\\"') return f'["{name}"]' diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py index 27e7f42..0b704bc 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -639,12 +639,34 @@ def test_calculated_measure_alias_on_rhs_of_operator(): query_expected = ( '["SalesData"]' '| summarize ["Measure 1"] = count() ' - '| extend ["Doubled"] = ["2"] * "Measure 1"' + '| extend ["Doubled"] = 2 * ["Measure 1"]' '| project ["Measure 1"], ["Doubled"]' ) assert query_compiled == query_expected +def test_calculated_measure_forward_reference(): + """Test that a calculated measure listed before its aggregate dependency works. + + When Measure 2 (which references Measure 1) appears before Measure 1 in the + select list, the two-pass alias collection ensures it is still classified as + post_extend. + """ + measure_2 = literal_column('"Measure 1" * 2').label("Measure 2") + measure_1 = literal_column("count(*)").label("Measure 1") + query = select([measure_2, measure_1]).select_from(text("SalesData")) + query_compiled = str( + query.compile(engine, compile_kwargs={"literal_binds": True}) + ).replace("\n", "") + query_expected = ( + '["SalesData"]' + '| summarize ["Measure 1"] = count() ' + '| extend ["Measure 2"] = ["Measure 1"] * 2' + '| project ["Measure 2"], ["Measure 1"]' + ) + assert query_compiled == query_expected + + @pytest.mark.parametrize( ("query_table_name", "expected_table_name"), [