You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@superset.apache.org by be...@apache.org on 2024/09/10 20:03:32 UTC
(superset) branch organize-sql-files created (now cb4e5aa30e)
This is an automated email from the ASF dual-hosted git repository.
beto pushed a change to branch organize-sql-files
in repository https://gitbox.apache.org/repos/asf/superset.git
at cb4e5aa30e chore: organize SQL parsing files
This branch includes the following new commits:
new cb4e5aa30e chore: organize SQL parsing files
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/01: chore: organize SQL parsing files
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch organize-sql-files
in repository https://gitbox.apache.org/repos/asf/superset.git
commit cb4e5aa30e3df0d4647f57ab3abf34630b800d3b
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Tue Sep 10 16:03:20 2024 -0400
chore: organize SQL parsing files
---
superset/db_engine_specs/base.py | 3 +-
superset/db_engine_specs/postgres.py | 2 +-
superset/models/helpers.py | 3 +-
superset/sql/__init__.py | 16 +
superset/sql/parse.py | 630 +++++++++++++++++++++++++++++++++++
superset/sql_lab.py | 3 +-
superset/sql_parse.py | 623 +---------------------------------
superset/sqllab/api.py | 2 +-
superset/tasks/thumbnails.py | 1 -
tests/unit_tests/sql_parse_tests.py | 12 +-
10 files changed, 664 insertions(+), 631 deletions(-)
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 2b32d156df..514fc35287 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -63,7 +63,8 @@ from superset.constants import TimeGrain as TimeGrainConstants
from superset.databases.utils import get_table_metadata, make_url_safe
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import DisallowedSQLFunction, OAuth2Error, OAuth2RedirectError
-from superset.sql_parse import ParsedQuery, SQLScript, Table
+from superset.sql.parse import SQLScript, Table
+from superset.sql_parse import ParsedQuery
from superset.superset_typing import (
OAuth2ClientConfig,
OAuth2State,
diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index 8525ea05da..0a638f65fe 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -35,7 +35,7 @@ from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import SupersetException, SupersetSecurityException
from superset.models.sql_lab import Query
-from superset.sql_parse import SQLScript
+from superset.sql.parse import SQLScript
from superset.utils import core as utils, json
from superset.utils.core import GenericDataType
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index 295ecea70e..80e66f5027 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -68,13 +68,12 @@ from superset.exceptions import (
)
from superset.extensions import feature_flag_manager
from superset.jinja_context import BaseTemplateProcessor
+from superset.sql.parse import SQLScript, SQLStatement
from superset.sql_parse import (
has_table_query,
insert_rls_in_predicate,
ParsedQuery,
sanitize_clause,
- SQLScript,
- SQLStatement,
)
from superset.superset_typing import (
AdhocMetric,
diff --git a/superset/sql/__init__.py b/superset/sql/__init__.py
new file mode 100644
index 0000000000..13a83393a9
--- /dev/null
+++ b/superset/sql/__init__.py
@@ -0,0 +1,16 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
new file mode 100644
index 0000000000..4b3d32a0a0
--- /dev/null
+++ b/superset/sql/parse.py
@@ -0,0 +1,630 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from __future__ import annotations
+
+import enum
+import logging
+import re
+import urllib.parse
+from collections.abc import Iterable
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar
+
+import sqlglot
+from sqlglot import exp, parse_one
+from sqlglot.dialects.dialect import Dialect, Dialects
+from sqlglot.errors import ParseError
+from sqlglot.optimizer.scope import Scope, ScopeType, traverse_scope
+
+from superset.exceptions import SupersetParseError
+
+logger = logging.getLogger(__name__)
+
+
+# mapping between DB engine specs and sqlglot dialects
+SQLGLOT_DIALECTS = {
+ "ascend": Dialects.HIVE,
+ "awsathena": Dialects.PRESTO,
+ "bigquery": Dialects.BIGQUERY,
+ "clickhouse": Dialects.CLICKHOUSE,
+ "clickhousedb": Dialects.CLICKHOUSE,
+ "cockroachdb": Dialects.POSTGRES,
+ "couchbase": Dialects.MYSQL,
+ # "crate": ???
+ # "databend": ???
+ "databricks": Dialects.DATABRICKS,
+ # "db2": ???
+ # "dremio": ???
+ "drill": Dialects.DRILL,
+ # "druid": ???
+ "duckdb": Dialects.DUCKDB,
+ # "dynamodb": ???
+ # "elasticsearch": ???
+ # "exa": ???
+ # "firebird": ???
+ # "firebolt": ???
+ "gsheets": Dialects.SQLITE,
+ "hana": Dialects.POSTGRES,
+ "hive": Dialects.HIVE,
+ # "ibmi": ???
+ # "impala": ???
+ # "kustokql": ???
+ # "kylin": ???
+ "mssql": Dialects.TSQL,
+ "mysql": Dialects.MYSQL,
+ "netezza": Dialects.POSTGRES,
+ # "ocient": ???
+ # "odelasticsearch": ???
+ "oracle": Dialects.ORACLE,
+ # "pinot": ???
+ "postgresql": Dialects.POSTGRES,
+ "presto": Dialects.PRESTO,
+ "pydoris": Dialects.DORIS,
+ "redshift": Dialects.REDSHIFT,
+ # "risingwave": ???
+ # "rockset": ???
+ "shillelagh": Dialects.SQLITE,
+ "snowflake": Dialects.SNOWFLAKE,
+ # "solr": ???
+ "spark": Dialects.SPARK,
+ "sqlite": Dialects.SQLITE,
+ "starrocks": Dialects.STARROCKS,
+ "superset": Dialects.SQLITE,
+ "teradatasql": Dialects.TERADATA,
+ "trino": Dialects.TRINO,
+ "vertica": Dialects.POSTGRES,
+}
+
+
+@dataclass(eq=True, frozen=True)
+class Table:
+ """
+ A fully qualified SQL table conforming to [[catalog.]schema.]table.
+ """
+
+ table: str
+ schema: str | None = None
+ catalog: str | None = None
+
+ def __str__(self) -> str:
+ """
+ Return the fully qualified SQL table name.
+
+ Should not be used for SQL generation, only for logging and debugging, since the
+ quoting is not engine-specific.
+ """
+ return ".".join(
+ urllib.parse.quote(part, safe="").replace(".", "%2E")
+ for part in [self.catalog, self.schema, self.table]
+ if part
+ )
+
+ def __eq__(self, other: Any) -> bool:
+ return str(self) == str(other)
+
+
+# To avoid unnecessary parsing/formatting of queries, the statement has the concept of
+# an "internal representation", which is the AST of the SQL statement. For most of the
+# engines supported by Superset this is `sqlglot.exp.Expression`, but there is a special
+# case: KustoKQL uses a different syntax and there are no Python parsers for it, so we
+# store the AST as a string (the original query), and manipulate it with regular
+# expressions.
+InternalRepresentation = TypeVar("InternalRepresentation")
+
+# The base type. This helps type checking the `split_query` method correctly, since each
+# derived class has a more specific return type (the class itself). This will no longer
+# be needed once Python 3.11 is the lowest version supported. See PEP 673 for more
+# information: https://peps.python.org/pep-0673/
+TBaseSQLStatement = TypeVar("TBaseSQLStatement") # pylint: disable=invalid-name
+
+
+class BaseSQLStatement(Generic[InternalRepresentation]):
+ """
+ Base class for SQL statements.
+
+ The class can be instantiated with a string representation of the query or, for
+ efficiency reasons, with a pre-parsed AST. This is useful with `sqlglot.parse`,
+ which will split a query in multiple already parsed statements.
+
+ The `engine` parameters comes from the `engine` attribute in a Superset DB engine
+ spec.
+ """
+
+ def __init__(
+ self,
+ statement: str | InternalRepresentation,
+ engine: str,
+ ):
+ self._parsed: InternalRepresentation = (
+ self._parse_statement(statement, engine)
+ if isinstance(statement, str)
+ else statement
+ )
+ self.engine = engine
+ self.tables = self._extract_tables_from_statement(self._parsed, self.engine)
+
+ @classmethod
+ def split_query(
+ cls: type[TBaseSQLStatement],
+ query: str,
+ engine: str,
+ ) -> list[TBaseSQLStatement]:
+ """
+ Split a query into multiple instantiated statements.
+
+ This is a helper function to split a full SQL query into multiple
+ `BaseSQLStatement` instances. It's used by `SQLScript` when instantiating the
+ statements within a query.
+ """
+ raise NotImplementedError()
+
+ @classmethod
+ def _parse_statement(
+ cls,
+ statement: str,
+ engine: str,
+ ) -> InternalRepresentation:
+ """
+ Parse a string containing a single SQL statement, and returns the parsed AST.
+
+ Derived classes should not assume that `statement` contains a single statement,
+ and MUST explicitly validate that. Since this validation is parser dependent the
+ responsibility is left to the children classes.
+ """
+ raise NotImplementedError()
+
+ @classmethod
+ def _extract_tables_from_statement(
+ cls,
+ parsed: InternalRepresentation,
+ engine: str,
+ ) -> set[Table]:
+ """
+ Extract all table references in a given statement.
+ """
+ raise NotImplementedError()
+
+ def format(self, comments: bool = True) -> str:
+ """
+ Format the statement, optionally ommitting comments.
+ """
+ raise NotImplementedError()
+
+ def get_settings(self) -> dict[str, str | bool]:
+ """
+ Return any settings set by the statement.
+
+ For example, for this statement:
+
+ sql> SET foo = 'bar';
+
+ The method should return `{"foo": "'bar'"}`. Note the single quotes.
+ """
+ raise NotImplementedError()
+
+ def is_mutating(self) -> bool:
+ """
+ Check if the statement mutates data (DDL/DML).
+
+ :return: True if the statement mutates data.
+ """
+ raise NotImplementedError()
+
+ def __str__(self) -> str:
+ return self.format()
+
+
+class SQLStatement(BaseSQLStatement[exp.Expression]):
+ """
+ A SQL statement.
+
+ This class is used for all engines with dialects that can be parsed using sqlglot.
+ """
+
+ def __init__(
+ self,
+ statement: str | exp.Expression,
+ engine: str,
+ ):
+ self._dialect = SQLGLOT_DIALECTS.get(engine)
+ super().__init__(statement, engine)
+
+ @classmethod
+ def split_query(
+ cls,
+ query: str,
+ engine: str,
+ ) -> list[SQLStatement]:
+ dialect = SQLGLOT_DIALECTS.get(engine)
+
+ try:
+ statements = sqlglot.parse(query, dialect=dialect)
+ except sqlglot.errors.ParseError as ex:
+ raise SupersetParseError("Unable to split query") from ex
+
+ return [cls(statement, engine) for statement in statements if statement]
+
+ @classmethod
+ def _parse_statement(
+ cls,
+ statement: str,
+ engine: str,
+ ) -> exp.Expression:
+ """
+ Parse a single SQL statement.
+ """
+ dialect = SQLGLOT_DIALECTS.get(engine)
+
+ # We could parse with `sqlglot.parse_one` to get a single statement, but we need
+ # to verify that the string contains exactly one statement.
+ try:
+ statements = sqlglot.parse(statement, dialect=dialect)
+ except sqlglot.errors.ParseError as ex:
+ raise SupersetParseError("Unable to split query") from ex
+
+ statements = [statement for statement in statements if statement]
+ if len(statements) != 1:
+ raise SupersetParseError("SQLStatement should have exactly one statement")
+
+ return statements[0]
+
+ @classmethod
+ def _extract_tables_from_statement(
+ cls,
+ parsed: exp.Expression,
+ engine: str,
+ ) -> set[Table]:
+ """
+ Find all referenced tables.
+ """
+ dialect = SQLGLOT_DIALECTS.get(engine)
+ return extract_tables_from_statement(parsed, dialect)
+
+ def is_mutating(self) -> bool:
+ """
+ Check if the statement mutates data (DDL/DML).
+
+ :return: True if the statement mutates data.
+ """
+ for node in self._parsed.walk():
+ if isinstance(
+ node,
+ (
+ exp.Insert,
+ exp.Update,
+ exp.Delete,
+ exp.Merge,
+ exp.Create,
+ exp.Drop,
+ exp.TruncateTable,
+ ),
+ ):
+ return True
+
+ if isinstance(node, exp.Command) and node.name == "ALTER":
+ return True
+
+ # Postgres runs DMLs prefixed by `EXPLAIN ANALYZE`, see
+ # https://www.postgresql.org/docs/current/sql-explain.html
+ if (
+ self._dialect == Dialects.POSTGRES
+ and isinstance(self._parsed, exp.Command)
+ and self._parsed.name == "EXPLAIN"
+ and self._parsed.expression.name.upper().startswith("ANALYZE ")
+ ):
+ analyzed_sql = self._parsed.expression.name[len("ANALYZE ") :]
+ return SQLStatement(analyzed_sql, self.engine).is_mutating()
+
+ return False
+
+ def format(self, comments: bool = True) -> str:
+ """
+ Pretty-format the SQL statement.
+ """
+ write = Dialect.get_or_raise(self._dialect)
+ return write.generate(self._parsed, copy=False, comments=comments, pretty=True)
+
+ def get_settings(self) -> dict[str, str | bool]:
+ """
+ Return the settings for the SQL statement.
+
+ >>> statement = SQLStatement("SET foo = 'bar'")
+ >>> statement.get_settings()
+ {"foo": "'bar'"}
+
+ """
+ return {
+ eq.this.sql(): eq.expression.sql()
+ for set_item in self._parsed.find_all(exp.SetItem)
+ for eq in set_item.find_all(exp.EQ)
+ }
+
+
+class KQLSplitState(enum.Enum):
+ """
+ State machine for splitting a KQL query.
+
+ The state machine keeps track of whether we're inside a string or not, so we
+ don't split the query in a semi-colon that's part of a string.
+ """
+
+ OUTSIDE_STRING = enum.auto()
+ INSIDE_SINGLE_QUOTED_STRING = enum.auto()
+ INSIDE_DOUBLE_QUOTED_STRING = enum.auto()
+ INSIDE_MULTILINE_STRING = enum.auto()
+
+
+def split_kql(kql: str) -> list[str]:
+ """
+ Custom function for splitting KQL statements.
+ """
+ statements = []
+ state = KQLSplitState.OUTSIDE_STRING
+ statement_start = 0
+ query = kql if kql.endswith(";") else kql + ";"
+ for i, character in enumerate(query):
+ if state == KQLSplitState.OUTSIDE_STRING:
+ if character == ";":
+ statements.append(query[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 query[i - 2 : i] == "``":
+ state = KQLSplitState.INSIDE_MULTILINE_STRING
+
+ elif (
+ state == KQLSplitState.INSIDE_SINGLE_QUOTED_STRING
+ and character == "'"
+ and query[i - 1] != "\\"
+ ):
+ state = KQLSplitState.OUTSIDE_STRING
+
+ elif (
+ state == KQLSplitState.INSIDE_DOUBLE_QUOTED_STRING
+ and character == '"'
+ and query[i - 1] != "\\"
+ ):
+ state = KQLSplitState.OUTSIDE_STRING
+
+ elif (
+ state == KQLSplitState.INSIDE_MULTILINE_STRING
+ and character == "`"
+ and query[i - 2 : i] == "``"
+ ):
+ state = KQLSplitState.OUTSIDE_STRING
+
+ return statements
+
+
+class KustoKQLStatement(BaseSQLStatement[str]):
+ """
+ Special class for Kusto KQL.
+
+ Kusto KQL is a SQL-like language, but it's not supported by sqlglot. Queries look
+ like this:
+
+ StormEvents
+ | summarize PropertyDamage = sum(DamageProperty) by State
+ | join kind=innerunique PopulationData on State
+ | project State, PropertyDamagePerCapita = PropertyDamage / Population
+ | sort by PropertyDamagePerCapita
+
+ See https://learn.microsoft.com/en-us/azure/data-explorer/kusto/query/ for more
+ details about it.
+ """
+
+ @classmethod
+ def split_query(
+ cls,
+ query: str,
+ engine: str,
+ ) -> list[KustoKQLStatement]:
+ """
+ Split a query at semi-colons.
+
+ Since we don't have a parser, we use a simple state machine based function. See
+ https://learn.microsoft.com/en-us/azure/data-explorer/kusto/query/scalar-data-types/string
+ for more information.
+ """
+ return [cls(statement, engine) for statement in split_kql(query)]
+
+ @classmethod
+ def _parse_statement(
+ cls,
+ statement: str,
+ engine: str,
+ ) -> str:
+ if engine != "kustokql":
+ raise SupersetParseError(f"Invalid engine: {engine}")
+
+ statements = split_kql(statement)
+ if len(statements) != 1:
+ raise SupersetParseError("SQLStatement should have exactly one statement")
+
+ return statements[0].strip()
+
+ @classmethod
+ def _extract_tables_from_statement(cls, parsed: str, engine: str) -> set[Table]:
+ """
+ Extract all tables referenced in the statement.
+
+ StormEvents
+ | where InjuriesDirect + InjuriesIndirect > 50
+ | join (PopulationData) on State
+ | project State, Population, TotalInjuries = InjuriesDirect + InjuriesIndirect
+
+ """
+ logger.warning(
+ "Kusto KQL doesn't support table extraction. This means that data access "
+ "roles will not be enforced by Superset in the database."
+ )
+ return set()
+
+ def format(self, comments: bool = True) -> str:
+ """
+ Pretty-format the SQL statement.
+ """
+ return self._parsed
+
+ def get_settings(self) -> dict[str, str | bool]:
+ """
+ Return the settings for the SQL statement.
+
+ >>> statement = KustoKQLStatement("set querytrace;")
+ >>> statement.get_settings()
+ {"querytrace": True}
+
+ """
+ set_regex = r"^set\s+(?P<name>\w+)(?:\s*=\s*(?P<value>\w+))?$"
+ if match := re.match(set_regex, self._parsed, re.IGNORECASE):
+ return {match.group("name"): match.group("value") or True}
+
+ return {}
+
+ def is_mutating(self) -> bool:
+ """
+ Check if the statement mutates data (DDL/DML).
+
+ :return: True if the statement mutates data.
+ """
+ return self._parsed.startswith(".") and not self._parsed.startswith(".show")
+
+
+class SQLScript:
+ """
+ A SQL script, with 0+ statements.
+ """
+
+ # Special engines that can't be parsed using sqlglot. Supporting non-SQL engines
+ # adds a lot of complexity to Superset, so we should avoid adding new engines to
+ # this data structure.
+ special_engines = {
+ "kustokql": KustoKQLStatement,
+ }
+
+ def __init__(
+ self,
+ query: str,
+ engine: str,
+ ):
+ statement_class = self.special_engines.get(engine, SQLStatement)
+ self.statements = statement_class.split_query(query, engine)
+
+ def format(self, comments: bool = True) -> str:
+ """
+ Pretty-format the SQL query.
+ """
+ return ";\n".join(statement.format(comments) for statement in self.statements)
+
+ def get_settings(self) -> dict[str, str | bool]:
+ """
+ Return the settings for the SQL query.
+
+ >>> statement = SQLScript("SET foo = 'bar'; SET foo = 'baz'")
+ >>> statement.get_settings()
+ {"foo": "'baz'"}
+
+ """
+ settings: dict[str, str | bool] = {}
+ for statement in self.statements:
+ settings.update(statement.get_settings())
+
+ return settings
+
+ def has_mutation(self) -> bool:
+ """
+ Check if the script contains mutating statements.
+
+ :return: True if the script contains mutating statements
+ """
+ return any(statement.is_mutating() for statement in self.statements)
+
+
+def extract_tables_from_statement(
+ statement: exp.Expression,
+ dialect: Dialects | None,
+) -> set[Table]:
+ """
+ Extract all table references in a single statement.
+
+ Please not that this is not trivial; consider the following queries:
+
+ DESCRIBE some_table;
+ SHOW PARTITIONS FROM some_table;
+ WITH masked_name AS (SELECT * FROM some_table) SELECT * FROM masked_name;
+
+ See the unit tests for other tricky cases.
+ """
+ sources: Iterable[exp.Table]
+
+ if isinstance(statement, exp.Describe):
+ # A `DESCRIBE` query has no sources in sqlglot, so we need to explicitly
+ # query for all tables.
+ sources = statement.find_all(exp.Table)
+ elif isinstance(statement, exp.Command):
+ # Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a
+ # `SELECT` statetement in order to extract tables.
+ literal = statement.find(exp.Literal)
+ if not literal:
+ return set()
+
+ try:
+ pseudo_query = parse_one(f"SELECT {literal.this}", dialect=dialect)
+ except ParseError:
+ return set()
+ sources = pseudo_query.find_all(exp.Table)
+ else:
+ sources = [
+ source
+ for scope in traverse_scope(statement)
+ for source in scope.sources.values()
+ if isinstance(source, exp.Table) and not is_cte(source, scope)
+ ]
+
+ return {
+ Table(
+ source.name,
+ source.db if source.db != "" else None,
+ source.catalog if source.catalog != "" else None,
+ )
+ for source in sources
+ }
+
+
+def is_cte(source: exp.Table, scope: Scope) -> bool:
+ """
+ Is the source a CTE?
+
+ CTEs in the parent scope look like tables (and are represented by
+ exp.Table objects), but should not be considered as such;
+ otherwise a user with access to table `foo` could access any table
+ with a query like this:
+
+ WITH foo AS (SELECT * FROM target_table) SELECT * FROM foo
+
+ """
+ parent_sources = scope.parent.sources if scope.parent else {}
+ ctes_in_scope = {
+ name
+ for name, parent_scope in parent_sources.items()
+ if isinstance(parent_scope, Scope) and parent_scope.scope_type == ScopeType.CTE
+ }
+
+ return source.name in ctes_in_scope
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 3c35899706..3d3b2898fa 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -51,13 +51,12 @@ from superset.extensions import celery_app, event_logger
from superset.models.core import Database
from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
+from superset.sql.parse import SQLStatement, Table
from superset.sql_parse import (
CtasMethod,
insert_rls_as_subquery,
insert_rls_in_predicate,
ParsedQuery,
- SQLStatement,
- Table,
)
from superset.sqllab.limiting_factor import LimitingFactor
from superset.sqllab.utils import write_ipc_buffer
diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index 59d6d643eb..f4dda2b5dc 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -19,23 +19,18 @@
from __future__ import annotations
-import enum
import logging
import re
-import urllib.parse
-from collections.abc import Iterable, Iterator
-from dataclasses import dataclass
-from typing import Any, cast, Generic, TYPE_CHECKING, TypeVar
+from collections.abc import Iterator
+from typing import Any, cast, TYPE_CHECKING
-import sqlglot
import sqlparse
from flask_babel import gettext as __
from jinja2 import nodes
from sqlalchemy import and_
-from sqlglot import exp, parse, parse_one
-from sqlglot.dialects.dialect import Dialect, Dialects
+from sqlglot import parse
+from sqlglot.dialects.dialect import Dialects
from sqlglot.errors import ParseError, SqlglotError
-from sqlglot.optimizer.scope import Scope, ScopeType, traverse_scope
from sqlparse import keywords
from sqlparse.lexer import Lexer
from sqlparse.sql import (
@@ -65,9 +60,9 @@ from sqlparse.utils import imt
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
QueryClauseValidationException,
- SupersetParseError,
SupersetSecurityException,
)
+from superset.sql.parse import extract_tables_from_statement, Table
from superset.utils.backports import StrEnum
try:
@@ -255,543 +250,6 @@ def strip_comments_from_sql(statement: str, engine: str | None = None) -> str:
)
-@dataclass(eq=True, frozen=True)
-class Table:
- """
- A fully qualified SQL table conforming to [[catalog.]schema.]table.
- """
-
- table: str
- schema: str | None = None
- catalog: str | None = None
-
- def __str__(self) -> str:
- """
- Return the fully qualified SQL table name.
- """
-
- return ".".join(
- urllib.parse.quote(part, safe="").replace(".", "%2E")
- for part in [self.catalog, self.schema, self.table]
- if part
- )
-
- def __eq__(self, __o: object) -> bool:
- return str(self) == str(__o)
-
-
-def extract_tables_from_statement(
- statement: exp.Expression,
- dialect: Dialects | None,
-) -> set[Table]:
- """
- Extract all table references in a single statement.
-
- Please not that this is not trivial; consider the following queries:
-
- DESCRIBE some_table;
- SHOW PARTITIONS FROM some_table;
- WITH masked_name AS (SELECT * FROM some_table) SELECT * FROM masked_name;
-
- See the unit tests for other tricky cases.
- """
- sources: Iterable[exp.Table]
-
- if isinstance(statement, exp.Describe):
- # A `DESCRIBE` query has no sources in sqlglot, so we need to explicitly
- # query for all tables.
- sources = statement.find_all(exp.Table)
- elif isinstance(statement, exp.Command):
- # Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a
- # `SELECT` statetement in order to extract tables.
- literal = statement.find(exp.Literal)
- if not literal:
- return set()
-
- try:
- pseudo_query = parse_one(f"SELECT {literal.this}", dialect=dialect)
- except ParseError:
- return set()
- sources = pseudo_query.find_all(exp.Table)
- else:
- sources = [
- source
- for scope in traverse_scope(statement)
- for source in scope.sources.values()
- if isinstance(source, exp.Table) and not is_cte(source, scope)
- ]
-
- return {
- Table(
- source.name,
- source.db if source.db != "" else None,
- source.catalog if source.catalog != "" else None,
- )
- for source in sources
- }
-
-
-def is_cte(source: exp.Table, scope: Scope) -> bool:
- """
- Is the source a CTE?
-
- CTEs in the parent scope look like tables (and are represented by
- exp.Table objects), but should not be considered as such;
- otherwise a user with access to table `foo` could access any table
- with a query like this:
-
- WITH foo AS (SELECT * FROM target_table) SELECT * FROM foo
-
- """
- parent_sources = scope.parent.sources if scope.parent else {}
- ctes_in_scope = {
- name
- for name, parent_scope in parent_sources.items()
- if isinstance(parent_scope, Scope) and parent_scope.scope_type == ScopeType.CTE
- }
-
- return source.name in ctes_in_scope
-
-
-# To avoid unnecessary parsing/formatting of queries, the statement has the concept of
-# an "internal representation", which is the AST of the SQL statement. For most of the
-# engines supported by Superset this is `sqlglot.exp.Expression`, but there is a special
-# case: KustoKQL uses a different syntax and there are no Python parsers for it, so we
-# store the AST as a string (the original query), and manipulate it with regular
-# expressions.
-InternalRepresentation = TypeVar("InternalRepresentation")
-
-# The base type. This helps type checking the `split_query` method correctly, since each
-# derived class has a more specific return type (the class itself). This will no longer
-# be needed once Python 3.11 is the lowest version supported. See PEP 673 for more
-# information: https://peps.python.org/pep-0673/
-TBaseSQLStatement = TypeVar("TBaseSQLStatement") # pylint: disable=invalid-name
-
-
-class BaseSQLStatement(Generic[InternalRepresentation]):
- """
- Base class for SQL statements.
-
- The class can be instantiated with a string representation of the query or, for
- efficiency reasons, with a pre-parsed AST. This is useful with `sqlglot.parse`,
- which will split a query in multiple already parsed statements.
-
- The `engine` parameters comes from the `engine` attribute in a Superset DB engine
- spec.
- """
-
- def __init__(
- self,
- statement: str | InternalRepresentation,
- engine: str,
- ):
- self._parsed: InternalRepresentation = (
- self._parse_statement(statement, engine)
- if isinstance(statement, str)
- else statement
- )
- self.engine = engine
- self.tables = self._extract_tables_from_statement(self._parsed, self.engine)
-
- @classmethod
- def split_query(
- cls: type[TBaseSQLStatement],
- query: str,
- engine: str,
- ) -> list[TBaseSQLStatement]:
- """
- Split a query into multiple instantiated statements.
-
- This is a helper function to split a full SQL query into multiple
- `BaseSQLStatement` instances. It's used by `SQLScript` when instantiating the
- statements within a query.
- """
- raise NotImplementedError()
-
- @classmethod
- def _parse_statement(
- cls,
- statement: str,
- engine: str,
- ) -> InternalRepresentation:
- """
- Parse a string containing a single SQL statement, and returns the parsed AST.
-
- Derived classes should not assume that `statement` contains a single statement,
- and MUST explicitly validate that. Since this validation is parser dependent the
- responsibility is left to the children classes.
- """
- raise NotImplementedError()
-
- @classmethod
- def _extract_tables_from_statement(
- cls,
- parsed: InternalRepresentation,
- engine: str,
- ) -> set[Table]:
- """
- Extract all table references in a given statement.
- """
- raise NotImplementedError()
-
- def format(self, comments: bool = True) -> str:
- """
- Format the statement, optionally ommitting comments.
- """
- raise NotImplementedError()
-
- def get_settings(self) -> dict[str, str | bool]:
- """
- Return any settings set by the statement.
-
- For example, for this statement:
-
- sql> SET foo = 'bar';
-
- The method should return `{"foo": "'bar'"}`. Note the single quotes.
- """
- raise NotImplementedError()
-
- def is_mutating(self) -> bool:
- """
- Check if the statement mutates data (DDL/DML).
-
- :return: True if the statement mutates data.
- """
- raise NotImplementedError()
-
- def __str__(self) -> str:
- return self.format()
-
-
-class SQLStatement(BaseSQLStatement[exp.Expression]):
- """
- A SQL statement.
-
- This class is used for all engines with dialects that can be parsed using sqlglot.
- """
-
- def __init__(
- self,
- statement: str | exp.Expression,
- engine: str,
- ):
- self._dialect = SQLGLOT_DIALECTS.get(engine)
- super().__init__(statement, engine)
-
- @classmethod
- def split_query(
- cls,
- query: str,
- engine: str,
- ) -> list[SQLStatement]:
- dialect = SQLGLOT_DIALECTS.get(engine)
-
- try:
- statements = sqlglot.parse(query, dialect=dialect)
- except sqlglot.errors.ParseError as ex:
- raise SupersetParseError("Unable to split query") from ex
-
- return [cls(statement, engine) for statement in statements if statement]
-
- @classmethod
- def _parse_statement(
- cls,
- statement: str,
- engine: str,
- ) -> exp.Expression:
- """
- Parse a single SQL statement.
- """
- dialect = SQLGLOT_DIALECTS.get(engine)
-
- # We could parse with `sqlglot.parse_one` to get a single statement, but we need
- # to verify that the string contains exactly one statement.
- try:
- statements = sqlglot.parse(statement, dialect=dialect)
- except sqlglot.errors.ParseError as ex:
- raise SupersetParseError("Unable to split query") from ex
-
- statements = [statement for statement in statements if statement]
- if len(statements) != 1:
- raise SupersetParseError("SQLStatement should have exactly one statement")
-
- return statements[0]
-
- @classmethod
- def _extract_tables_from_statement(
- cls,
- parsed: exp.Expression,
- engine: str,
- ) -> set[Table]:
- """
- Find all referenced tables.
- """
- dialect = SQLGLOT_DIALECTS.get(engine)
- return extract_tables_from_statement(parsed, dialect)
-
- def is_mutating(self) -> bool:
- """
- Check if the statement mutates data (DDL/DML).
-
- :return: True if the statement mutates data.
- """
- for node in self._parsed.walk():
- if isinstance(
- node,
- (
- exp.Insert,
- exp.Update,
- exp.Delete,
- exp.Merge,
- exp.Create,
- exp.Drop,
- exp.TruncateTable,
- ),
- ):
- return True
-
- if isinstance(node, exp.Command) and node.name == "ALTER":
- return True
-
- # Postgres runs DMLs prefixed by `EXPLAIN ANALYZE`, see
- # https://www.postgresql.org/docs/current/sql-explain.html
- if (
- self._dialect == Dialects.POSTGRES
- and isinstance(self._parsed, exp.Command)
- and self._parsed.name == "EXPLAIN"
- and self._parsed.expression.name.upper().startswith("ANALYZE ")
- ):
- analyzed_sql = self._parsed.expression.name[len("ANALYZE ") :]
- return SQLStatement(analyzed_sql, self.engine).is_mutating()
-
- return False
-
- def format(self, comments: bool = True) -> str:
- """
- Pretty-format the SQL statement.
- """
- write = Dialect.get_or_raise(self._dialect)
- return write.generate(self._parsed, copy=False, comments=comments, pretty=True)
-
- def get_settings(self) -> dict[str, str | bool]:
- """
- Return the settings for the SQL statement.
-
- >>> statement = SQLStatement("SET foo = 'bar'")
- >>> statement.get_settings()
- {"foo": "'bar'"}
-
- """
- return {
- eq.this.sql(): eq.expression.sql()
- for set_item in self._parsed.find_all(exp.SetItem)
- for eq in set_item.find_all(exp.EQ)
- }
-
-
-class KQLSplitState(enum.Enum):
- """
- State machine for splitting a KQL query.
-
- The state machine keeps track of whether we're inside a string or not, so we
- don't split the query in a semi-colon that's part of a string.
- """
-
- OUTSIDE_STRING = enum.auto()
- INSIDE_SINGLE_QUOTED_STRING = enum.auto()
- INSIDE_DOUBLE_QUOTED_STRING = enum.auto()
- INSIDE_MULTILINE_STRING = enum.auto()
-
-
-def split_kql(kql: str) -> list[str]:
- """
- Custom function for splitting KQL statements.
- """
- statements = []
- state = KQLSplitState.OUTSIDE_STRING
- statement_start = 0
- query = kql if kql.endswith(";") else kql + ";"
- for i, character in enumerate(query):
- if state == KQLSplitState.OUTSIDE_STRING:
- if character == ";":
- statements.append(query[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 query[i - 2 : i] == "``":
- state = KQLSplitState.INSIDE_MULTILINE_STRING
-
- elif (
- state == KQLSplitState.INSIDE_SINGLE_QUOTED_STRING
- and character == "'"
- and query[i - 1] != "\\"
- ):
- state = KQLSplitState.OUTSIDE_STRING
-
- elif (
- state == KQLSplitState.INSIDE_DOUBLE_QUOTED_STRING
- and character == '"'
- and query[i - 1] != "\\"
- ):
- state = KQLSplitState.OUTSIDE_STRING
-
- elif (
- state == KQLSplitState.INSIDE_MULTILINE_STRING
- and character == "`"
- and query[i - 2 : i] == "``"
- ):
- state = KQLSplitState.OUTSIDE_STRING
-
- return statements
-
-
-class KustoKQLStatement(BaseSQLStatement[str]):
- """
- Special class for Kusto KQL.
-
- Kusto KQL is a SQL-like language, but it's not supported by sqlglot. Queries look
- like this:
-
- StormEvents
- | summarize PropertyDamage = sum(DamageProperty) by State
- | join kind=innerunique PopulationData on State
- | project State, PropertyDamagePerCapita = PropertyDamage / Population
- | sort by PropertyDamagePerCapita
-
- See https://learn.microsoft.com/en-us/azure/data-explorer/kusto/query/ for more
- details about it.
- """
-
- @classmethod
- def split_query(
- cls,
- query: str,
- engine: str,
- ) -> list[KustoKQLStatement]:
- """
- Split a query at semi-colons.
-
- Since we don't have a parser, we use a simple state machine based function. See
- https://learn.microsoft.com/en-us/azure/data-explorer/kusto/query/scalar-data-types/string
- for more information.
- """
- return [cls(statement, engine) for statement in split_kql(query)]
-
- @classmethod
- def _parse_statement(
- cls,
- statement: str,
- engine: str,
- ) -> str:
- if engine != "kustokql":
- raise SupersetParseError(f"Invalid engine: {engine}")
-
- statements = split_kql(statement)
- if len(statements) != 1:
- raise SupersetParseError("SQLStatement should have exactly one statement")
-
- return statements[0].strip()
-
- @classmethod
- def _extract_tables_from_statement(cls, parsed: str, engine: str) -> set[Table]:
- """
- Extract all tables referenced in the statement.
-
- StormEvents
- | where InjuriesDirect + InjuriesIndirect > 50
- | join (PopulationData) on State
- | project State, Population, TotalInjuries = InjuriesDirect + InjuriesIndirect
-
- """
- logger.warning(
- "Kusto KQL doesn't support table extraction. This means that data access "
- "roles will not be enforced by Superset in the database."
- )
- return set()
-
- def format(self, comments: bool = True) -> str:
- """
- Pretty-format the SQL statement.
- """
- return self._parsed
-
- def get_settings(self) -> dict[str, str | bool]:
- """
- Return the settings for the SQL statement.
-
- >>> statement = KustoKQLStatement("set querytrace;")
- >>> statement.get_settings()
- {"querytrace": True}
-
- """
- set_regex = r"^set\s+(?P<name>\w+)(?:\s*=\s*(?P<value>\w+))?$"
- if match := re.match(set_regex, self._parsed, re.IGNORECASE):
- return {match.group("name"): match.group("value") or True}
-
- return {}
-
- def is_mutating(self) -> bool:
- """
- Check if the statement mutates data (DDL/DML).
-
- :return: True if the statement mutates data.
- """
- return self._parsed.startswith(".") and not self._parsed.startswith(".show")
-
-
-class SQLScript:
- """
- A SQL script, with 0+ statements.
- """
-
- # Special engines that can't be parsed using sqlglot. Supporting non-SQL engines
- # adds a lot of complexity to Superset, so we should avoid adding new engines to
- # this data structure.
- special_engines = {
- "kustokql": KustoKQLStatement,
- }
-
- def __init__(
- self,
- query: str,
- engine: str,
- ):
- statement_class = self.special_engines.get(engine, SQLStatement)
- self.statements = statement_class.split_query(query, engine)
-
- def format(self, comments: bool = True) -> str:
- """
- Pretty-format the SQL query.
- """
- return ";\n".join(statement.format(comments) for statement in self.statements)
-
- def get_settings(self) -> dict[str, str | bool]:
- """
- Return the settings for the SQL query.
-
- >>> statement = SQLScript("SET foo = 'bar'; SET foo = 'baz'")
- >>> statement.get_settings()
- {"foo": "'baz'"}
-
- """
- settings: dict[str, str | bool] = {}
- for statement in self.statements:
- settings.update(statement.get_settings())
-
- return settings
-
- def has_mutation(self) -> bool:
- """
- Check if the script contains mutating statements.
-
- :return: True if the script contains mutating statements
- """
- return any(statement.is_mutating() for statement in self.statements)
-
-
class ParsedQuery:
def __init__(
self,
@@ -883,77 +341,6 @@ class ParsedQuery:
if statement
}
- def _extract_tables_from_statement(self, statement: exp.Expression) -> set[Table]:
- """
- Extract all table references in a single statement.
-
- Please not that this is not trivial; consider the following queries:
-
- DESCRIBE some_table;
- SHOW PARTITIONS FROM some_table;
- WITH masked_name AS (SELECT * FROM some_table) SELECT * FROM masked_name;
-
- See the unit tests for other tricky cases.
- """
- sources: Iterable[exp.Table]
-
- if isinstance(statement, exp.Describe):
- # A `DESCRIBE` query has no sources in sqlglot, so we need to explicitly
- # query for all tables.
- sources = statement.find_all(exp.Table)
- elif isinstance(statement, exp.Command):
- # Commands, like `SHOW COLUMNS FROM foo`, have to be converted into a
- # `SELECT` statetement in order to extract tables.
- if not (literal := statement.find(exp.Literal)):
- return set()
-
- try:
- pseudo_query = parse_one(
- f"SELECT {literal.this}",
- dialect=self._dialect,
- )
- sources = pseudo_query.find_all(exp.Table)
- except SqlglotError:
- return set()
- else:
- sources = [
- source
- for scope in traverse_scope(statement)
- for source in scope.sources.values()
- if isinstance(source, exp.Table) and not self._is_cte(source, scope)
- ]
-
- return {
- Table(
- source.name,
- source.db if source.db != "" else None,
- source.catalog if source.catalog != "" else None,
- )
- for source in sources
- }
-
- def _is_cte(self, source: exp.Table, scope: Scope) -> bool:
- """
- Is the source a CTE?
-
- CTEs in the parent scope look like tables (and are represented by
- exp.Table objects), but should not be considered as such;
- otherwise a user with access to table `foo` could access any table
- with a query like this:
-
- WITH foo AS (SELECT * FROM target_table) SELECT * FROM foo
-
- """
- parent_sources = scope.parent.sources if scope.parent else {}
- ctes_in_scope = {
- name
- for name, parent_scope in parent_sources.items()
- if isinstance(parent_scope, Scope)
- and parent_scope.scope_type == ScopeType.CTE
- }
-
- return source.name in ctes_in_scope
-
@property
def limit(self) -> int | None:
return self._limit
diff --git a/superset/sqllab/api.py b/superset/sqllab/api.py
index f7d66ed4e1..2403a36583 100644
--- a/superset/sqllab/api.py
+++ b/superset/sqllab/api.py
@@ -35,8 +35,8 @@ from superset.daos.query import QueryDAO
from superset.extensions import event_logger
from superset.jinja_context import get_template_processor
from superset.models.sql_lab import Query
+from superset.sql.parse import SQLScript
from superset.sql_lab import get_sql_results
-from superset.sql_parse import SQLScript
from superset.sqllab.command_status import SqlJsonExecutionStatus
from superset.sqllab.exceptions import (
QueryIsForbiddenToAccessException,
diff --git a/superset/tasks/thumbnails.py b/superset/tasks/thumbnails.py
index eb64ac49cb..e37e874355 100644
--- a/superset/tasks/thumbnails.py
+++ b/superset/tasks/thumbnails.py
@@ -107,7 +107,6 @@ def cache_dashboard_thumbnail(
)
-# pylint: disable=too-many-arguments
@celery_app.task(name="cache_dashboard_screenshot", soft_time_limit=300)
def cache_dashboard_screenshot(
dashboard_id: int,
diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py
index 216d896a5e..be9ebd38ae 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/tests/unit_tests/sql_parse_tests.py
@@ -30,6 +30,13 @@ from superset.exceptions import (
QueryClauseValidationException,
SupersetSecurityException,
)
+from superset.sql.parse import (
+ KustoKQLStatement,
+ split_kql,
+ SQLScript,
+ SQLStatement,
+ Table,
+)
from superset.sql_parse import (
add_table_name,
check_sql_functions_exist,
@@ -39,14 +46,9 @@ 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,
- Table,
)