diff --git a/sqlalchemy_kusto/dialect_kql.py b/sqlalchemy_kusto/dialect_kql.py index 5be6be3..9bb678f 100644 --- a/sqlalchemy_kusto/dialect_kql.py +++ b/sqlalchemy_kusto/dialect_kql.py @@ -67,6 +67,46 @@ } AGGREGATE_PATTERN = r"(\w+)\s*\(\s*(DISTINCT|distinct\s*)?\(?\s*(\*|\[?\"?\'?\w+\"?\]?)\s*(,.+)*\)?\s*\)" +# Pre-compiled regex for aggregate function matching (performance optimization). +# Compiled once at module load to avoid recompiling on every call, which significantly +# improves performance for query-heavy workloads. +KQL_AGG_PATTERN = re.compile(r'\b(' + '|'.join(kql_aggregates) + r')\s*\(', re.IGNORECASE) + + +class _ParseState: + """Tracks parsing state while scanning through text.""" + __slots__ = ('in_double_quote', 'in_single_quote', 'in_bracket', 'paren_depth') + + def __init__(self): + self.in_double_quote = False + self.in_single_quote = False + self.in_bracket = False + self.paren_depth = 0 + + def update(self, ch: str, prev_ch: str | None) -> None: + """Update state based on current and previous character.""" + # Handle quotes (only if not escaped and not in conflicting context) + if ch == '"' and prev_ch != "\\" and not self.in_single_quote and not self.in_bracket: + self.in_double_quote = not self.in_double_quote + elif ch == "'" and prev_ch != "\\" and not self.in_double_quote and not self.in_bracket: + self.in_single_quote = not self.in_single_quote + # Handle brackets and parens (only if not in quotes) + elif not self.in_double_quote and not self.in_single_quote: + if ch == "[": + self.in_bracket = True + elif ch == "]": + self.in_bracket = False + elif not self.in_bracket: + if ch == "(": + self.paren_depth += 1 + elif ch == ")": + self.paren_depth -= 1 + + @property + def in_quotes_or_brackets(self) -> bool: + """Check if currently inside quotes or brackets.""" + return self.in_double_quote or self.in_single_quote or self.in_bracket + class UniversalSet: def __contains__(self, item): @@ -89,6 +129,89 @@ class KustoKqlCompiler(compiler.SQLCompiler): visit_sequence = None sort_with_clause_parts = 2 + @staticmethod + def _find_top_level_operator(text: str, operator: str) -> int: + """Find position of operator at depth 0 (not inside quotes, brackets, or parens). + + Args: + text: The string to search in + operator: The single-character operator to find (e.g., '+', '-', '*', '/') + + Returns: + The position of the operator at depth 0, or -1 if not found. + Returns -1 when the operator only appears inside quotes, brackets, or nested parens. + """ + state = _ParseState() + for i, ch in enumerate(text): + if ch == operator and state.paren_depth == 0 and not state.in_quotes_or_brackets: + return i + state.update(ch, text[i - 1] if i > 0 else None) + return -1 + + @staticmethod + def _find_matching_paren(text: str, start_pos: int) -> int: + """Find the matching closing parenthesis for an opening paren at start_pos. + + Args: + text: The string containing parentheses + start_pos: The position of the opening parenthesis + + Returns: + The position of the matching closing parenthesis, or -1 if: + - start_pos is out of bounds + - The character at start_pos is not '(' + - No matching closing parenthesis is found + + Note: + This function properly handles: + - Nested parentheses + - Parentheses inside quoted strings (ignored) + - Escaped quotes + - Parentheses inside brackets (ignored) + """ + if start_pos >= len(text) or text[start_pos] != '(': + return -1 + + state = _ParseState() + state.paren_depth = 1 # Start with depth 1 since we're at opening paren + + for i in range(start_pos + 1, len(text)): + ch = text[i] + state.update(ch, text[i - 1] if i > 0 else None) + if state.paren_depth == 0: + return i + return -1 + + @staticmethod + def _is_inside_quotes_or_brackets(text: str, pos: int) -> bool: + """Check if a position in text is inside quotes or brackets. + + Args: + text: The string to check + pos: The position to check (0-based index) + + Returns: + True if the position is inside double quotes, single quotes, or brackets. + False otherwise, or if pos is out of bounds. + + Note: + This function properly handles escaped quotes (preceded by backslash). + Quotes inside brackets don't affect bracket state tracking. + + Example: + >>> _is_inside_quotes_or_brackets('a + "b" + c', 5) + True # Position 5 is inside quotes + >>> _is_inside_quotes_or_brackets('a + "b" + c', 9) + False # Position 9 is outside quotes + """ + if pos >= len(text): + return False + + state = _ParseState() + for i in range(pos): + state.update(text[i], text[i - 1] if i > 0 else None) + return state.in_quotes_or_brackets + def visit_select( self, select_stmt: selectable.Select, @@ -107,7 +230,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 +265,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) @@ -192,6 +321,141 @@ def _legacy_join(self, select_stmt: selectable.Select, **kwargs): def visit_join(self, join, asfrom=True, from_linter=None, **kwargs): return "" + @staticmethod + def _count_outer_parens(text: str) -> tuple[int, str]: + """Count and strip outer parentheses from text. Returns (count, stripped_text).""" + text = text.strip() + count = 0 + while len(text) >= 2 and text[0] == "(" and text[-1] == ")": # noqa: PLR2004 + depth = 0 + for ch in text[:-1]: # Scan all but last char + depth += (ch == "(") - (ch == ")") + if depth == 0: + return count, text # First '(' closed before end + count += 1 + text = text[1:-1].strip() + return count, text + + @staticmethod + def _has_operators_outside_quotes(expr: str) -> bool: + """Check if expression has arithmetic operators outside of quoted strings and brackets.""" + return any(KustoKqlCompiler._find_top_level_operator(expr, op) != -1 for op in "+-*/") + + @staticmethod + def _wrap_column_refs_in_parens(expr: str) -> str: + """Wrap bracket-quoted column refs in parens for arithmetic precedence, unless already wrapped.""" + def wrap_col_ref(m: re.Match[str]) -> str: + if m.start() > 0 and expr[m.start() - 1] == "(": + return m.group(1) + return f"({m.group(1)})" + return re.sub(r'(\["(?:[^"\\]|\\.)*"\])', wrap_col_ref, expr) + + @staticmethod + def _extract_aggregates_from_expression( + expr: str, measure_name: str, existing_aggs: dict[str, str] | None = None + ) -> tuple[str, list[tuple[str, str]]]: + r""" + Extract aggregate functions from an expression and replace with references. + + This function uses a regex-based approach with a pre-compiled pattern (KQL_AGG_PATTERN) + for performance. It properly handles escaped quotes, nested parentheses, and correctly + distinguishes between aggregate functions and quoted text containing aggregate keywords. + + Args: + expr: The expression to process (may contain aggregates, operators, quoted strings) + measure_name: Name of the parent measure (used for generating new ref names) + existing_aggs: Dict mapping kql_agg (lowercase) -> ref_name for reuse. + Allows multiple expressions to share the same aggregate computation. + + Returns: + A tuple of: + - modified expression with aggregates replaced by references like ["__measure_1"] + - list of (ref_name, kql_aggregate) tuples to add to summarize (only NEW ones) + + Example: + Input: "((COUNT(UserInfo_Ring)))", "Measure 4", {"count([\"userinfo_ring\"])": '["Existing"]'} + Output: ('((["Existing"]))', []) # Reuses existing, no new aggregates + + Input: "count(x) + sum(y)", "MyMeasure", {} + Output: ('["__MyMeasure_1"] + ["__MyMeasure_2"]', + [('["__MyMeasure_1"]', 'count(["x"])'), + ('["__MyMeasure_2"]', 'sum(["y"])')] + + Complexity: + O(n * m) where n = len(expr), m = number of aggregate matches + """ + if existing_aggs is None: + existing_aggs = {} + + new_aggregates = [] + agg_counter = 0 + + # Collect replacements: (start, end, ref_name) + replacements = [] + + # Use pre-compiled pattern for performance (avoids recompiling on every call) + for match in KQL_AGG_PATTERN.finditer(expr): + start = match.start() + + # Skip if inside quotes or brackets + if KustoKqlCompiler._is_inside_quotes_or_brackets(expr, start): + continue + + # Find matching closing paren + paren_start = match.end() - 1 + paren_end = KustoKqlCompiler._find_matching_paren(expr, paren_start) + if paren_end == -1: + continue + + # Extract and convert the function call + func_call = expr[start:paren_end + 1] + kql_agg = KustoKqlCompiler._extract_maybe_agg_column_parts(func_call) + if not kql_agg: + continue + + kql_agg_lower = kql_agg.lower() + + # Reuse existing aggregate or create new one + if kql_agg_lower in existing_aggs: + ref_name = existing_aggs[kql_agg_lower] + else: + agg_counter += 1 + clean_name = measure_name.strip('[]"') + ref_name = f'["__{clean_name}_{agg_counter}"]' + existing_aggs[kql_agg_lower] = ref_name + new_aggregates.append((ref_name, kql_agg)) + + replacements.append((start, paren_end + 1, ref_name)) + + # Apply replacements from right to left so positions stay valid + result = expr + for start, end, ref_name in reversed(replacements): + result = result[:start] + ref_name + result[end:] + + return result, new_aggregates + + @staticmethod + def _contains_aggregate_function(expr: str) -> bool: + """Check if expression contains an aggregate function call (even inside parens). + + This is an optimized check that uses the pre-compiled pattern and only checks + for existence without doing full extraction (no reference creation or dict updates). + Much faster than _extract_aggregates_from_expression when you only need a boolean. + + Args: + expr: The expression to check + + Returns: + True if the expression contains at least one aggregate function call + outside of quotes/brackets. False otherwise. + """ + # Optimized: just check if pattern matches outside quotes/brackets + # No need to do full extraction with reference creation + for match in KQL_AGG_PATTERN.finditer(expr): + if not KustoKqlCompiler._is_inside_quotes_or_brackets(expr, match.start()): + return True + return False + def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, str]: """Builds the ending part of the query either project or summarize.""" columns = select.inner_columns @@ -212,47 +476,86 @@ 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() + # Convert to list to allow multiple iterations + columns_list = [c for c in columns if c.name != "*"] + + summarize_columns = [] # Use list to maintain order + extend_columns = [] # Use list to maintain order projection_columns = [] - for column in [c for c in columns if c.name != "*"]: + # Track intermediary measures (should not appear in project) + intermediary_aliases = set() + # Track existing aggregates: kql_agg (lowercase) -> ref_name + # This allows reuse of already-defined aggregates + existing_aggs: dict[str, str] = {} + + # Process all columns in a single pass + for column in columns_list: + # Extract and normalize column information 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 ? - # Do we have aggregate columns ? + column_name = re.sub( + r'(?:[a-zA-Z_][a-zA-Z0-9_]*|\["[^"]+"\])\.', "", column_name + ) + column_alias_escaped = self._escape_and_quote_columns(column_alias, True) + + # Analyze column expression kql_agg = self._extract_maybe_agg_column_parts(column_name) - if kql_agg: + has_operators = self._has_operators_outside_quotes(column_name) + contains_agg = self._contains_aggregate_function(column_name) + + # Case 1: Simple aggregate (e.g., count(), sum(col)) + if kql_agg and not has_operators: has_aggregates = True - summarize_columns.add( - self._build_column_projection(kql_agg, 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) - ) - if column_alias: - projection_columns.append( - self._escape_and_quote_columns(column_alias, True) - ) + summarize_entry = self._build_column_projection(kql_agg, column_alias_escaped) + if summarize_entry not in summarize_columns: + summarize_columns.append(summarize_entry) + projection_columns.append(column_alias_escaped) + # Register this aggregate for reuse by later columns + if column_alias_escaped: + existing_aggs[kql_agg.lower()] = column_alias_escaped + + # Case 2 & 3: Expressions with aggregates or aliased columns (both go to extend) + elif contains_agg or ( + column_alias_escaped + and column_alias_escaped != self._escape_and_quote_columns(column_name) + ): + # If contains aggregates, extract them first + if contains_agg: + has_aggregates = True + column_name, extracted_aggs = self._extract_aggregates_from_expression( + column_name, column_alias or "expr", existing_aggs + ) + + # Add extracted aggregates to summarize + for ref_name, kql_agg in extracted_aggs: + summarize_entry = f"{ref_name} = {kql_agg}" + if summarize_entry not in summarize_columns: + summarize_columns.append(summarize_entry) + intermediary_aliases.add(ref_name) + + # Build extend entry (common for both cases) + escaped_expr = self._escape_and_quote_columns(column_name) + if has_operators: + escaped_expr = self._wrap_column_refs_in_parens(escaped_expr) + extend_entry = f"{column_alias_escaped} = {escaped_expr}" + if extend_entry not in extend_columns: + extend_columns.append(extend_entry) + projection_columns.append(column_alias_escaped) + + # Case 4: Simple column reference else: - projection_columns.append( - self._escape_and_quote_columns(column_name) - ) + projection_columns.append(self._escape_and_quote_columns(column_name)) + # group by columns by_columns = self._group_by(group_by_cols) - if has_aggregates or bool( - by_columns - ): # Summarize can happen with or without aggregate being created + if has_aggregates or bool(by_columns): summarize_statement = f"| summarize {', '.join(summarize_columns)} " if by_columns: - summarize_statement = ( - f"{summarize_statement} by {', '.join(by_columns)}" - ) + summarize_statement = f"{summarize_statement} by {', '.join(by_columns)}" if extend_columns: - extend_statement = f"| extend {', '.join(sorted(extend_columns))}" + extend_statement = f"| extend {', '.join(extend_columns)}" + + # Filter out intermediary aliases from projection + projection_columns = [p for p in projection_columns if p not in intermediary_aliases] project_statement = ( f"| project {', '.join(projection_columns)}" if projection_columns @@ -271,10 +574,10 @@ def _get_projection_or_summarize(self, select: selectable.Select) -> dict[str, s @staticmethod def _extract_maybe_agg_column_parts(column_name) -> str | None: + # First check AGGREGATE_PATTERN which handles SQL-style aggregates (count, sum, etc.) + # including count(distinct X) syntax match_agg_cols = re.match(AGGREGATE_PATTERN, column_name, re.IGNORECASE) if match_agg_cols and match_agg_cols.groups(): - # Check if the aggregate function is count_distinct. This is case from superset - # where we can use count(distinct or count_distinct) aggregate_func, distinct_keyword, agg_column_name, extra_params = ( match_agg_cols.groups() ) @@ -286,8 +589,23 @@ def _extract_maybe_agg_column_parts(column_name) -> str | None: ) return kql_agg - maybe_aggregation_function = column_name.lower().split("(")[0] + # Check if it's a KQL-specific aggregate function not covered by AGGREGATE_PATTERN + maybe_aggregation_function = column_name.lower().split("(")[0].strip() if maybe_aggregation_function in kql_aggregates: + # Multi-arg KQL aggregate (e.g., percentile(col, 99), dcountif(col, predicate)) + # Match func(first_arg, rest...) pattern + match_multi = re.match( + r"(\w+)\s*\(\s*([^,]+?)\s*,\s*(.+)\s*\)$", column_name, re.IGNORECASE + ) + if match_multi: + func_name = match_multi.group(1).lower() + first_arg = match_multi.group(2).strip() + rest_args = match_multi.group(3).strip() + # Escape first arg if it's a column name (not a number/literal) + if not KustoKqlCompiler._is_number_literal(first_arg): + first_arg = KustoKqlCompiler._escape_and_quote_columns(first_arg) + return f"{func_name}({first_arg}, {rest_args})" + # Single-arg KQL aggregate (e.g., countif(predicate)) - return as-is return column_name return None @@ -345,7 +663,8 @@ def replacer(match): # Apply transformation modified_expression = re.sub(pattern, replacer, kql_expression) - return modified_expression + # Convert remaining standalone "col" -> ["col"], skip already bracketed + return re.sub(r'(? str: @@ -358,23 +677,33 @@ def _escape_and_quote_columns(name: str | None, is_alias=False) -> str: return name if name.startswith('"') and name.endswith('"'): name = name[1:-1] - # First, check if the name is already wrapped in ["ColumnName"] (escaped format) if name.startswith('["') and name.endswith('"]'): return name # Return as is if already properly escaped - # Remove surrounding spaces - # Handle mathematical operations (wrap only the column part before operators) - # Find the position of the first operator or space that separates the column name + # Handle arithmetic expressions by recursively processing operands if not is_alias: + outer_paren_count, inner = KustoKqlCompiler._count_outer_parens(name) for operator in ["/", "+", "-", "*"]: - if operator in name: - # Split the name at the first operator and wrap the left part - parts = name.split(operator, 1) - # Remove quotes if they exist at the edges - col_part = parts[0].strip() - 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 + pos = KustoKqlCompiler._find_top_level_operator(inner, operator) + if pos != -1: + left = KustoKqlCompiler._escape_and_quote_columns( + inner[:pos].strip() + ) + right = KustoKqlCompiler._escape_and_quote_columns( + inner[pos + 1 :].strip() + ) + return ( + "(" * outer_paren_count + + left + + " " + + operator + + " " + + right + + ")" * outer_paren_count + ) + # No operators - recurse on inner content if we stripped parens + if outer_paren_count > 0: + inner_result = KustoKqlCompiler._escape_and_quote_columns(inner) + return "(" * outer_paren_count + inner_result + ")" * outer_paren_count # No operators found, just wrap the entire name name = name.replace('"', '\\"') return f'["{name}"]' @@ -547,8 +876,11 @@ def _is_kql_function(name: str) -> bool: @staticmethod def _is_number_literal(s: str) -> bool: - pattern = r"^[0-9]+$" - return bool(re.match(pattern, s)) + """Check if string is a numeric literal (integer or floating point).""" + # Match integers, decimals, and scientific notation + # Examples: 5, 0.5, .5, 5., 5.0, 1e10, 1.5e-3 + pattern = r"^-?(\d+\.?\d*|\d*\.?\d+)([eE][+-]?\d+)?$" + return bool(re.match(pattern, s.strip())) def _get_most_inner_element(self, clause): """Finds the most nested element in clause.""" @@ -644,8 +976,8 @@ def _sql_to_kql_aggregate( return_value = None # The count function is a special case because it can be used with or without a column name # We can also use it in count(Distinct column_name) format. This has to be handled separately - if sql_agg and sql_agg in ("count", "COUNT"): - if "*" in sql_agg or column_name in ("*", "1"): + if sql_agg and sql_agg.lower() in ("count",): + if "*" in str(column_name) or column_name in ("*", "1"): return_value = aggregates_sql_to_kql["count(*)"] elif is_distinct: return_value = f"dcount({column_name_escaped})" @@ -661,9 +993,7 @@ def _sql_to_kql_aggregate( if sql_to_kql_aggregate_function: return_value = f"{sql_to_kql_aggregate_function}({column_name_escaped})" elif aggregation_function in kql_aggregates: - return_value = ( - f"{aggregation_function}({column_name_escaped}{extra_params})" - ) + return_value = f"{aggregation_function}({column_name_escaped}{extra_params if extra_params else ''})" return return_value diff --git a/tests/unit/test_dialect_kql.py b/tests/unit/test_dialect_kql.py index 04b8fc0..e159839 100644 --- a/tests/unit/test_dialect_kql.py +++ b/tests/unit/test_dialect_kql.py @@ -9,6 +9,7 @@ column, create_engine, distinct, + func, literal_column, select, text, @@ -24,14 +25,10 @@ def test_compiler_with_projection(): statement_str = "logs | take 10" stmt = TextAsFrom(sa.text(statement_str), []).alias("virtual_table") query = sa.select( - from_obj=stmt, - columns=[ - column("Id").label("id"), - column("TypeId").label("tId"), - column("Type"), - ], - ) - query = query.select_from(stmt) + column("Id").label("id"), + column("TypeId").label("tId"), + column("Type"), + ).select_from(stmt) query = query.limit(10) query_compiled = str(query.compile(engine)).replace("\n", "") @@ -49,11 +46,7 @@ def test_compiler_with_projection(): def test_compiler_with_star(): statement_str = "logs | take 10" stmt = TextAsFrom(sa.text(statement_str), []).alias("virtual_table") - query = sa.select( - "*", - from_obj=stmt, - ) - query = query.select_from(stmt) + query = sa.select("*").select_from(stmt) query = query.limit(10) query_compiled = str(query.compile(engine)).replace("\n", "") query_expected = ( @@ -66,9 +59,7 @@ def test_compiler_with_star(): def test_select_from_text(): query = ( - select([column("Field1"), column("Field2")]) - .select_from(text("logs")) - .limit(100) + select(column("Field1"), column("Field2")).select_from(text("logs")).limit(100) ) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) @@ -148,7 +139,7 @@ def test_select_from_text(): ) def test_where_predicates(f, expected): query = ( - select([column("Field1"), column("Field2")]).select_from(text("logs")).where(f) + select(column("Field1"), column("Field2")).select_from(text("logs")).where(f) ).limit(100) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) @@ -164,7 +155,7 @@ def test_group_by_text(): event_col = literal_column('"EventInfo_Time" / time(1d)').label("EventInfo_Time") active_users_col = literal_column("ActiveUsers").label("ActiveUserMetric") query = ( - select([event_col, active_users_col]) + select(event_col, active_users_col) .select_from(text("ActiveUsersLastMonth")) .group_by(literal_column('"EventInfo_Time" / time(1d)')) .order_by(text("ActiveUserMetric DESC")) @@ -173,15 +164,16 @@ 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 - extend column order follows select order query_expected = ( - '["ActiveUsersLastMonth"]| extend ["ActiveUserMetric"] = ["ActiveUsers"], ' - '["EventInfo_Time"] = ["EventInfo_Time"] / time(1d)' - '| summarize by ["EventInfo_Time"] / time(1d)' + '["ActiveUsersLastMonth"]| summarize by ["EventInfo_Time"] / time(1d)' + '| extend ["EventInfo_Time"] = (["EventInfo_Time"]) / time(1d), ' + '["ActiveUserMetric"] = ["ActiveUsers"]' '| project ["EventInfo_Time"], ["ActiveUserMetric"]' '| order by ["ActiveUserMetric"] desc' ) assert query_compiled == query_expected + assert query_compiled == query_expected @pytest.mark.parametrize( @@ -195,18 +187,19 @@ def test_function_text(f: str, expected: str): # create a query from select_query_text creating clause event_col = literal_column(f).label("EventInfo_Time") active_users_col = literal_column("ActiveUsers").label("ActiveUserMetric") - query = select([event_col, active_users_col]).select_from( + query = select(event_col, active_users_col).select_from( text("ActiveUsersLastMonth") ) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) ).replace("\n", "") + # extend columns follow select order 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 @@ -215,7 +208,7 @@ def test_group_by_text_vaccine_dataset(): # SQL: SELECT country_name AS country_name FROM superset."CovidVaccineData" GROUP BY country_name # ORDER BY country_name ASC - this is a simple query to get distinct country names query = ( - select([literal_column("country_name").label("country_name")]) + select(literal_column("country_name").label("country_name")) .select_from(text('superset."CovidVaccineData"')) .group_by(literal_column("country_name")) .order_by(text("country_name ASC")) @@ -225,7 +218,6 @@ def test_group_by_text_vaccine_dataset(): ).replace("\n", "") query_expected = ( 'database("superset").["CovidVaccineData"]| ' - 'extend ["country_name"] = ["country_name"]| ' 'summarize by ["country_name"]| ' 'project ["country_name"]| order by ["country_name"] asc' ) @@ -233,11 +225,9 @@ def test_group_by_text_vaccine_dataset(): 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")""" @@ -247,9 +237,7 @@ def test_is_kql_function(): def test_percentile_by_text(): event_col = literal_column("percentile(quantity_ordered, 99)").label("Measure 1") query = select( - [ - event_col, - ] + event_col, ).select_from(text("SalesData")) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) @@ -268,9 +256,7 @@ def test_dcountif_by_text(): "dcountif(year, city == 'Paris' or city in ('Madrid'))" ).label("Measure 1") query = select( - [ - event_col, - ] + event_col, ).select_from(text("SalesData")) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) @@ -289,14 +275,12 @@ def test_countif_by_text(): "Measure 1" ) query = select( - [ - event_col, - ] + event_col, ).select_from(text("SalesData")) query_compiled = str( query.compile(engine, compile_kwargs={"literal_binds": True}) ).replace("\n", "") - # raw query text from query + # raw query text from query - column names in predicate passed as-is query_expected = ( '["SalesData"]' "| summarize [\"Measure 1\"] = countif(city == 'Paris' OR city in ('Madrid')) " @@ -313,10 +297,8 @@ def test_distinct_count_by_text(): active_users_col = literal_column("ActiveUsers") query = ( select( - [ - event_col, - sa.func.count(distinct(active_users_col)).label("DistinctUsers"), - ] + event_col, + sa.func.count(distinct(active_users_col)).label("DistinctUsers"), ) .select_from(text("ActiveUsersLastMonth")) .group_by(literal_column('"EventInfo_Time" / time(1d)')) @@ -328,8 +310,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' ) @@ -343,7 +325,7 @@ def test_distinct_count_alt_by_text(): event_col = literal_column("EventInfo_Time / time(1d)").label("EventInfo_Time") active_users_col = literal_column("COUNT_DISTINCT(ActiveUsers)") query = ( - select([event_col, active_users_col.label("DistinctUsers")]) + select(event_col, active_users_col.label("DistinctUsers")) .select_from(text("ActiveUsersLastMonth")) .group_by(literal_column("EventInfo_Time / time(1d)")) .order_by(text("ActiveUserMetric DESC")) @@ -354,8 +336,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' ) @@ -412,7 +394,7 @@ def test_select_count(): kql_query = "logs" column_count = literal_column("count(*)").label("total-count") query = ( - select([column_count]) + select(column_count) .select_from(TextAsFrom(text(kql_query), ["*"]).alias("inner_qry")) .where(text("Field1 > 1")) .where(text("Field2 < 2")) @@ -576,3 +558,815 @@ def test_schema_from_query(query_table_name: str, expected_table_name: str): query_expected = f"let inner_qry = ({expected_table_name});inner_qry| take 5" assert query_compiled == query_expected + + +class TestCalculatedMeasures: + """Tests for calculated measures (arithmetic expressions with aggregates).""" + + @pytest.fixture + def events_table(self): + metadata = MetaData() + return Table( + "events", + metadata, + Column("region", String), + Column("ring", String), + Column("value", Integer), + ) + + @pytest.fixture + def pt_search_table(self): + """Table matching the Superset PT_Search_scenario use case.""" + metadata = MetaData() + return Table( + "PT_Search_scenario", + metadata, + Column("UserInfo_Ring", String), + Column("UserInfo_Region", String), + schema="bc3902d8132f43e3ae086a009979fa88", + ) + + def test_multi_aggregate_expression(self, events_table): + """Test that expressions with multiple aggregates generate correct KQL.""" + query = select( + (func.count(events_table.c.region) + func.count(events_table.c.ring)).label( + "multi_agg" + ) + ).select_from(events_table) + + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should have summarize with both aggregates + assert "summarize" in compiled + # Should have extend for the calculated measure + assert "extend" in compiled + # Should project the alias + assert '["multi_agg"]' in compiled + + def test_arithmetic_expression_with_columns(self, events_table): + """Test arithmetic expressions with column references.""" + query = select( + (events_table.c.value / literal_column("100")).label("percentage") + ).select_from(events_table) + + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should handle division operator + assert "/" in compiled + assert '["percentage"]' in compiled + + def test_escape_and_quote_columns_with_arithmetic(self): + """Test _escape_and_quote_columns handles arithmetic expressions.""" + result = KustoKqlCompiler._escape_and_quote_columns("col1 + col2") + assert '["col1"]' in result + assert '["col2"]' in result + assert "+" in result + + def test_escape_and_quote_columns_with_parentheses(self): + """Test _escape_and_quote_columns handles parenthesized expressions.""" + result = KustoKqlCompiler._escape_and_quote_columns("(col1 + col2)") + assert result.startswith("(") + assert result.endswith(")") + assert '["col1"]' in result + assert '["col2"]' in result + + def test_has_operators_outside_quotes(self): + """Test detection of arithmetic operators outside quoted strings.""" + assert KustoKqlCompiler._has_operators_outside_quotes("a + b") is True + assert KustoKqlCompiler._has_operators_outside_quotes("a - b") is True + assert KustoKqlCompiler._has_operators_outside_quotes("a * b") is True + assert KustoKqlCompiler._has_operators_outside_quotes("a / b") is True + assert KustoKqlCompiler._has_operators_outside_quotes('["col"]') is False + assert KustoKqlCompiler._has_operators_outside_quotes('"a + b"') is False + + def test_count_outer_parens(self): + """Test counting and stripping outer parentheses.""" + count, inner = KustoKqlCompiler._count_outer_parens("((a + b))") + assert count == 2 # noqa: PLR2004 + assert inner == "a + b" + + count, inner = KustoKqlCompiler._count_outer_parens("(a) + (b)") + assert count == 0 + assert inner == "(a) + (b)" + + def test_predefined_measures_lowercase(self, pt_search_table): + """Test that predefined measures (aggregates) compile to lowercase KQL functions.""" + userinfo_ring_count = func.COUNT(pt_search_table.c.UserInfo_Ring).label( + "UserInfo_Ring Count" + ) + userinfo_region_count = func.COUNT(pt_search_table.c.UserInfo_Region).label( + "UserInfo_Region Count" + ) + + query = select(userinfo_ring_count, userinfo_region_count).select_from( + pt_search_table + ) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should use lowercase count function in KQL + assert 'count(["UserInfo_Ring"])' in compiled or "count([" in compiled + # Should NOT have uppercase COUNT + assert "COUNT(" not in compiled + + def test_calculated_measure_simple_reference(self, pt_search_table): + """Test a calculated measure that's just a reference to another measure.""" + measure_15 = literal_column('"UserInfo_Region Count"').label("Measure 15") + + query = select(measure_15).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should convert quoted identifier to bracket notation + assert '["Measure 15"]' in compiled + assert '["UserInfo_Region Count"]' in compiled + + def test_calculated_measure_single_paren(self, pt_search_table): + """Test a calculated measure with single parentheses wrapper.""" + measure_16 = literal_column('("UserInfo_Ring Count")').label("Measure 16") + + query = select(measure_16).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 16"]' in compiled + + def test_calculated_measure_double_paren(self, pt_search_table): + """Test a calculated measure with double parentheses wrapper.""" + measure_3 = literal_column('(("Measure 1"))').label("Measure 3") + + query = select(measure_3).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 3"]' in compiled + # Should preserve double parens + assert "((" in compiled + assert "))" in compiled + + def test_calculated_measure_multiply_by_constant(self, pt_search_table): + """Test a calculated measure that multiplies a reference by a constant.""" + measure_9 = literal_column('"UserInfo_Ring Count" * 2').label("Measure 9") + + query = select(measure_9).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 9"]' in compiled + assert "* 2" in compiled + + def test_calculated_measure_addition(self, pt_search_table): + """Test a calculated measure that adds two measure references.""" + measure_14 = literal_column( + '"UserInfo_Region Count" + "UserInfo_Ring Count"' + ).label("Measure 14") + + query = select(measure_14).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 14"]' in compiled + assert "+" in compiled + assert '["UserInfo_Region Count"]' in compiled + assert '["UserInfo_Ring Count"]' in compiled + + def test_calculated_measure_parens_addition(self, pt_search_table): + """Test a calculated measure with parenthesized addition.""" + measure_11 = literal_column('("Measure 1") + ("Measure 2")').label("Measure 11") + + query = select(measure_11).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 11"]' in compiled + assert "+" in compiled + + def test_calculated_measure_complex_expression(self, pt_search_table): + """Test a complex calculated measure with nested parens and multiplication.""" + measure_8 = literal_column( + '("UserInfo_Ring Count" + "UserInfo_Region Count") * 2' + ).label("Measure 8") + + query = select(measure_8).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 8"]' in compiled + assert "* 2" in compiled + assert "+" in compiled + + def test_calculated_measure_plus_constant(self, pt_search_table): + """Test a calculated measure that adds a constant.""" + measure_20 = literal_column('"Measure 1" + 1').label("Measure 20") + + query = select(measure_20).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Measure 20"]' in compiled + assert "+ 1" in compiled + + def test_no_double_bracketing(self, pt_search_table): + """Test that there's no double bracketing like [["col"]].""" + measure = literal_column('"UserInfo_Ring Count"').label("Test Measure") + + query = select(measure).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should not have double brackets + assert '[["' not in compiled + assert '"]]' not in compiled + + def test_standalone_quoted_identifier(self): + """Test that standalone quoted identifiers are converted to bracket notation.""" + metadata = MetaData() + test_table = Table( + "TestTable", + metadata, + Column("Revenue", String), + Column("Cost", String), + schema="test_schema", + ) + + measure_standalone = literal_column('"Revenue"').label("Standalone Quote") + + query = select(measure_standalone).select_from(test_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Standalone Quote"]' in compiled + assert '["Revenue"]' in compiled + + def test_standalone_quoted_expression(self): + """Test standalone expression with quoted identifiers.""" + metadata = MetaData() + test_table = Table( + "TestTable", + metadata, + Column("Revenue", String), + Column("Cost", String), + schema="test_schema", + ) + + measure_expr = literal_column('"Revenue" + "Cost"').label( + "Standalone Expression" + ) + + query = select(measure_expr).select_from(test_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert '["Standalone Expression"]' in compiled + assert '["Revenue"]' in compiled + assert '["Cost"]' in compiled + assert "+" in compiled + + def test_calculated_measure_references_simple_measures(self, pt_search_table): + """Test that calculated measures can reference other measures by name. + + This simulates the Superset UI where: + - Measure 1 = count() + - Measure 4 = (("Measure 1")) # References Measure 1 by name + + The extend should reference the measure name, not the raw SQL. + """ + # Simple measures + measure_1 = literal_column("count()").label("Measure 1") + measure_2 = literal_column("count()").label("Measure 2") + + # Calculated measure that references Measure 1 by name + measure_4 = literal_column('(("Measure 1"))').label("Measure 4") + + query = select(measure_1, measure_2, measure_4).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Summarize should have the base measures + assert 'summarize' in compiled + assert '["Measure 1"] = count()' in compiled + assert '["Measure 2"] = count()' in compiled + + # Extend should reference ["Measure 1"] not count() + assert 'extend' in compiled + assert '["Measure 4"]' in compiled + # Should reference the measure, not raw SQL + assert '(["Measure 1"])' in compiled or '((["Measure 1"]))' in compiled + + def test_calculated_measure_with_arithmetic_on_measure_refs(self, pt_search_table): + """Test calculated measures with arithmetic on measure references.""" + measure_1 = literal_column("count()").label("Measure 1") + measure_2 = literal_column("count()").label("Measure 2") + + # Measure 6 = Measure 1 + Measure 2 + measure_6 = literal_column('"Measure 1" + "Measure 2"').label("Measure 6") + + query = select(measure_1, measure_2, measure_6).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + assert 'summarize' in compiled + assert 'extend' in compiled + # The calculated measure should reference the measure names + assert '["Measure 6"]' in compiled + assert '["Measure 1"]' in compiled + assert '["Measure 2"]' in compiled + + def test_no_aggregates_in_extend(self, pt_search_table): + """Verify that aggregate functions don't appear in extend statements.""" + measure_1 = literal_column("count()").label("Measure 1") + measure_4 = literal_column('(("Measure 1"))').label("Measure 4") + + query = select(measure_1, measure_4).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Find the extend part + extend_idx = compiled.find('extend') + if extend_idx != -1: + project_idx = compiled.find('| project') + extend_part = compiled[extend_idx:project_idx] if project_idx != -1 else compiled[extend_idx:] + # Should not have count() in extend - it should reference ["Measure 1"] + assert 'count()' not in extend_part.lower() + # Should have the measure reference instead + assert '["Measure 1"]' in extend_part + + def test_measure_name_with_aggregate_keyword(self, pt_search_table): + """Test that measure names containing aggregate keywords (like 'Count') aren't parsed as aggregates. + + This tests the case where a measure is named "UserInfo_Ring Count" - the word "Count" + should NOT be treated as an aggregate function. + """ + # Base measures with aggregate keywords in their names + ring_count = func.COUNT(pt_search_table.c.UserInfo_Ring).label("UserInfo_Ring Count") + region_count = func.COUNT(pt_search_table.c.UserInfo_Region).label("UserInfo_Region Count") + + # Calculated measure referencing measure with "Count" in its name + measure_4 = literal_column('(("UserInfo_Ring Count"))').label("Measure 4") + + query = select(ring_count, region_count, measure_4).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Find the extend part + extend_idx = compiled.find('extend') + if extend_idx != -1: + project_idx = compiled.find('| project') + extend_part = compiled[extend_idx:project_idx] if project_idx != -1 else compiled[extend_idx:] + + # Should NOT have COUNT(UserInfo_Ring) in extend - that's the bug we're fixing + assert 'COUNT(' not in extend_part + assert 'count(' not in extend_part + # Should reference the measure name, not the raw SQL + assert '["UserInfo_Ring Count"]' in extend_part + + def test_measure_name_with_sum_keyword(self, pt_search_table): + """Test that measure names containing 'Sum' aren't parsed as aggregates.""" + # A measure named "Total Sum" should not have "Sum" treated as an aggregate + base_measure = literal_column("count()").label("Total Sum") + calc_measure = literal_column('"Total Sum" * 2').label("Double Sum") + + query = select(base_measure, calc_measure).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # The extend should reference ["Total Sum"], not try to parse "Sum" as aggregate + extend_idx = compiled.find('extend') + if extend_idx != -1: + project_idx = compiled.find('| project') + extend_part = compiled[extend_idx:project_idx] if project_idx != -1 else compiled[extend_idx:] + assert '["Total Sum"]' in extend_part + # Should not have sum() function call in extend + assert 'sum(' not in extend_part.lower() or '["Total Sum"]' in extend_part + + def test_aggregate_in_quoted_string_not_extracted(self): + """Test that aggregates inside quoted strings are not extracted.""" + # Expression with "Count" inside a quoted measure name + expr = '(("UserInfo_Ring Count"))' + result, new_aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Test") + + # Should NOT extract any aggregates - "Count" is inside quotes + assert len(new_aggs) == 0 + # Expression should be unchanged (except for normal bracket escaping) + assert 'count(' not in result.lower() + + def test_real_aggregate_still_extracted(self): + """Test that real aggregate functions are still properly extracted.""" + # Expression with actual aggregate function + expr = 'count(col1) + sum(col2)' + result, new_aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Test") + + # Should extract both aggregates + assert len(new_aggs) == 2 + # Result should have references, not the original aggregates + assert 'count(' not in result.lower() + assert 'sum(' not in result.lower() + + def test_mixed_quoted_and_real_aggregates(self): + """Test expression with both quoted measure names and real aggregates.""" + # "Ring Count" is a measure name (quoted), count(col) is a real aggregate + expr = '"Ring Count" + count(col)' + result, new_aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Test") + + # Should extract only the real aggregate, not the one in quotes + assert len(new_aggs) == 1 + agg_sql = new_aggs[0][1] + assert 'count(' in agg_sql.lower() + + def test_bracket_notation_not_extracted(self): + """Test that aggregates in bracket notation are not extracted.""" + # Expression with "Count" inside bracket notation + expr = '["UserInfo_Ring Count"] * 2' + result, new_aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Test") + + # Should NOT extract any aggregates - "Count" is inside brackets + assert len(new_aggs) == 0 + + def test_wrapped_aggregate_extracted_correctly(self, pt_search_table): + """Test that aggregates wrapped in parens (like ((COUNT(col)))) are extracted correctly.""" + # This is what Superset sends when a user writes (("UserInfo_Ring Count")) + # Superset resolves the measure reference to the actual SQL + measure_4 = literal_column("((COUNT(UserInfo_Ring)))").label("Measure 4") + + query = select(measure_4).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Should have summarize with the aggregate + assert 'summarize' in compiled + + # Find the extend part + extend_idx = compiled.find('extend') + if extend_idx != -1: + project_idx = compiled.find('| project') + extend_part = compiled[extend_idx:project_idx] if project_idx != -1 else compiled[extend_idx:] + + # Should NOT have COUNT() in extend + assert 'COUNT(' not in extend_part + assert 'count(' not in extend_part + # Should have a reference + assert '["Measure 4"]' in extend_part + + def test_floating_point_numbers(self, pt_search_table): + """Test that floating point numbers are preserved correctly.""" + # Measure with floating point multiplier + measure_1 = literal_column("count()").label("Measure 1") + measure_2 = literal_column('"Measure 1" * 0.5').label("Measure 2") + measure_3 = literal_column('"Measure 1" * 1.25').label("Measure 3") + measure_4 = literal_column('"Measure 1" / 0.1').label("Measure 4") + + query = select(measure_1, measure_2, measure_3, measure_4).select_from(pt_search_table) + compiled = str(query.compile(engine, compile_kwargs={"literal_binds": True})) + + # Floating point numbers should be preserved, not wrapped in brackets + assert '* 0.5' in compiled + assert '* 1.25' in compiled + assert '/ 0.1' in compiled + # Should NOT have bracketed numbers + assert '["0.5"]' not in compiled + assert '["1.25"]' not in compiled + assert '["0.1"]' not in compiled + + def test_is_number_literal(self): + """Test _is_number_literal handles various number formats.""" + # Integers + assert KustoKqlCompiler._is_number_literal("5") is True + assert KustoKqlCompiler._is_number_literal("123") is True + assert KustoKqlCompiler._is_number_literal("0") is True + + # Floating point + assert KustoKqlCompiler._is_number_literal("0.5") is True + assert KustoKqlCompiler._is_number_literal("1.25") is True + assert KustoKqlCompiler._is_number_literal(".5") is True + assert KustoKqlCompiler._is_number_literal("5.") is True + assert KustoKqlCompiler._is_number_literal("0.0") is True + + # Negative numbers + assert KustoKqlCompiler._is_number_literal("-5") is True + assert KustoKqlCompiler._is_number_literal("-0.5") is True + + # Scientific notation + assert KustoKqlCompiler._is_number_literal("1e10") is True + assert KustoKqlCompiler._is_number_literal("1.5e-3") is True + + # Not numbers + assert KustoKqlCompiler._is_number_literal("abc") is False + assert KustoKqlCompiler._is_number_literal("1.2.3") is False + assert KustoKqlCompiler._is_number_literal("") is False + + +def test_find_top_level_operator_with_single_quotes(): + """Test that _find_top_level_operator handles single-quoted strings correctly.""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Operator outside quotes should be found + assert KustoKqlCompiler._find_top_level_operator("a + b", "+") == 2 + + # Operator inside double quotes should NOT be found + assert KustoKqlCompiler._find_top_level_operator('"a + b"', "+") == -1 + + # Operator inside single quotes should NOT be found (KQL string literals) + assert KustoKqlCompiler._find_top_level_operator("'value-with-minus'", "-") == -1 + assert KustoKqlCompiler._find_top_level_operator("col + 'test-value'", "-") == -1 + + # Operator outside single quotes should be found + assert KustoKqlCompiler._find_top_level_operator("col + 'test'", "+") == 4 + + # Mixed quotes + assert KustoKqlCompiler._find_top_level_operator("\"col\" + 'value'", "+") == 6 + assert KustoKqlCompiler._find_top_level_operator("'a-b' + \"c-d\"", "+") == 6 + assert KustoKqlCompiler._find_top_level_operator("'a-b' + \"c-d\"", "-") == -1 + + +# ============================================================================== +# Tests for dialect_kql.py improvements (Performance & Correctness) +# ============================================================================== + +def test_precompiled_pattern_exists_and_works(): + """IMPROVEMENT: Verify KQL_AGG_PATTERN is pre-compiled (performance optimization). + + This tests the fix for the performance regression where the regex pattern + was being compiled on every function call. Now it's pre-compiled as a + module-level constant. + """ + from sqlalchemy_kusto.dialect_kql import KQL_AGG_PATTERN + import re + + # Must be a pre-compiled Pattern object, not a string + assert isinstance(KQL_AGG_PATTERN, re.Pattern) + + # Should work correctly + assert KQL_AGG_PATTERN.search("count(x)") is not None + assert KQL_AGG_PATTERN.search("SUM(revenue)") is not None + assert KQL_AGG_PATTERN.search("dcount(users)") is not None + + # Should respect word boundaries + assert KQL_AGG_PATTERN.search("mycount(x)") is None + + +def test_is_inside_quotes_or_brackets_handles_escaped_quotes(): + """IMPROVEMENT: Test that _is_inside_quotes_or_brackets handles escaped quotes correctly. + + This tests the bug fix where escaped quotes were not being properly handled, + which could cause incorrect detection of whether a position is inside quotes. + """ + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Fixed: Escaped double quote should not close the string + text = r'x + "a\"b" + y' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 5) is True # at 'a' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 7) is True # at escaped quote + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 8) is True # at 'b' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 12) is False # at '+' + + # Fixed: Escaped single quote should not close the string + text = r"x + 'a\'b' + y" + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 5) is True # at 'a' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 7) is True # at escaped quote + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 8) is True # at 'b' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(text, 12) is False # at '+' + + # Without escape handling, this would incorrectly think position 12 is inside quotes + + +def test_contains_aggregate_no_unnecessary_extraction(): + """IMPROVEMENT: Test that _contains_aggregate_function is optimized. + + This tests the optimization where _contains_aggregate_function no longer + does a full extraction (creating references and modifying dicts), but just + checks if an aggregate exists outside quotes/brackets. + """ + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Should detect presence of aggregates + assert KustoKqlCompiler._contains_aggregate_function("count(x)") is True + assert KustoKqlCompiler._contains_aggregate_function("sum(a) + avg(b)") is True + assert KustoKqlCompiler._contains_aggregate_function("((COUNT(users)))") is True + + # Should correctly skip aggregates in quotes (bug would return True) + assert KustoKqlCompiler._contains_aggregate_function('"count(x)"') is False + assert KustoKqlCompiler._contains_aggregate_function("'sum is a word'") is False + assert KustoKqlCompiler._contains_aggregate_function('["Count Column"]') is False + + # Mixed: real aggregate + quoted text containing aggregate keywords + assert KustoKqlCompiler._contains_aggregate_function('"Count Text" + count(x)') is True + + # Should not detect non-aggregates + assert KustoKqlCompiler._contains_aggregate_function("column_name") is False + assert KustoKqlCompiler._contains_aggregate_function('"Measure 1" + "Measure 2"') is False + + +def test_extract_aggregates_uses_precompiled_pattern(): + """IMPROVEMENT: Verify _extract_aggregates_from_expression uses pre-compiled pattern. + + This ensures the performance optimization is actually being used by the + extraction function. + """ + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler, KQL_AGG_PATTERN + + # Extract aggregates + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression( + "count(x) + sum(y) + avg(z)", "measure" + ) + + # Should extract all three aggregates + assert len(aggs) == 3 + + # Verify each aggregate is in the list of KQL aggregates + for ref_name, kql_agg in aggs: + # The aggregate should match our pre-compiled pattern + # (this indirectly verifies the function uses the pattern) + assert KQL_AGG_PATTERN.search(kql_agg) is not None + + +def test_escaped_quotes_in_aggregate_extraction(): + """IMPROVEMENT: Test that aggregate extraction handles escaped quotes correctly. + + This ensures the fix for escaped quote handling is applied in the + extraction logic. + """ + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Expression with escaped quotes - aggregate should still be extracted + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression( + r'"text\"more" + count(x)', "measure" + ) + assert len(aggs) == 1 + assert "count" in aggs[0][1].lower() + + # Aggregate inside string with escaped quote should NOT be extracted + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression( + r'"count(\"x\")" + y', "measure" + ) + assert len(aggs) == 0 + + # Both escaped quote and real aggregate + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression( + r'"escaped\"quote" + sum(revenue)', "measure" + ) + assert len(aggs) == 1 + assert "sum" in aggs[0][1].lower() + + +def test_performance_no_regex_recompilation(): + """IMPROVEMENT: Verify regex pattern is not recompiled on each call. + + This is a regression test to ensure the performance fix stays in place. + Multiple calls should use the same compiled pattern object. + """ + from sqlalchemy_kusto.dialect_kql import KQL_AGG_PATTERN, KustoKqlCompiler + + # Get the pattern object id before any calls + pattern_id_before = id(KQL_AGG_PATTERN) + + # Make multiple calls to functions that use the pattern + for _ in range(100): + KustoKqlCompiler._contains_aggregate_function("count(x)") + KustoKqlCompiler._extract_aggregates_from_expression("sum(y)", "test") + + # Pattern object should be the same (not recompiled) + pattern_id_after = id(KQL_AGG_PATTERN) + assert pattern_id_before == pattern_id_after + + +def test_complex_expression_with_all_improvements(): + """INTEGRATION: Test complex expression uses all improvements correctly. + + This integration test verifies that all improvements work together: + - Pre-compiled pattern for performance + - Escaped quote handling for correctness + - Optimized aggregate detection + """ + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Complex expression with: escaped quotes, aggregates, operators, and quoted measure names + expr = r'"Total\"Count" + (count(x) + sum(y)) / avg(z) * "Factor"' + + # Should detect aggregates correctly + contains_agg = KustoKqlCompiler._contains_aggregate_function(expr) + assert contains_agg is True + + # Should extract aggregates correctly + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "ComplexMeasure") + + # Should extract 3 aggregates (count, sum, avg) + assert len(aggs) == 3 + + # Should NOT extract quoted measure names + assert any("count" in agg[1].lower() for agg in aggs) + assert any("sum" in agg[1].lower() for agg in aggs) + assert any("avg" in agg[1].lower() for agg in aggs) + + # Result should have references, not original aggregates + assert "count(x)" not in result.lower() + assert "sum(y)" not in result.lower() + assert "avg(z)" not in result.lower() + + # But should preserve quoted strings + assert r'"Total\"Count"' in result or r'["Total\"Count"]' in result + assert r'"Factor"' in result or r'["Factor"]' in result + + +# ============================================================================== +# Edge Case Tests (From PR Review Feedback) +# ============================================================================== + +def test_large_expression_stress_test(): + """EDGE CASE: Test performance with very large expressions (1000+ characters).""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Generate a large expression with many aggregates + parts = [f"count(col{i})" for i in range(50)] + large_expr = " + ".join(parts) + + # Should handle large expression without errors + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(large_expr, "LargeMeasure") + + # Should extract all 50 aggregates + assert len(aggs) == 50 + + # Should also detect correctly + assert KustoKqlCompiler._contains_aggregate_function(large_expr) is True + + +def test_deeply_nested_parentheses(): + """EDGE CASE: Test handling of deeply nested parentheses (10+ levels).""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Create deeply nested expression (balanced parens: 9 outer + 1 from count() = 10 total opening, 10 closing with extra trailing) + expr = "(((((((((count(x)))))))))))" + + # Should extract the aggregate + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "DeepNest") + assert len(aggs) == 1 + + # Result should maintain outer parentheses minus the function's opening paren + # Original: (((((((((count(x))))))))))) has 10 '(' and 11 ')' + # After extraction: (((((((((ref))))))))))) has 9 '(' and 10 ')' + # The aggregate "count(x)" is replaced with "ref", removing one ( and one ) + assert result.count("(") == 9 + assert result.count(")") == 10 + + # Test _find_matching_paren with deep nesting (using balanced expression) + balanced_expr = "(((((((((())))))))))" # 10 levels deep, balanced: 10 '(' and 10 ')' + assert KustoKqlCompiler._find_matching_paren(balanced_expr, 0) == len(balanced_expr) - 1 # Outermost match + assert KustoKqlCompiler._find_matching_paren(balanced_expr, 5) == len(balanced_expr) - 6 # 5th level match + + +def test_unicode_characters_in_column_names(): + """EDGE CASE: Test unicode characters in column names.""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Unicode column names (e.g., Chinese, Arabic, emoji) + expr = 'count(价格) + sum(المبلغ) + avg(🔥column)' + + # Should handle unicode correctly + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Unicode") + assert len(aggs) == 3 + + # Should escape unicode column names properly + for ref_name, kql_agg in aggs: + assert kql_agg.startswith(("count(", "sum(", "avg(")) + + +def test_multiple_consecutive_escaped_quotes(): + """EDGE CASE: Test multiple consecutive escaped quotes.""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Expression with multiple escaped quotes: "a\\"b\\"c" + expr = r'"a\\"b\\"c" + count(x)' + + # Should correctly identify that count is outside quotes + contains_agg = KustoKqlCompiler._contains_aggregate_function(expr) + assert contains_agg is True + + # Should extract the aggregate + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "EscapedQuotes") + assert len(aggs) == 1 + + # Test _is_inside_quotes_or_brackets at various positions + assert KustoKqlCompiler._is_inside_quotes_or_brackets(expr, 2) is True # at 'a' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(expr, 8) is True # at 'c' + assert KustoKqlCompiler._is_inside_quotes_or_brackets(expr, 13) is False # at '+' + + +def test_empty_and_edge_inputs(): + """EDGE CASE: Test empty strings and boundary conditions.""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Empty string + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression("", "Empty") + assert len(aggs) == 0 + assert result == "" + + # Just an aggregate, no operators + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression("count(x)", "Simple") + assert len(aggs) == 1 + + # Just a column reference + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression("column_name", "Column") + assert len(aggs) == 0 + + # Out of bounds checks for helper functions + assert KustoKqlCompiler._is_inside_quotes_or_brackets("abc", 100) is False + assert KustoKqlCompiler._find_matching_paren("(abc)", 100) == -1 + + +def test_mixed_aggregate_and_string_patterns(): + """EDGE CASE: Test expressions with aggregate keywords in various contexts.""" + from sqlalchemy_kusto.dialect_kql import KustoKqlCompiler + + # Aggregate keyword in column name, measure name, and real aggregate + expr = '"Total Count" + ["Count Column"] + count(actual_count)' + + # Should only detect the real aggregate + result, aggs = KustoKqlCompiler._extract_aggregates_from_expression(expr, "Mixed") + assert len(aggs) == 1 + assert "count" in aggs[0][1].lower() + + # Should preserve quoted strings and bracket notation + assert '"Total Count"' in result or '["Total Count"]' in result + assert '["Count Column"]' in result +