Fix support for calculated measures in KQL compiler

[AlisonNeedsCopilot]

PR: Fix KQL Dialect for Calculated Measures Support

Overview

The KQL (Kusto Query Language) SQLAlchemy dialect in sqlalchemy-kusto does not properly handle calculated measures. When users create calculated measures that combine multiple aggregates with arithmetic operators (e.g., "Measure 1" + "Measure 2"), the generated KQL is malformed and fails to execute.

Issues fixed:

1. Parentheses not preserved - Expressions like (("Measure 1")) were losing their parentheses, breaking operator precedence
2. Quoted identifiers not converted - SQL-style "column" wasn't being converted to KQL-style ["column"]
3. Uppercase function names - KQL requires lowercase functions (count() not COUNT())
4. Incorrect clause ordering - extend appeared before summarize, causing reference errors since extend columns reference values created in summarize

Additional improvements:

• Multi-arg KQL function handling - Improved column argument extraction for functions like percentile(col, 99) and dcountif(col, predicate)

---

Technical Details

New Functions Added

Function	Type	Description
_find_top_level_operator(text, operator)	Module-level helper	Finds the position of an operator at depth 0 (not inside quotes, brackets, or parentheses). Returns -1 if not found. Essential for correctly parsing nested arithmetic expressions like (A + B) * C.
_count_outer_parens(text)	Static method	Counts and strips matched outer parentheses from text. Returns (count, stripped_text). Used to preserve parentheses for operator precedence in calculated measures.
_has_operators_outside_quotes(expr)	Static method	Checks if an expression has arithmetic operators (+, -, *, /) outside of quoted strings and brackets. Used to detect calculated measures vs simple aggregates.

Functions Updated

Function	Changes
visit_select()	Added explicit ordering to ensure summarize clause is added before extend clause. This is required because extend columns reference values created in summarize.
_get_projection_or_summarize()	Added calculated measures detection using _has_operators_outside_quotes(). Added inline aggregate extraction to summarize for expressions containing aggregates. Added column reference wrapping in parentheses for arithmetic precedence.
_extract_maybe_agg_column_parts()	Reordered logic to check AGGREGATE_PATTERN first (handles SQL-style aggregates including count(distinct X)). Added support for multi-arg KQL aggregates like percentile(col, 99) and dcountif(col, predicate) - now properly escapes the first (column) argument while leaving predicates/numeric args untouched.
_escape_and_quote_columns()	Added recursive handling of arithmetic expressions using _find_top_level_operator(). Preserves parentheses count using _count_outer_parens(). Properly processes nested expressions like (A + B) * C by recursively escaping each operand.
_convert_quoted_columns()	Added regex to convert standalone quoted identifiers ("col") to bracket notation (["col"]), not just those inside function calls.
_sql_to_kql_aggregate()	Fixed bug: Changed "*" in sql_agg to "*" in str(column_name) for proper star detection. Added null check for extra_params to prevent None from being concatenated.

---

Example

Input SQL expression:

("UserInfo_Ring Count" + "UserInfo_Region Count") * 2

Before (broken KQL):

| extend ["Calculated Measure"] = "UserInfo_Ring Count" + "UserInfo_Region Count" * 2
| summarize ["UserInfo_Ring Count"] = COUNT(["UserInfo_Ring"]), ...

After (correct KQL):

| summarize ["UserInfo_Ring Count"] = count(["UserInfo_Ring"]), ["UserInfo_Region Count"] = count(["UserInfo_Region"])
| extend ["Calculated Measure"] = (["UserInfo_Ring Count"] + ["UserInfo_Region Count"]) * 2

UI Change

---

Testing

• Added comprehensive test suite with 112 tests covering all KQL dialect functionality
• All existing tests continue to pass
• New tests added for:
  • Calculated measures with arithmetic operators
  • Parentheses preservation (single, double, nested)
  • Standalone quoted identifier conversion
  • Multi-arg KQL functions (percentile, dcountif, countif, etc.)
  • Lowercase function names

[Copilot]

Pull request overview

This PR fixes critical issues in the KQL (Kusto Query Language) SQLAlchemy dialect to properly support calculated measures - arithmetic expressions that combine multiple aggregates. The main issues addressed are incorrect clause ordering (extend before summarize), missing parentheses preservation, quoted identifier conversion, and uppercase function names.

Changes:

• Added helper functions for parsing arithmetic expressions and detecting operators outside quotes/brackets
• Fixed clause ordering to ensure summarize appears before extend in generated KQL
• Enhanced aggregate detection and column escaping to handle calculated measures with arithmetic operators
• Improved multi-argument KQL function support (percentile, dcountif, etc.)
• Modernized test syntax from SQLAlchemy 1.x to 2.x patterns
• Added comprehensive test suite (112 tests) covering calculated measures functionality

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 6 comments.

File	Description
sqlalchemy_kusto/dialect_kql.py	Core implementation of calculated measures support including new helper functions _find_top_level_operator, _count_outer_parens, _has_operators_outside_quotes, and updates to query compilation logic
tests/unit/test_dialect_kql.py	Test modernization (SQLAlchemy 2.x syntax) and new TestCalculatedMeasures class with comprehensive tests for arithmetic expressions, parentheses handling, and quoted identifiers

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.