diff --git a/superset/db_engine_specs/kusto.py b/superset/db_engine_specs/kusto.py index 7177cd13fb2..85b563262bd 100644 --- a/superset/db_engine_specs/kusto.py +++ b/superset/db_engine_specs/kusto.py @@ -14,11 +14,12 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import logging import re from datetime import datetime -from typing import Any, Optional +from typing import Any, Optional, TYPE_CHECKING -from sqlalchemy import types +from sqlalchemy import func, types from sqlalchemy.dialects.mssql.base import SMALLDATETIME from superset.constants import TimeGrain @@ -28,8 +29,45 @@ from superset.db_engine_specs.exceptions import ( SupersetDBAPIOperationalError, SupersetDBAPIProgrammingError, ) -from superset.sql.parse import LimitMethod -from superset.utils.core import GenericDataType + +if TYPE_CHECKING: + from superset.models.core import Database + +from superset.sql.parse import KQLTokenType, LimitMethod, tokenize_kql +from superset.utils.core import FilterOperator, GenericDataType + +logger = logging.getLogger(__name__) + +_OPENING_BRACKET = [ + (KQLTokenType.WORD, "ARRAY"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, "["), +] +_CLOSING_BRACKET = [(KQLTokenType.OTHER, "]"), (KQLTokenType.OTHER, ")")] + + +def strip_array_brackets(kql: str) -> str: + """ + Replace ``ARRAY([...])`` wrappers with ``[...]`` using the KQL tokenizer. + + SQLAlchemy sometimes wraps bracket-quoted KQL identifiers in ARRAY(), + which is invalid KQL. This strips the wrapper while preserving the contents. + """ + tokens = tokenize_kql(kql) + + to_remove: set[int] = set() + depth = 0 + for i in range(len(tokens)): + if tokens[i : i + 3] == _OPENING_BRACKET: + to_remove.add(i) + to_remove.add(i + 1) + depth += 1 + elif depth > 0 and tokens[i : i + 2] == _CLOSING_BRACKET: + to_remove.add(i + 1) + depth -= 1 + + tokens = [token for i, token in enumerate(tokens) if i not in to_remove] + return "".join(val for _, val in tokens) class KustoSqlEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method @@ -190,6 +228,14 @@ class KustoKqlEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method type_code_map: dict[int, str] = {} # loaded from get_datatype only if needed + column_type_mappings = ( + ( + re.compile(r"^array.*", re.IGNORECASE), + types.String(), + GenericDataType.STRING, + ), + ) + @classmethod def get_dbapi_exception_mapping(cls) -> dict[type[Exception], type[Exception]]: # pylint: disable=import-outside-toplevel,import-error @@ -201,6 +247,44 @@ class KustoKqlEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method kusto_exceptions.ProgrammingError: SupersetDBAPIProgrammingError, } + @classmethod + def handle_null_filter( + cls, + sqla_col: Any, + op: FilterOperator, + ) -> Any: + """ + Handle null/not null filter operations for KQL. + + In KQL, null checks use functions: + - isnull(col) for IS NULL + - isnotnull(col) for IS NOT NULL + + :param sqla_col: SQLAlchemy column element + :param op: Filter operator (IS_NULL or IS_NOT_NULL) + :return: SQLAlchemy expression for the null filter + """ + if op == FilterOperator.IS_NULL: + return func.isnull(sqla_col) + if op == FilterOperator.IS_NOT_NULL: + return func.isnotnull(sqla_col) + + raise ValueError(f"Invalid null filter operator: {op}") + + @classmethod + def epoch_to_dttm(cls) -> str: + """ + Convert from number of seconds since the epoch to a timestamp. + """ + return "unixtime_seconds_todatetime({col})" + + @classmethod + def epoch_ms_to_dttm(cls) -> str: + """ + Convert from number of milliseconds since the epoch to a timestamp. + """ + return "unixtime_milliseconds_todatetime({col})" + @classmethod def convert_dttm( cls, target_type: str, dttm: datetime, db_extra: Optional[dict[str, Any]] = None @@ -213,3 +297,22 @@ class KustoKqlEngineSpec(BaseEngineSpec): # pylint: disable=abstract-method return f"""datetime({dttm.isoformat(timespec="microseconds")})""" return None + + @classmethod + def execute( + cls, + cursor: Any, + query: str, + database: "Database", + **kwargs: Any, + ) -> None: + """ + Execute a KQL query, fixing ARRAY() wrappers around + bracket-quoted identifiers. + + Example: + ARRAY(["age"]) -> ["age"] + ARRAY(["user_name"]) -> ["user_name"] + """ + processed_query = strip_array_brackets(query) + super().execute(cursor, processed_query, database, **kwargs) diff --git a/tests/unit_tests/db_engine_specs/test_kusto.py b/tests/unit_tests/db_engine_specs/test_kusto.py index a21e82a5676..568e90153a8 100644 --- a/tests/unit_tests/db_engine_specs/test_kusto.py +++ b/tests/unit_tests/db_engine_specs/test_kusto.py @@ -139,3 +139,88 @@ def test_timegrain_expressions(in_duration: str, expected_result: str) -> None: col=col, pdf=None, time_grain=in_duration ) assert str(actual_result) == expected_result + + +def test_epoch_to_dttm() -> None: + """ + Test that KQL engine spec returns correct epoch to datetime conversion template. + """ + result = KustoKqlEngineSpec.epoch_to_dttm() + assert result == "unixtime_seconds_todatetime({col})" + + +def test_epoch_ms_to_dttm() -> None: + """ + Test that KQL engine spec returns correct epoch milliseconds to + datetime conversion template. + """ + result = KustoKqlEngineSpec.epoch_ms_to_dttm() + assert result == "unixtime_milliseconds_todatetime({col})" + + +def test_handle_null_filter() -> None: + """ + Test that KQL engine spec uses isnull/isnotnull functions for null filters. + """ + from superset.utils.core import FilterOperator + + test_col = column("test_column") + + # Test IS_NULL - should return isnull(col) + result_null = KustoKqlEngineSpec.handle_null_filter( + test_col, FilterOperator.IS_NULL + ) + assert str(result_null) == "isnull(test_column)" + + # Test IS_NOT_NULL - should return isnotnull(col) + result_not_null = KustoKqlEngineSpec.handle_null_filter( + test_col, FilterOperator.IS_NOT_NULL + ) + assert str(result_not_null) == "isnotnull(test_column)" + + # Test invalid operator - should raise ValueError + with pytest.raises(ValueError, match="Invalid null filter operator"): + KustoKqlEngineSpec.handle_null_filter(test_col, "INVALID_OPERATOR") # type: ignore[arg-type] + + +@pytest.mark.parametrize( + ("raw_query", "expected_query"), + [ + ( + 'database("superset").["FreeCodeCamp"] | extend ["age"] = ARRAY(["age"]) ' + '| project ["age"] | take 100', + 'database("superset").["FreeCodeCamp"] | extend ["age"] = ["age"] ' + '| project ["age"] | take 100', + ), + ( + 'database("superset").["FreeCodeCamp"] | project ["age"] | take 100', + 'database("superset").["FreeCodeCamp"] | project ["age"] | take 100', + ), + ( + 'database("superset").["VideoGameSales"]' + ' | where ["rank"]<= 25' + ' | summarize ["SUM(Global_Sales)"] = sum(["global_sales"])' + ' by ["publisher"]' + ' | project ["publisher"], ["SUM(Global_Sales)"]' + ' | order by ["SUM(Global_Sales)"] desc' + " | take 50000", + 'database("superset").["VideoGameSales"]' + ' | where ["rank"]<= 25' + ' | summarize ["SUM(Global_Sales)"] = sum(["global_sales"])' + ' by ["publisher"]' + ' | project ["publisher"], ["SUM(Global_Sales)"]' + ' | order by ["SUM(Global_Sales)"] desc' + " | take 50000", + ), + ], +) +def test_kql_execute_array_processing(raw_query: str, expected_query: str) -> None: + """Ensure `execute` replaces ARRAY wrappers and leaves other queries unchanged.""" + from unittest.mock import Mock + + mock_cursor = Mock() + mock_db = Mock() + + KustoKqlEngineSpec.execute(mock_cursor, raw_query, mock_db) + + mock_cursor.execute.assert_called_once_with(expected_query) diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 206ebc6ea72..bec1849a196 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -4121,6 +4121,274 @@ def test_kustokql_statement_check_tables_present() -> None: ), ("'test'", [(KQLTokenType.STRING, "'test'")]), ("```test```", [(KQLTokenType.STRING, "```test```")]), + # Double-quoted strings + ('"hello"', [(KQLTokenType.STRING, '"hello"')]), + # Single-quoted string with escaped quote + ( + "'it\\'s a test'", + [(KQLTokenType.STRING, "'it\\'s a test'")], + ), + # Double-quoted string with escaped quote + ( + '"say \\"hi\\""', + [(KQLTokenType.STRING, '"say \\"hi\\""')], + ), + # Semicolon token + ( + "a; b", + [ + (KQLTokenType.WORD, "a"), + (KQLTokenType.SEMICOLON, ";"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "b"), + ], + ), + # Semicolon inside string is not a SEMICOLON token + ( + "'a;b'", + [(KQLTokenType.STRING, "'a;b'")], + ), + # Numbers + ( + "42", + [(KQLTokenType.NUMBER, "42")], + ), + # Other/punctuation tokens + ( + "()", + [ + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, ")"), + ], + ), + # Empty input + ("", []), + # ARRAY bracket pattern used in Kusto engine spec + ( + 'ARRAY(["age"])', + [ + (KQLTokenType.WORD, "ARRAY"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, "["), + (KQLTokenType.STRING, '"age"'), + (KQLTokenType.OTHER, "]"), + (KQLTokenType.OTHER, ")"), + ], + ), + # Mixed identifiers, operators, and strings + ( + "tbl | where name == 'Alice' | take 5", + [ + (KQLTokenType.WORD, "tbl"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "where"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "name"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "="), + (KQLTokenType.OTHER, "="), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.STRING, "'Alice'"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "take"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.NUMBER, "5"), + ], + ), + # Underscore in identifiers + ( + "my_table", + [(KQLTokenType.WORD, "my_table")], + ), + # Identifiers starting with underscore + ( + "_col1", + [(KQLTokenType.WORD, "_col1")], + ), + # Multiline string with semicolons and quotes + ( + "```select 'x'; drop```", + [(KQLTokenType.STRING, "```select 'x'; drop```")], + ), + # Adjacent strings without whitespace + ( + "'a''b'", + [ + (KQLTokenType.STRING, "'a'"), + (KQLTokenType.STRING, "'b'"), + ], + ), + # Dot operator + ( + "db.table", + [ + (KQLTokenType.WORD, "db"), + (KQLTokenType.OTHER, "."), + (KQLTokenType.WORD, "table"), + ], + ), + # Bracket-quoted identifier (KQL style) + ( + '["column name"]', + [ + (KQLTokenType.OTHER, "["), + (KQLTokenType.STRING, '"column name"'), + (KQLTokenType.OTHER, "]"), + ], + ), + # Whitespace variants (tab, newline) + ( + "a\t\nb", + [ + (KQLTokenType.WORD, "a"), + (KQLTokenType.WHITESPACE, "\t\n"), + (KQLTokenType.WORD, "b"), + ], + ), + # Summarize with count aggregation + ( + "T | summarize count() by State", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "summarize"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "count"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, ")"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "by"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "State"), + ], + ), + # Aliased aggregation with avg + ( + "T | summarize avg_val = avg(price) by category", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "summarize"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "avg_val"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "="), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "avg"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.WORD, "price"), + (KQLTokenType.OTHER, ")"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "by"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "category"), + ], + ), + # Multiple aggregations with dcount + ( + "T | summarize cnt = count(), uniq = dcount(user_id)", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "summarize"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "cnt"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "="), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "count"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, ")"), + (KQLTokenType.OTHER, ","), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "uniq"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "="), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "dcount"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.WORD, "user_id"), + (KQLTokenType.OTHER, ")"), + ], + ), + # Summarize with bin time bucketing + ( + "T | summarize count() by bin(ts, 1h)", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "summarize"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "count"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.OTHER, ")"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "by"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "bin"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.WORD, "ts"), + (KQLTokenType.OTHER, ","), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.NUMBER, "1"), + (KQLTokenType.WORD, "h"), + (KQLTokenType.OTHER, ")"), + ], + ), + ( + "T | summarize dcountif(user_id, status == 'active') by region", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "summarize"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "dcountif"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.WORD, "user_id"), + (KQLTokenType.OTHER, ","), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "status"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "="), + (KQLTokenType.OTHER, "="), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.STRING, "'active'"), + (KQLTokenType.OTHER, ")"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "by"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "region"), + ], + ), + ( + "T | project tostring(value)", + [ + (KQLTokenType.WORD, "T"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.OTHER, "|"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "project"), + (KQLTokenType.WHITESPACE, " "), + (KQLTokenType.WORD, "tostring"), + (KQLTokenType.OTHER, "("), + (KQLTokenType.WORD, "value"), + (KQLTokenType.OTHER, ")"), + ], + ), ], ) def test_tokenize_kql(kql: str, expected: list[tuple[KQLTokenType, str]]) -> None: