diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 5be6be3..ce58fa7 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)});" @@ -142,9 +142,12 @@ def visit_select( ) compiled_query_lines.append(f"| where {converted_where_clause}") - 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(): if statement_part: compiled_query_lines.append(statement_part) @@ -198,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 @@ -213,9 +217,17 @@ 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 = [] - 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 ? @@ -230,9 +242,21 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s # 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 + ) + + # Calculated measures (aggregate-dependent) go after summarize; + # calculated columns (no aggregate dependency) go before summarize + target = ( + post_extend_columns + if self._expression_references_aliases( + column_name, all_agg_aliases_raw + ) + else pre_extend_columns ) + target.append(extend_entry) + if column_alias: projection_columns.append( self._escape_and_quote_columns(column_alias, True) @@ -251,8 +275,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 @@ -263,7 +289,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, @@ -347,6 +374,45 @@ 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. + + 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: @@ -371,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 04b8fc0..0b704bc 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -173,10 +173,11 @@ 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 ["ActiveUserMetric"] = ["ActiveUsers"], ' - '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' + '["ActiveUsersLastMonth"]' + '| extend ["EventInfo_Time"] = ["EventInfo_Time"] / time(1d), ' + '["ActiveUserMetric"] = ["ActiveUsers"]' '| summarize by ["EventInfo_Time"] / time(1d)' '| project ["EventInfo_Time"], ["ActiveUserMetric"]' '| order by ["ActiveUserMetric"] desc' @@ -201,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 @@ -224,20 +226,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"]' + '| extend ["country_name"] = ["country_name"]' + '| summarize by ["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")""" @@ -549,10 +550,126 @@ 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 + + +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 + + +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 + + +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"), [ - ("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"]'),