Move extend clause after summarize for calculated measures

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}"]'

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"]'),