WIP
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index e6c52d6..67d2dcd 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py
@@ -2170,10 +2170,6 @@ return False @classmethod - def parse_sql(cls, sql: str) -> list[str]: - return [str(s).strip(" ;") for s in sqlparse.parse(sql)] - - @classmethod def get_impersonation_key(cls, user: User | None) -> Any: """ Construct an impersonation key, by default it's the given username.
diff --git a/superset/db_engine_specs/kusto.py b/superset/db_engine_specs/kusto.py index 696faf7..abceb6a 100644 --- a/superset/db_engine_specs/kusto.py +++ b/superset/db_engine_specs/kusto.py
@@ -156,11 +156,3 @@ @classmethod def is_select_query(cls, parsed_query: ParsedQuery) -> bool: return not parsed_query.sql.startswith(".") - - @classmethod - def parse_sql(cls, sql: str) -> list[str]: - """ - Kusto supports a single query statement, but it could include sub queries - and variables declared via let keyword. - """ - return [sql]
diff --git a/superset/models/core.py b/superset/models/core.py index 4181412..35aedb1 100755 --- a/superset/models/core.py +++ b/superset/models/core.py
@@ -74,6 +74,7 @@ ) from superset.models.helpers import AuditMixinNullable, ImportExportMixin from superset.result_set import SupersetResultSet +from superset.sql.parse import SQLScript from superset.sql_parse import Table from superset.superset_typing import OAuth2ClientConfig, ResultSetColumnType from superset.utils import cache as cache_util, core as utils, json @@ -674,7 +675,7 @@ schema: str | None = None, mutator: Callable[[pd.DataFrame], None] | None = None, ) -> pd.DataFrame: - sqls = self.db_engine_spec.parse_sql(sql) + parsed_script = SQLScript(sql, engine=self.db_engine_spec.engine) with self.get_sqla_engine(catalog=catalog, schema=schema) as engine: engine_url = engine.url @@ -691,8 +692,9 @@ with self.get_raw_connection(catalog=catalog, schema=schema) as conn: cursor = conn.cursor() df = None - for i, sql_ in enumerate(sqls): - sql_ = self.mutate_sql_based_on_config(sql_, is_split=True) + for i, statement in enumerate(parsed_script.statements): + # pylint: disable=protected-access + sql_ = self.mutate_sql_based_on_config(statement._sql, is_split=True) _log_query(sql_) with event_logger.log_context( action="execute_sql", @@ -700,7 +702,7 @@ object_ref=__name__, ): self.db_engine_spec.execute(cursor, sql_, self) - if i < len(sqls) - 1: + if i < len(parsed_script.statements) - 1: # If it's not the last, we don't keep the results cursor.fetchall() else:
diff --git a/superset/sql/parse.py b/superset/sql/parse.py index 2c7276b..453dbd1 100644 --- a/superset/sql/parse.py +++ b/superset/sql/parse.py
@@ -20,10 +20,11 @@ import enum import logging import re +import string import urllib.parse from collections.abc import Iterable from dataclasses import dataclass -from typing import Any, Generic, TypeVar +from typing import Any, Generic, Iterator, TypeVar import sqlglot import sqlparse @@ -226,6 +227,12 @@ """ raise NotImplementedError() + def is_select(self) -> bool: + """ + Check if the statement is a `SELECT` statement. + """ + raise NotImplementedError() + def __str__(self) -> str: return self.format() @@ -382,6 +389,12 @@ return False + def is_select(self) -> bool: + """ + Check if the statement is a `SELECT` statement. + """ + return isinstance(self._parsed, exp.Select) + def format(self, comments: bool = True) -> str: """ Pretty-format the SQL statement. @@ -431,60 +444,115 @@ } -class KQLSplitState(enum.Enum): +class KQLTokenizeState(enum.Enum): """ - State machine for splitting a KQL script. + State machine for tokenizing a KQL script. The state machine keeps track of whether we're inside a string or not, so we don't split the script in a semi-colon that's part of a string. """ - OUTSIDE_STRING = enum.auto() + OUTSIDE = enum.auto() INSIDE_SINGLE_QUOTED_STRING = enum.auto() INSIDE_DOUBLE_QUOTED_STRING = enum.auto() INSIDE_MULTILINE_STRING = enum.auto() + INSIDE_SINGLE_QUOTED_IDENTIFIER = enum.auto() + INSIDE_DOUBLE_QUOTED_IDENTIFIER = enum.auto() + + +def tokenize_kql(kql: str) -> Iterator[str]: + """ + Tokenize a KQL script. + """ + valid_identifier_chars = set(string.ascii_letters + string.digits + "_") + valid_quoted_identifier_chars = valid_identifier_chars | set(" .-") + + script = kql if kql.endswith(";") else kql + ";" + + cursor = 0 + while cursor < len(script): + rest = script[cursor:] + + # quoted identifiers + if rest[:2] in {"['", '["'}: + match = "']" if rest[:2] == "['" else '"]' + if match not in rest[2:]: + raise SupersetParseError( + script, + "kustokql", + message="Unclosed quoted identifier", + ) + + token = rest[: rest.index(match, 2) + 2] + + if any(char not in valid_quoted_identifier_chars for char in token[2:-2]): + raise SupersetParseError( + script, + "kustokql", + message="Invalid quoted identifier", + ) + + yield token + cursor += len(token) + + # multi-line strings + elif rest[:3] == "```": + if "```" not in rest[3:]: + raise SupersetParseError( + script, + "kustokql", + message="Unclosed multi-line string", + ) + + token = rest[: rest.index("```", 3) + 3] + yield token + cursor += len(token) + + # single-quoted strings + elif rest[0] in {'"', "'"}: + match = rest[0] + # find first unescaped quote + start = 1 + while True: + if match not in rest[start:]: + raise SupersetParseError( + script, + "kustokql", + message="Unclosed string", + ) + index = rest.index(match, start) + if rest[index - 1] != "\\": + break + + start = index + 1 + + token = rest[: index + 1] + yield token + cursor += len(token) + + # identifiers and keywords + else: + for i, char in enumerate(rest): + if char not in valid_identifier_chars: + if i > 0: + yield rest[:i] + yield char + cursor += i + 1 + break def split_kql(kql: str) -> list[str]: """ Custom function for splitting KQL statements. """ - statements = [] - state = KQLSplitState.OUTSIDE_STRING - statement_start = 0 - script = kql if kql.endswith(";") else kql + ";" - for i, character in enumerate(script): - if state == KQLSplitState.OUTSIDE_STRING: - if character == ";": - statements.append(script[statement_start:i]) - statement_start = i + 1 - elif character == "'": - state = KQLSplitState.INSIDE_SINGLE_QUOTED_STRING - elif character == '"': - state = KQLSplitState.INSIDE_DOUBLE_QUOTED_STRING - elif character == "`" and script[i - 2 : i] == "``": - state = KQLSplitState.INSIDE_MULTILINE_STRING - - elif ( - state == KQLSplitState.INSIDE_SINGLE_QUOTED_STRING - and character == "'" - and script[i - 1] != "\\" - ): - state = KQLSplitState.OUTSIDE_STRING - - elif ( - state == KQLSplitState.INSIDE_DOUBLE_QUOTED_STRING - and character == '"' - and script[i - 1] != "\\" - ): - state = KQLSplitState.OUTSIDE_STRING - - elif ( - state == KQLSplitState.INSIDE_MULTILINE_STRING - and character == "`" - and script[i - 2 : i] == "``" - ): - state = KQLSplitState.OUTSIDE_STRING + statements: list[str] = [] + statement: list[str] = [] + for token in tokenize_kql(kql): + if token == ";": + statements.append("".join(statement)) + statement = [] + else: + statement.append(token) return statements @@ -506,6 +574,14 @@ details about it. """ + def __init__( + self, + statement: str, + engine: str = "kustokql", + ast: str | None = None, + ): + super().__init__(statement, engine, ast) + @classmethod def split_script( cls, @@ -588,6 +664,56 @@ """ return self._parsed.startswith(".") and not self._parsed.startswith(".show") + def is_select(self) -> bool: + """ + Check if the statement is a `SELECT` statement. + """ + if not self._parsed or self.is_mutating(): + return False + + # strip comments + kql = "\n".join( + line + for line in self._parsed.split("\n") + if not line.strip().startswith("//") + ).strip() + + first_token = next(tokenize_kql(kql), None) + if not first_token: + return False + + return first_token == "|" or self._is_identifier(first_token) + + @staticmethod + def _is_identifier(identifier: str) -> bool: + """ + Validates if a given string is a valid KQL identifier. + + From the documentation: + + Identifiers are case-sensitive. Database names are case-insensitive, and + therefore an exception to this rule. + Identifiers must be between 1 and 1024 characters long. + Identifiers may contain letters, digits, and underscores (_). + Identifiers may contain certain special characters: spaces, dots (.), and + dashes (-). For information on how to reference identifiers with special + characters, see Reference identifiers in queries. + + """ + valid_chars = set(string.ascii_letters + string.digits + "_") + + # Identifiers names that (1) include special character, (2) are language + # keywords, or (3) are literals must be enclosed using [' and '] or [" and "]. + if (identifier.startswith("['") and identifier.endswith("']")) or ( + identifier.startswith('["') and identifier.endswith('"]') + ): + identifier = identifier[2:-2] + valid_chars.update(" .-") + + return 1 <= len(identifier) <= 1024 and all( + char in valid_chars for char in identifier + ) + class SQLScript: """ @@ -642,6 +768,24 @@ """ return any(statement.is_mutating() for statement in self.statements) + def is_valid_ctas(self) -> bool: + """ + Check if the script contains a valid CTAS statement. + + CTAS (`CREATE TABLE AS SELECT`) can only be run with scripts where the last + statement is a `SELECT`. + """ + return self.statements[-1].is_select() + + def is_valid_cvas(self) -> bool: + """ + Check if the script contains a valid CVAS statement. + + CVAS (`CREATE VIEW AS SELECT`) can only be run with scripts with a single + `SELECT` statement. + """ + return len(self.statements) == 1 and self.statements[0].is_select() + def extract_tables_from_statement( statement: exp.Expression, @@ -650,7 +794,7 @@ """ Extract all table references in a single statement. - Please not that this is not trivial; consider the following queries: + Please note that this is not trivial; consider the following queries: DESCRIBE some_table; SHOW PARTITIONS FROM some_table;
diff --git a/tests/unit_tests/db_engine_specs/test_base.py b/tests/unit_tests/db_engine_specs/test_base.py index d8e632c..805089c 100644 --- a/tests/unit_tests/db_engine_specs/test_base.py +++ b/tests/unit_tests/db_engine_specs/test_base.py
@@ -49,32 +49,6 @@ assert text_clause.text == "SELECT foo FROM tbl WHERE foo = '123\\:456')" -def test_parse_sql_single_statement() -> None: - """ - `parse_sql` should properly strip leading and trailing spaces and semicolons - """ - - from superset.db_engine_specs.base import BaseEngineSpec - - queries = BaseEngineSpec.parse_sql(" SELECT foo FROM tbl ; ") - assert queries == ["SELECT foo FROM tbl"] - - -def test_parse_sql_multi_statement() -> None: - """ - For string with multiple SQL-statements `parse_sql` method should return list - where each element represents the single SQL-statement - """ - - from superset.db_engine_specs.base import BaseEngineSpec - - queries = BaseEngineSpec.parse_sql("SELECT foo FROM tbl1; SELECT bar FROM tbl2;") - assert queries == [ - "SELECT foo FROM tbl1", - "SELECT bar FROM tbl2", - ] - - def test_validate_db_uri(mocker: MockerFixture) -> None: """ Ensures that the `validate_database_uri` method invokes the validator correctly
diff --git a/tests/unit_tests/db_engine_specs/test_kusto.py b/tests/unit_tests/db_engine_specs/test_kusto.py index 68330ed..88dfcff 100644 --- a/tests/unit_tests/db_engine_specs/test_kusto.py +++ b/tests/unit_tests/db_engine_specs/test_kusto.py
@@ -20,99 +20,11 @@ import pytest -from superset.sql.parse import SQLScript -from superset.sql_parse import ParsedQuery from tests.unit_tests.db_engine_specs.utils import assert_convert_dttm from tests.unit_tests.fixtures.common import dttm # noqa: F401 @pytest.mark.parametrize( - "sql,expected", - [ - ("SELECT foo FROM tbl", False), - ("SHOW TABLES", False), - ("EXPLAIN SELECT foo FROM tbl", False), - ("INSERT INTO tbl (foo) VALUES (1)", True), - ], -) -def test_sql_has_mutation(sql: str, expected: bool) -> None: - """ - Make sure that SQL dialect consider only SELECT statements as read-only - """ - - from superset.db_engine_specs.kusto import KustoSqlEngineSpec - - assert ( - SQLScript( - sql, - engine=KustoSqlEngineSpec.engine, - ).has_mutation() - == expected - ) - - -@pytest.mark.parametrize( - "kql,expected", - [ - ("tbl | limit 100", True), - ("let foo = 1; tbl | where bar == foo", True), - (".show tables", False), - ], -) -def test_kql_is_select_query(kql: str, expected: bool) -> None: - """ - Make sure that KQL dialect consider only statements that do not start with "." (dot) - as a SELECT statements - """ - - from superset.db_engine_specs.kusto import KustoKqlEngineSpec - - parsed_query = ParsedQuery(kql) - assert KustoKqlEngineSpec.is_select_query(parsed_query) == expected - - -@pytest.mark.parametrize( - "kql,expected", - [ - ("tbl | limit 100", False), - ("let foo = 1; tbl | where bar == foo", False), - (".show tables", False), - ("print 1", False), - ("set querytrace; Events | take 100", False), - (".drop table foo", True), - (".set-or-append table foo <| bar", True), - ], -) -def test_kql_has_mutation(kql: str, expected: bool) -> None: - """ - Make sure that KQL dialect consider only SELECT statements as read-only - """ - - from superset.db_engine_specs.kusto import KustoKqlEngineSpec - - assert ( - SQLScript( - kql, - engine=KustoKqlEngineSpec.engine, - ).has_mutation() - == expected - ) - - -def test_kql_parse_sql() -> None: - """ - parse_sql method should always return a list with a single element - which is an original query - """ - - from superset.db_engine_specs.kusto import KustoKqlEngineSpec - - queries = KustoKqlEngineSpec.parse_sql("let foo = 1; tbl | where bar == foo") - - assert queries == ["let foo = 1; tbl | where bar == foo"] - - -@pytest.mark.parametrize( "target_type,expected_result", [ ("DateTime", "datetime(2019-01-02T03:04:05.678900)"),
diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index 4911d4c..c08bf1d 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py
@@ -945,6 +945,10 @@ ("kustokql", "set querytrace; Events | take 100", False), ("kustokql", ".drop table foo", True), ("kustokql", ".set-or-append table foo <| bar", True), + ("kustosql", "SELECT foo FROM tbl", False), + ("kustosql", "SHOW TABLES", False), + ("kustosql", "EXPLAIN SELECT foo FROM tbl", False), + ("kustosql", "INSERT INTO tbl (foo) VALUES (1)", True), ("base", "SHOW LOCKS test EXTENDED", False), ("base", "SET hivevar:desc='Legislators'", False), ("base", "UPDATE t1 SET col1 = NULL", True), @@ -1070,3 +1074,100 @@ "with source as ( select 1 as one ) select * from source", engine=engine, ).is_mutating() + + +@pytest.mark.parametrize( + "identifier,expected", + [ + # Rule: Identifiers are case-sensitive + ("myTable", True), + ("MYTABLE", True), + ("MyTable", True), + # Rule: Identifiers must be between 1 and 1024 characters long + ("a", True), + ("a" * 1024, True), + ("a" * 1025, False), + ("", False), + # Rule: Identifiers may contain letters, digits, and underscores + ("My_Table_123", True), + ("123Table", True), + # Rule: Identifiers may contain special characters: spaces, dots, dashes (when quoted) + ("['My Table']", True), + ("['My-Table']", True), + ("['My.Table']", True), + ("['Table-']", True), + ("['My Table Name']", True), + ("['My!Table']", False), + ("['MyTable ']", True), + (" MyTable", False), + ("MyTable ", False), + # Rule: Non-special identifiers don't require quoting + ("MyTable", True), + ("My-Table", False), + # Invalid quoting + ("['Invalid]", False), + ("['Invalid'Name']", False), + ("['']", False), + # Rule: Literal identifiers or language keywords + ("['select']", True), + ("select", True), + ], +) +def test_is_kql_identifier(identifier: str, expected: bool): + """ + Tests the _is_identifier method for various valid and invalid cases. + """ + assert KustoKQLStatement._is_identifier(identifier) == expected + + +@pytest.mark.parametrize( + "kql,expected", + [ + # Simple SELECT-like statements (non-mutating queries) + ("MyTable | count", True), + ("MyTable", True), + ("| count", True), + ("tbl | limit 100", True), + (".show tables", False), + # With comments (ensure comments are stripped out) + ("// Comment only", False), + ("MyTable // trailing comment", True), + ("// leading comment\nMyTable", True), + ("MyTable\n// intermediate comment\n| count", True), + # Mutating query (should return False) + (".drop MyTable", False), + ( + ".update MyTable set Column1 = 100", + False, + ), + (".alter MyTable", False), + # Edge cases for first token + ("", False), + (" ", False), + (".command", False), + ("['My Table']", True), + # Complex multi-line queries + ( + """ + // Initial comment + MyTable + | where Column1 > 100 + """, + True, + ), + ( + """ + MyTable + | where Column1 > 100 + | summarize by Column2 + // Final comment + """, + True, + ), + ], +) +def test_kql_is_select(kql: str, expected: bool): + """ + Tests the is_select method for various valid and invalid cases. + """ + assert KustoKQLStatement(kql).is_select() == expected