feat: support for KQL in SQLScript (#27522)

tests/unit_tests/sql_parse_tests.py

@@ -37,8 +37,10 @@ from superset.sql_parse import (
     has_table_query,
     insert_rls_as_subquery,
     insert_rls_in_predicate,
+    KustoKQLStatement,
     ParsedQuery,
     sanitize_clause,
+    split_kql,
     SQLScript,
     SQLStatement,
     strip_comments_from_sql,
@@ -1883,21 +1885,31 @@ def test_sqlquery() -> None:
     """
     Test the `SQLScript` class.
     """
-    script = SQLScript("SELECT 1; SELECT 2;")
+    script = SQLScript("SELECT 1; SELECT 2;", "sqlite")
 
     assert len(script.statements) == 2
     assert script.format() == "SELECT\n  1;\nSELECT\n  2"
     assert script.statements[0].format() == "SELECT\n  1"
 
-    script = SQLScript("SET a=1; SET a=2; SELECT 3;")
+    script = SQLScript("SET a=1; SET a=2; SELECT 3;", "sqlite")
     assert script.get_settings() == {"a": "2"}
 
+    query = SQLScript(
+        """set querytrace;
+Events | take 100""",
+        "kustokql",
+    )
+    assert query.get_settings() == {"querytrace": True}
+
 
 def test_sqlstatement() -> None:
     """
     Test the `SQLStatement` class.
     """
-    statement = SQLStatement("SELECT * FROM table1 UNION ALL SELECT * FROM table2")
+    statement = SQLStatement(
+        "SELECT * FROM table1 UNION ALL SELECT * FROM table2",
+        "sqlite",
+    )
 
     assert statement.tables == {
         Table(table="table1", schema=None, catalog=None),
@@ -1908,7 +1920,7 @@ def test_sqlstatement() -> None:
         == "SELECT\n  *\nFROM table1\nUNION ALL\nSELECT\n  *\nFROM table2"
     )
 
-    statement = SQLStatement("SET a=1")
+    statement = SQLStatement("SET a=1", "sqlite")
     assert statement.get_settings() == {"a": "1"}
 
 
@@ -1950,3 +1962,137 @@ def test_extract_tables_from_jinja_sql(
         extract_tables_from_jinja_sql(sql.format(engine=engine, macro=macro), engine)
         == expected
     )
+
+
+def test_kustokqlstatement_split_query() -> None:
+    """
+    Test the `KustoKQLStatement` split method.
+    """
+    statements = KustoKQLStatement.split_query(
+        """
+let totalPagesPerDay = PageViews
+| summarize by Page, Day = startofday(Timestamp)
+| summarize count() by Day;
+let materializedScope = PageViews
+| summarize by Page, Day = startofday(Timestamp);
+let cachedResult = materialize(materializedScope);
+cachedResult
+| project Page, Day1 = Day
+| join kind = inner
+(
+    cachedResult
+    | project Page, Day2 = Day
+)
+on Page
+| where Day2 > Day1
+| summarize count() by Day1, Day2
+| join kind = inner
+    totalPagesPerDay
+on $left.Day1 == $right.Day
+| project Day1, Day2, Percentage = count_*100.0/count_1
+        """,
+        "kustokql",
+    )
+    assert len(statements) == 4
+
+
+def test_kustokqlstatement_with_program() -> None:
+    """
+    Test the `KustoKQLStatement` split method when the KQL has a program.
+    """
+    statements = KustoKQLStatement.split_query(
+        """
+print program = ```
+  public class Program {
+    public static void Main() {
+      System.Console.WriteLine("Hello!");
+    }
+  }```
+        """,
+        "kustokql",
+    )
+    assert len(statements) == 1
+
+
+def test_kustokqlstatement_with_set() -> None:
+    """
+    Test the `KustoKQLStatement` split method when the KQL has a set command.
+    """
+    statements = KustoKQLStatement.split_query(
+        """
+set querytrace;
+Events | take 100
+        """,
+        "kustokql",
+    )
+    assert len(statements) == 2
+    assert statements[0].format() == "set querytrace"
+    assert statements[1].format() == "Events | take 100"
+
+
+@pytest.mark.parametrize(
+    "kql,statements",
+    [
+        ('print banner=strcat("Hello", ", ", "World!")', 1),
+        (r"print 'O\'Malley\'s'", 1),
+        (r"print 'O\'Mal;ley\'s'", 1),
+        ("print ```foo;\nbar;\nbaz;```\n", 1),
+    ],
+)
+def test_kustokql_statement_split_special(kql: str, statements: int) -> None:
+    assert len(KustoKQLStatement.split_query(kql, "kustokql")) == statements
+
+
+def test_split_kql() -> None:
+    """
+    Test the `split_kql` function.
+    """
+    kql = """
+let totalPagesPerDay = PageViews
+| summarize by Page, Day = startofday(Timestamp)
+| summarize count() by Day;
+let materializedScope = PageViews
+| summarize by Page, Day = startofday(Timestamp);
+let cachedResult = materialize(materializedScope);
+cachedResult
+| project Page, Day1 = Day
+| join kind = inner
+(
+    cachedResult
+    | project Page, Day2 = Day
+)
+on Page
+| where Day2 > Day1
+| summarize count() by Day1, Day2
+| join kind = inner
+    totalPagesPerDay
+on $left.Day1 == $right.Day
+| project Day1, Day2, Percentage = count_*100.0/count_1
+    """
+    assert split_kql(kql) == [
+        """
+let totalPagesPerDay = PageViews
+| summarize by Page, Day = startofday(Timestamp)
+| summarize count() by Day""",
+        """
+let materializedScope = PageViews
+| summarize by Page, Day = startofday(Timestamp)""",
+        """
+let cachedResult = materialize(materializedScope)""",
+        """
+cachedResult
+| project Page, Day1 = Day
+| join kind = inner
+(
+    cachedResult
+    | project Page, Day2 = Day
+)
+on Page
+| where Day2 > Day1
+| summarize count() by Day1, Day2
+| join kind = inner
+    totalPagesPerDay
+on $left.Day1 == $right.Day
+| project Day1, Day2, Percentage = count_*100.0/count_1
+    """,
+    ]