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/05/24 00:44:28 UTC
(superset) branch rls-sqlglot updated (46c38a6aeb -> d170de4217)
This is an automated email from the ASF dual-hosted git repository.
beto pushed a change to branch rls-sqlglot
in repository https://gitbox.apache.org/repos/asf/superset.git
discard 46c38a6aeb WIP
new d170de4217 WIP
This update added new revisions after undoing existing revisions.
That is to say, some revisions that were in the old version of the
branch are not in the new version. This situation occurs
when a user --force pushes a change and generates a repository
containing something like this:
* -- * -- B -- O -- O -- O (46c38a6aeb)
\
N -- N -- N refs/heads/rls-sqlglot (d170de4217)
You should already have received notification emails for all of the O
revisions, and so the following emails describe only the N revisions
from the common base, B.
Any revisions marked "omit" are not gone; other references still
refer to them. Any revisions marked "discard" are gone forever.
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.
Summary of changes:
superset/sql_lab.py | 70 +++++-----------
superset/sql_parse.py | 154 ++++++++++++++++++++++++++++++------
tests/unit_tests/sql_parse_tests.py | 2 +-
3 files changed, 152 insertions(+), 74 deletions(-)
(superset) 01/01: WIP
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
beto pushed a commit to branch rls-sqlglot
in repository https://gitbox.apache.org/repos/asf/superset.git
commit d170de4217d6a9514c0fe94c71038de91da21ea6
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Wed May 22 22:48:42 2024 -0700
WIP
---
superset/connectors/sqla/models.py | 49 +-----
superset/models/helpers.py | 6 +-
superset/sql_lab.py | 70 ++------
superset/sql_parse.py | 292 +++++++++++++++++++++++++++++--
tests/unit_tests/sql_parse_tests.py | 333 ++++++++++++++++++++++++++----------
5 files changed, 554 insertions(+), 196 deletions(-)
diff --git a/superset/connectors/sqla/models.py b/superset/connectors/sqla/models.py
index 8c74dfd589..6dd0d76b69 100644
--- a/superset/connectors/sqla/models.py
+++ b/superset/connectors/sqla/models.py
@@ -22,7 +22,6 @@ import dataclasses
import json
import logging
import re
-from collections import defaultdict
from collections.abc import Hashable
from dataclasses import dataclass, field
from datetime import datetime, timedelta
@@ -72,7 +71,7 @@ from sqlalchemy.sql.elements import ColumnClause, TextClause
from sqlalchemy.sql.expression import Label, TextAsFrom
from sqlalchemy.sql.selectable import Alias, TableClause
-from superset import app, db, is_feature_enabled, security_manager
+from superset import app, db, security_manager
from superset.commands.dataset.exceptions import DatasetNotFoundError
from superset.common.db_query_status import QueryStatus
from superset.connectors.sqla.utils import (
@@ -175,7 +174,9 @@ class DatasourceKind(StrEnum):
PHYSICAL = "physical"
-class BaseDatasource(AuditMixinNullable, ImportExportMixin): # pylint: disable=too-many-public-methods
+class BaseDatasource(
+ AuditMixinNullable, ImportExportMixin
+): # pylint: disable=too-many-public-methods
"""A common interface to objects that are queryable
(tables and datasources)"""
@@ -1605,48 +1606,6 @@ class SqlaTable(
if is_alias_used_in_orderby(col):
col.name = f"{col.name}__"
- def get_sqla_row_level_filters(
- self,
- template_processor: BaseTemplateProcessor,
- ) -> list[TextClause]:
- """
- Return the appropriate row level security filters for this table and the
- current user. A custom username can be passed when the user is not present in the
- Flask global namespace.
-
- :param template_processor: The template processor to apply to the filters.
- :returns: A list of SQL clauses to be ANDed together.
- """
- all_filters: list[TextClause] = []
- filter_groups: dict[int | str, list[TextClause]] = defaultdict(list)
- try:
- for filter_ in security_manager.get_rls_filters(self):
- clause = self.text(
- f"({template_processor.process_template(filter_.clause)})"
- )
- if filter_.group_key:
- filter_groups[filter_.group_key].append(clause)
- else:
- all_filters.append(clause)
-
- if is_feature_enabled("EMBEDDED_SUPERSET"):
- for rule in security_manager.get_guest_rls_filters(self):
- clause = self.text(
- f"({template_processor.process_template(rule['clause'])})"
- )
- all_filters.append(clause)
-
- grouped_filters = [or_(*clauses) for clauses in filter_groups.values()]
- all_filters.extend(grouped_filters)
- return all_filters
- except TemplateError as ex:
- raise QueryObjectValidationError(
- _(
- "Error in jinja expression in RLS filters: %(msg)s",
- msg=ex.message,
- )
- ) from ex
-
def text(self, clause: str) -> TextClause:
return self.db_engine_spec.get_text_clause(clause)
diff --git a/superset/models/helpers.py b/superset/models/helpers.py
index 5845a400c2..790ed7e76c 100644
--- a/superset/models/helpers.py
+++ b/superset/models/helpers.py
@@ -17,6 +17,8 @@
# pylint: disable=too-many-lines
"""a collection of model-related helper classes and functions"""
+from __future__ import annotations
+
import builtins
import dataclasses
import json
@@ -807,7 +809,7 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
def get_sqla_row_level_filters(
self,
- template_processor: BaseTemplateProcessor,
+ template_processor: BaseTemplateProcessor | None,
) -> list[TextClause]:
"""
Return the appropriate row level security filters for this table and the
@@ -817,6 +819,8 @@ class ExploreMixin: # pylint: disable=too-many-public-methods
:param template_processor: The template processor to apply to the filters.
:returns: A list of SQL clauses to be ANDed together.
"""
+ template_processor = template_processor or self.get_template_processor()
+
all_filters: list[TextClause] = []
filter_groups: dict[Union[int, str], list[TextClause]] = defaultdict(list)
try:
diff --git a/superset/sql_lab.py b/superset/sql_lab.py
index 3f8c1cc737..4c67e84ae9 100644
--- a/superset/sql_lab.py
+++ b/superset/sql_lab.py
@@ -52,9 +52,8 @@ from superset.models.sql_lab import Query
from superset.result_set import SupersetResultSet
from superset.sql_parse import (
CtasMethod,
- insert_rls_as_subquery,
- insert_rls_in_predicate,
ParsedQuery,
+ SQLStatement,
Table,
)
from superset.sqllab.limiting_factor import LimitingFactor
@@ -203,67 +202,49 @@ def execute_sql_statement( # pylint: disable=too-many-statements
database: Database = query.database
db_engine_spec = database.db_engine_spec
- parsed_query = ParsedQuery(sql_statement, engine=db_engine_spec.engine)
- if is_feature_enabled("RLS_IN_SQLLAB"):
- # There are two ways to insert RLS: either replacing the table with a subquery
- # that has the RLS, or appending the RLS to the ``WHERE`` clause. The former is
- # safer, but not supported in all databases.
- insert_rls = (
- insert_rls_as_subquery
- if database.db_engine_spec.allows_subqueries
- and database.db_engine_spec.allows_alias_in_select
- else insert_rls_in_predicate
- )
+ parsed_statement = SQLStatement(sql_statement, engine=db_engine_spec.engine)
- # Insert any applicable RLS predicates
- parsed_query = ParsedQuery(
- str(
- insert_rls(
- parsed_query._parsed[0], # pylint: disable=protected-access
- database.id,
- query.schema,
- )
- ),
- engine=db_engine_spec.engine,
- )
-
- sql = parsed_query.stripped()
-
- # This is a test to see if the query is being
- # limited by either the dropdown or the sql.
- # We are testing to see if more rows exist than the limit.
- increased_limit = None if query.limit is None else query.limit + 1
+ if is_feature_enabled("RLS_IN_SQLLAB"):
+ default_schema = database.get_default_schema_for_query(query)
+ parsed_statement = parsed_statement.apply_rls(query.catalog, default_schema)
- if not db_engine_spec.is_readonly_query(parsed_query) and not database.allow_dml:
+ if parsed_statement.is_dml() and not database.allow_dml:
raise SupersetErrorException(
SupersetError(
- message=__("Only SELECT statements are allowed against this database."),
+ message=__("Only read-only statements are allowed in this database."),
error_type=SupersetErrorType.DML_NOT_ALLOWED_ERROR,
level=ErrorLevel.ERROR,
)
)
+
if apply_ctas:
if not query.tmp_table_name:
start_dttm = datetime.fromtimestamp(query.start_time)
query.tmp_table_name = (
f'tmp_{query.user_id}_table_{start_dttm.strftime("%Y_%m_%d_%H_%M_%S")}'
)
- sql = parsed_query.as_create_table(
- query.tmp_table_name,
- schema_name=query.tmp_schema_name,
+ parsed_statement = parsed_statement.as_create_table(
+ Table(query.tmp_table_name, query.tmp_schema_name, query.catalog),
method=query.ctas_method,
)
query.select_as_cta_used = True
+ increased_limit = None if query.limit is None else query.limit + 1
+
# Do not apply limit to the CTA queries when SQLLAB_CTAS_NO_LIMIT is set to true
- if db_engine_spec.is_select_query(parsed_query) and not (
+ if parsed_statement.is_select() and not (
query.select_as_cta_used and SQLLAB_CTAS_NO_LIMIT
):
if SQL_MAX_ROW and (not query.limit or query.limit > SQL_MAX_ROW):
query.limit = SQL_MAX_ROW
- sql = apply_limit_if_exists(database, increased_limit, query, sql)
+
+ if query.limit:
+ # Increase limit by one so we can test if there are more rows when the
+ # database returns exactly the number of rows requested by the user.
+ parsed_statement = parsed_statement.apply_limit(increased_limit)
# Hook to allow environment-specific mutation (usually comments) to the SQL
+ sql = parsed_statement.format(strip=True)
sql = database.mutate_sql_based_on_config(sql)
try:
query.executed_sql = sql
@@ -331,19 +312,6 @@ def execute_sql_statement( # pylint: disable=too-many-statements
return SupersetResultSet(data, cursor_description, db_engine_spec)
-def apply_limit_if_exists(
- database: Database, increased_limit: Optional[int], query: Query, sql: str
-) -> str:
- if query.limit and increased_limit:
- # We are fetching one more than the requested limit in order
- # to test whether there are more rows than the limit. According to the DB
- # Engine support it will choose top or limit parse
- # Later, the extra row will be dropped before sending
- # the results back to the user.
- sql = database.apply_limit_to_sql(sql, increased_limit, force=True)
- return sql
-
-
def _serialize_payload(
payload: dict[Any, Any], use_msgpack: Optional[bool] = False
) -> Union[bytes, str]:
diff --git a/superset/sql_parse.py b/superset/sql_parse.py
index f32647042b..86d822ecf2 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -32,7 +32,7 @@ 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 import exp
from sqlglot.dialects.dialect import Dialect, Dialects
from sqlglot.errors import ParseError, SqlglotError
from sqlglot.optimizer.scope import Scope, ScopeType, traverse_scope
@@ -61,6 +61,7 @@ from sqlparse.tokens import (
)
from sqlparse.utils import imt
+from superset.db_engine_specs import load_engine_specs
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import (
QueryClauseValidationException,
@@ -294,7 +295,7 @@ def extract_tables_from_statement(
return set()
try:
- pseudo_query = parse_one(f"SELECT {literal.this}", dialect=dialect)
+ pseudo_query = sqlglot.parse_one(f"SELECT {literal.this}", dialect=dialect)
except ParseError:
return set()
sources = pseudo_query.find_all(exp.Table)
@@ -419,7 +420,7 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
raise NotImplementedError()
- def format(self, comments: bool = True) -> str:
+ def format(self, comments: bool = True, strip: bool = False) -> str:
"""
Format the statement, optionally ommitting comments.
"""
@@ -437,10 +438,93 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
raise NotImplementedError()
+ def apply_rls(
+ self,
+ catalog: str | None,
+ schema: str | None,
+ ) -> InternalRepresentation:
+ """
+ Apply Row Level Security to the SQL.
+
+ :param database: The database where the SQL will run
+ :param catalog: The default catalog for non-qualified table names
+ :param schema: The default schema for non-qualified table names
+ :return: The SQL with RLS applied
+ """
+ raise NotImplementedError()
+
def __str__(self) -> str:
return self.format()
+class RLSAsPredicate:
+ """
+ Apply Row Level Security role as a predicate.
+
+ This transformer will apply any RLS predicates to the relevant tables. For example,
+ given the RLS rule:
+
+ table: some_table
+ clause: id = 42
+
+ If a user subject to the rule runs the following query:
+
+ SELECT foo FROM some_table WHERE bar = 'baz'
+
+ The query will be modified to:
+
+ SELECT foo FROM some_table WHERE bar = 'baz' AND id = 42
+
+ This approach is probably less secure than using subqueries, so it's only used for
+ databases without support for subqueries.
+ """
+
+ def __init__(self, rules: dict[Table, str]) -> None:
+ self.rules = rules
+
+ def __call__(self, node: exp.Expression) -> exp.Expression:
+ if not isinstance(node, exp.Select):
+ return node
+
+ table_node = node.find(exp.Table)
+ if not table_node:
+ return node
+
+ table = Table(
+ str(table_node.this),
+ str(table_node.db) if table_node.db else None,
+ str(table_node.catalog) if table_node.catalog else None,
+ )
+ if predicate := self.rules.get(table):
+ if where := node.args.get("where"):
+ predicate = exp.And(this=predicate, expression=where.this)
+
+ node.set("where", exp.Where(this=predicate))
+
+ return node
+
+
+class RLSAsSubquery:
+ def __init__(self, rules: dict[Table, str]) -> None:
+ self.rules = rules
+
+ def __call__(self, node: exp.Expression) -> exp.Expression:
+ if not isinstance(node, exp.Table):
+ return node
+
+ table = Table(
+ str(node.this),
+ str(node.db) if node.db else None,
+ str(node.catalog) if node.catalog else None,
+ )
+ if predicate := self.rules.get(table):
+ alias = node.alias
+ node.set("alias", None)
+ return f"(SELECT * FROM {node} WHERE {predicate}) AS {alias}"
+
+ return node
+
+
class SQLStatement(BaseSQLStatement[exp.Expression]):
"""
A SQL statement.
@@ -507,12 +591,19 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
dialect = SQLGLOT_DIALECTS.get(engine)
return extract_tables_from_statement(parsed, dialect)
- def format(self, comments: bool = True) -> str:
+ def format(self, comments: bool = True, strip: bool = False) -> str:
"""
Pretty-format the SQL statement.
"""
write = Dialect.get_or_raise(self._dialect)
- return write.generate(self._parsed, copy=False, comments=comments, pretty=True)
+ output = write.generate(
+ self._parsed,
+ copy=False,
+ comments=comments,
+ pretty=True,
+ )
+
+ return output.strip(" \t\r\n;") if strip else output
def get_settings(self) -> dict[str, str | bool]:
"""
@@ -529,6 +620,184 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
for eq in set_item.find_all(exp.EQ)
}
+ def apply_rls(
+ self,
+ catalog: str | None,
+ schema: str | None,
+ ) -> SQLStatement:
+ """
+ Apply Row Level Security to the SQL.
+
+ :param catalog: The default catalog for non-qualified table names
+ :param schema: The default schema for non-qualified table names
+ :return: The SQL with RLS applied
+ """
+ statement = self._parsed.copy()
+
+ # collect all relevant RLS rules
+ rules = {}
+ for table in self.tables:
+ if rls := self._get_rls_for_table(table, catalog, schema):
+ rules[table] = rls
+
+ if not rules:
+ return statement
+
+ use_subquery = all(
+ engine_spec.allows_subqueries
+ for engine_spec in load_engine_specs()
+ if engine_spec.engine == self.engine
+ )
+ transformer = RLSAsSubquery(rules) if use_subquery else RLSAsPredicate(rules)
+
+ return SQLStatement(statement.transform(transformer), self.engine)
+
+ def _get_rls_for_table(
+ self,
+ database: Database,
+ table: Table,
+ catalog: str | None,
+ schema: str | None,
+ ) -> exp.Expression | None:
+ """
+ Get the RLS for a table.
+
+ :param table: The table to get the RLS for
+ :param catalog: The default catalog for non-qualified table names
+ :param schema: The default schema for non-qualified table names
+ :return: The RLS for the table
+ """
+ # pylint: disable=import-outside-toplevel
+ from superset import db
+ from superset.connectors.sqla.models import SqlaTable
+
+ dataset = db.session.query(SqlaTable).filter(
+ and_(
+ SqlaTable.database_id == database.id,
+ SqlaTable.catalog == table.catalog or catalog,
+ SqlaTable.schema == table.schema or schema,
+ SqlaTable.table_name == table.table,
+ ).one_or_none()
+ )
+ if not dataset:
+ return None
+
+ filters = dataset.get_sqla_row_level_filters()
+ if not filters:
+ return None
+
+ rls = and_(*filters).compile(
+ dialect=database.get_dialect(),
+ compile_kwargs={"literal_binds": True},
+ )
+ dialect = SQLGLOT_DIALECTS.get(self.engine)
+
+ return sqlglot.parse_one(str(rls), dialect=dialect)
+
+ def is_dml(self) -> bool:
+ """
+ Check if the statement is DML.
+
+ :return: True if the statement is DML
+ """
+ for node in self._parsed.walk():
+ if isinstance(
+ node,
+ (
+ exp.Insert,
+ exp.Update,
+ exp.Delete,
+ exp.Merge,
+ exp.Create,
+ exp.Alter,
+ exp.Drop,
+ exp.TruncateTable,
+ ),
+ ):
+ return True
+
+ return False
+
+ def as_create_table(self, table: Table, method: CtasMethod) -> SQLStatement:
+ """
+ Convert the statement to a CREATE TABLE statement.
+ """
+ create_table = exp.Create(
+ this=sqlglot.parse_one(table, into=exp.Table),
+ kind=method.value,
+ expression=self._parsed.copy(),
+ )
+
+ return SQLStatement(create_table, self.engine)
+
+ def is_select(self) -> bool:
+ """
+ Check if the statement is a SELECT statement.
+
+ :return: True if the statement is a SELECT statement
+ """
+ return isinstance(self._parsed, exp.Select)
+
+ def apply_limit(self, limit: int, force: bool = False) -> SQLStatement:
+ """
+ Apply a limit to the SQL.
+
+ There are 3 strategies to limit queries, defined in the DB engine spec:
+
+ 1. `FORCE_LIMIT`: a limit is added to the query, or the existing one is
+ replaced. This is the most efficient, since the database will produce at
+ most the number of rows that Superset will display.
+ 2. `WRAP_SQL`: the query is wrapped in a subquery, and the limit is applied
+ to the outer query. This might be inneficient, since the database
+ optimizer might not be able to push the limit down to the inner query.
+ 3. `FETCH_MANY`: no limit is applied, but only `LIMIT` rows are fetched from
+ the database. This is the least efficient, unless the database computes
+ rows as they are read by the cursor, which is unlikely.
+
+ :param limit: The limit to apply
+ :param force: Apply limit even when a lower one is present
+ :return: The SQL with the limit applied
+ """
+ from superset.db_engine_specs.base import LimitMethod
+
+ methods = {
+ engine_spec.limit_method
+ for engine_spec in load_engine_specs()
+ if engine_spec.engine == self.engine
+ }
+ if not methods:
+ methods = {LimitMethod.FETCH_MANY}
+
+ # When multiple methods are supported, we prefer the more generic one --
+ # usually less efficient.
+ preference = [
+ LimitMethod.FETCH_MANY,
+ LimitMethod.WRAP_SQL,
+ LimitMethod.FORCE_LIMIT,
+ ]
+ method = sorted(methods, key=preference.index)[0]
+
+ if not self.is_select() or method == LimitMethod.FETCH_MANY:
+ return SQLStatement(self._parsed.copy(), self.engine)
+
+ if method == LimitMethod.WRAP_SQL:
+ limited = exp.Select(
+ expressions=[exp.Star()],
+ from_=exp.Subquery(subquery=self._parsed.copy(), alias="inner_qry"),
+ limit=exp.Literal.number(limit),
+ )
+ return SQLStatement(limited, self.engine)
+
+ current_limit: int | None = None
+ for node in self._parsed.find_all(exp.Limit):
+ current_limit = int(node.expression.this)
+ break
+
+ if force or current_limit is None or limit < current_limit:
+ return SQLStatement(self._parsed.limit(limit), self.engine)
+
+ return SQLStatement(self._parsed.copy(), self.engine)
+
class KQLSplitState(enum.Enum):
"""
@@ -652,11 +921,11 @@ class KustoKQLStatement(BaseSQLStatement[str]):
)
return set()
- def format(self, comments: bool = True) -> str:
+ def format(self, comments: bool = True, strip: bool = False) -> str:
"""
Pretty-format the SQL statement.
"""
- return self._parsed
+ return self._parsed.strip(" \t\r\n;") if strip else self._parsed
def get_settings(self) -> dict[str, str | bool]:
"""
@@ -698,7 +967,10 @@ class SQLScript:
"""
Pretty-format the SQL query.
"""
- return ";\n".join(statement.format(comments) for statement in self.statements)
+ return (
+ ";\n".join(statement.format(comments) for statement in self.statements)
+ + ";"
+ )
def get_settings(self) -> dict[str, str | bool]:
"""
@@ -750,7 +1022,7 @@ class ParsedQuery:
Note: this uses sqlglot, since it's better at catching more edge cases.
"""
try:
- statements = parse(self.stripped(), dialect=self._dialect)
+ statements = sqlglot.parse(self.stripped(), dialect=self._dialect)
except SqlglotError as ex:
logger.warning("Unable to parse SQL (%s): %s", self._dialect, self.sql)
@@ -804,7 +1076,7 @@ class ParsedQuery:
return set()
try:
- pseudo_query = parse_one(
+ pseudo_query = sqlglot.parse_one(
f"SELECT {literal.this}",
dialect=self._dialect,
)
diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py
index 3b80b7e01d..64ad8650d4 100644
--- a/tests/unit_tests/sql_parse_tests.py
+++ b/tests/unit_tests/sql_parse_tests.py
@@ -20,6 +20,7 @@ from typing import Optional
from unittest.mock import Mock
import pytest
+import sqlglot
import sqlparse
from pytest_mock import MockerFixture
from sqlalchemy import text
@@ -40,6 +41,8 @@ from superset.sql_parse import (
insert_rls_in_predicate,
KustoKQLStatement,
ParsedQuery,
+ RLSAsPredicate,
+ RLSAsSubquery,
sanitize_clause,
split_kql,
SQLScript,
@@ -118,8 +121,9 @@ def test_extract_tables_subselect() -> None:
"""
Test that tables inside subselects are parsed correctly.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT sub.*
FROM (
SELECT *
@@ -128,10 +132,13 @@ FROM (
) sub, s2.t2
WHERE sub.resolution = 'NONE'
"""
- ) == {Table("t1", "s1"), Table("t2", "s2")}
+ )
+ == {Table("t1", "s1"), Table("t2", "s2")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT sub.*
FROM (
SELECT *
@@ -140,10 +147,13 @@ FROM (
) sub
WHERE sub.resolution = 'NONE'
"""
- ) == {Table("t1", "s1")}
+ )
+ == {Table("t1", "s1")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT * FROM t1
WHERE s11 > ANY (
SELECT COUNT(*) /* no hint */ FROM t2
@@ -155,7 +165,9 @@ WHERE s11 > ANY (
)
)
"""
- ) == {Table("t1"), Table("t2"), Table("t3"), Table("t4")}
+ )
+ == {Table("t1"), Table("t2"), Table("t3"), Table("t4")}
+ )
def test_extract_tables_select_in_expression() -> None:
@@ -226,24 +238,30 @@ def test_extract_tables_select_array() -> None:
"""
Test that queries selecting arrays work as expected.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT ARRAY[1, 2, 3] AS my_array
FROM t1 LIMIT 10
"""
- ) == {Table("t1")}
+ )
+ == {Table("t1")}
+ )
def test_extract_tables_select_if() -> None:
"""
Test that queries with an ``IF`` work as expected.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT IF(CARDINALITY(my_array) >= 3, my_array[3], NULL)
FROM t1 LIMIT 10
"""
- ) == {Table("t1")}
+ )
+ == {Table("t1")}
+ )
def test_extract_tables_with_catalog() -> None:
@@ -311,29 +329,38 @@ def test_extract_tables_where_subquery() -> None:
"""
Test that tables in a ``WHERE`` subquery are parsed correctly.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT name
FROM t1
WHERE regionkey = (SELECT max(regionkey) FROM t2)
"""
- ) == {Table("t1"), Table("t2")}
+ )
+ == {Table("t1"), Table("t2")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT name
FROM t1
WHERE regionkey IN (SELECT regionkey FROM t2)
"""
- ) == {Table("t1"), Table("t2")}
+ )
+ == {Table("t1"), Table("t2")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT name
FROM t1
WHERE EXISTS (SELECT 1 FROM t2 WHERE t1.regionkey = t2.regionkey);
"""
- ) == {Table("t1"), Table("t2")}
+ )
+ == {Table("t1"), Table("t2")}
+ )
def test_extract_tables_describe() -> None:
@@ -347,12 +374,15 @@ def test_extract_tables_show_partitions() -> None:
"""
Test ``SHOW PARTITIONS``.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SHOW PARTITIONS FROM orders
WHERE ds >= '2013-01-01' ORDER BY ds DESC
"""
- ) == {Table("orders")}
+ )
+ == {Table("orders")}
+ )
def test_extract_tables_join() -> None:
@@ -364,8 +394,9 @@ def test_extract_tables_join() -> None:
Table("t2"),
}
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT a.date, b.name
FROM left_table a
JOIN (
@@ -376,10 +407,13 @@ JOIN (
) b
ON a.date = b.date
"""
- ) == {Table("left_table"), Table("right_table")}
+ )
+ == {Table("left_table"), Table("right_table")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT a.date, b.name
FROM left_table a
LEFT INNER JOIN (
@@ -390,10 +424,13 @@ LEFT INNER JOIN (
) b
ON a.date = b.date
"""
- ) == {Table("left_table"), Table("right_table")}
+ )
+ == {Table("left_table"), Table("right_table")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT a.date, b.name
FROM left_table a
RIGHT OUTER JOIN (
@@ -404,10 +441,13 @@ RIGHT OUTER JOIN (
) b
ON a.date = b.date
"""
- ) == {Table("left_table"), Table("right_table")}
+ )
+ == {Table("left_table"), Table("right_table")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT a.date, b.name
FROM left_table a
FULL OUTER JOIN (
@@ -418,15 +458,18 @@ FULL OUTER JOIN (
) b
ON a.date = b.date
"""
- ) == {Table("left_table"), Table("right_table")}
+ )
+ == {Table("left_table"), Table("right_table")}
+ )
def test_extract_tables_semi_join() -> None:
"""
Test ``LEFT SEMI JOIN``.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT a.date, b.name
FROM left_table a
LEFT SEMI JOIN (
@@ -437,15 +480,18 @@ LEFT SEMI JOIN (
) b
ON a.data = b.date
"""
- ) == {Table("left_table"), Table("right_table")}
+ )
+ == {Table("left_table"), Table("right_table")}
+ )
def test_extract_tables_combinations() -> None:
"""
Test a complex case with nested queries.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT * FROM t1
WHERE s11 > ANY (
SELECT * FROM t1 UNION ALL SELECT * FROM (
@@ -459,10 +505,13 @@ WHERE s11 > ANY (
)
)
"""
- ) == {Table("t1"), Table("t3"), Table("t4"), Table("t6")}
+ )
+ == {Table("t1"), Table("t3"), Table("t4"), Table("t6")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT * FROM (
SELECT * FROM (
SELECT * FROM (
@@ -471,45 +520,56 @@ SELECT * FROM (
) AS S2
) AS S3
"""
- ) == {Table("EmployeeS")}
+ )
+ == {Table("EmployeeS")}
+ )
def test_extract_tables_with() -> None:
"""
Test ``WITH``.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
WITH
x AS (SELECT a FROM t1),
y AS (SELECT a AS b FROM t2),
z AS (SELECT b AS c FROM t3)
SELECT c FROM z
"""
- ) == {Table("t1"), Table("t2"), Table("t3")}
+ )
+ == {Table("t1"), Table("t2"), Table("t3")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
WITH
x AS (SELECT a FROM t1),
y AS (SELECT a AS b FROM x),
z AS (SELECT b AS c FROM y)
SELECT c FROM z
"""
- ) == {Table("t1")}
+ )
+ == {Table("t1")}
+ )
def test_extract_tables_reusing_aliases() -> None:
"""
Test that the parser follows aliases.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
with q1 as ( select key from q2 where key = '5'),
q2 as ( select key from src where key = '5')
select * from (select key from q1) a
"""
- ) == {Table("src")}
+ )
+ == {Table("src")}
+ )
# weird query with circular dependency
assert (
@@ -546,8 +606,9 @@ def test_extract_tables_complex() -> None:
"""
Test a few complex queries.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT sum(m_examples) AS "sum__m_example"
FROM (
SELECT
@@ -568,23 +629,29 @@ FROM (
ORDER BY "sum__m_example" DESC
LIMIT 10;
"""
- ) == {
- Table("my_l_table"),
- Table("my_b_table"),
- Table("my_t_table"),
- Table("inner_table"),
- }
+ )
+ == {
+ Table("my_l_table"),
+ Table("my_b_table"),
+ Table("my_t_table"),
+ Table("inner_table"),
+ }
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT *
FROM table_a AS a, table_b AS b, table_c as c
WHERE a.id = b.id and b.id = c.id
"""
- ) == {Table("table_a"), Table("table_b"), Table("table_c")}
+ )
+ == {Table("table_a"), Table("table_b"), Table("table_c")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT somecol AS somecol
FROM (
WITH bla AS (
@@ -628,51 +695,63 @@ FROM (
LIMIT 50000
)
"""
- ) == {Table("a"), Table("b"), Table("c"), Table("d"), Table("e"), Table("f")}
+ )
+ == {Table("a"), Table("b"), Table("c"), Table("d"), Table("e"), Table("f")}
+ )
def test_extract_tables_mixed_from_clause() -> None:
"""
Test that the parser handles a ``FROM`` clause with table and subselect.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
SELECT *
FROM table_a AS a, (select * from table_b) AS b, table_c as c
WHERE a.id = b.id and b.id = c.id
"""
- ) == {Table("table_a"), Table("table_b"), Table("table_c")}
+ )
+ == {Table("table_a"), Table("table_b"), Table("table_c")}
+ )
def test_extract_tables_nested_select() -> None:
"""
Test that the parser handles selects inside functions.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
select (extractvalue(1,concat(0x7e,(select GROUP_CONCAT(TABLE_NAME)
from INFORMATION_SCHEMA.COLUMNS
WHERE TABLE_SCHEMA like "%bi%"),0x7e)));
""",
- "mysql",
- ) == {Table("COLUMNS", "INFORMATION_SCHEMA")}
+ "mysql",
+ )
+ == {Table("COLUMNS", "INFORMATION_SCHEMA")}
+ )
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
select (extractvalue(1,concat(0x7e,(select GROUP_CONCAT(COLUMN_NAME)
from INFORMATION_SCHEMA.COLUMNS
WHERE TABLE_NAME="bi_achievement_daily"),0x7e)));
""",
- "mysql",
- ) == {Table("COLUMNS", "INFORMATION_SCHEMA")}
+ "mysql",
+ )
+ == {Table("COLUMNS", "INFORMATION_SCHEMA")}
+ )
def test_extract_tables_complex_cte_with_prefix() -> None:
"""
Test that the parser handles CTEs with prefixes.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
WITH CTE__test (SalesPersonID, SalesOrderID, SalesYear)
AS (
SELECT SalesPersonID, SalesOrderID, YEAR(OrderDate) AS SalesYear
@@ -684,21 +763,26 @@ FROM CTE__test
GROUP BY SalesYear, SalesPersonID
ORDER BY SalesPersonID, SalesYear;
"""
- ) == {Table("SalesOrderHeader")}
+ )
+ == {Table("SalesOrderHeader")}
+ )
def test_extract_tables_identifier_list_with_keyword_as_alias() -> None:
"""
Test that aliases that are keywords are parsed correctly.
"""
- assert extract_tables(
- """
+ assert (
+ extract_tables(
+ """
WITH
f AS (SELECT * FROM foo),
match AS (SELECT * FROM f)
SELECT * FROM match
"""
- ) == {Table("foo")}
+ )
+ == {Table("foo")}
+ )
def test_update() -> None:
@@ -1815,7 +1899,7 @@ def test_sqlquery() -> None:
script = SQLScript("SELECT 1; SELECT 2;", "sqlite")
assert len(script.statements) == 2
- assert script.format() == "SELECT\n 1;\nSELECT\n 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;", "sqlite")
@@ -2032,3 +2116,74 @@ on $left.Day1 == $right.Day
| project Day1, Day2, Percentage = count_*100.0/count_1
""",
]
+
+
+@pytest.mark.parametrize(
+ "sql,rules,expected",
+ [
+ (
+ "SELECT t.foo FROM some_table AS t",
+ {Table("some_table"): "id = 42"},
+ "SELECT t.foo FROM (SELECT * FROM some_table WHERE id = 42) AS t",
+ ),
+ (
+ "SELECT t.foo FROM some_table AS t WHERE bar = 'baz'",
+ {Table("some_table"): "id = 42"},
+ (
+ "SELECT t.foo FROM (SELECT * FROM some_table WHERE id = 42) AS t "
+ "WHERE bar = 'baz'"
+ ),
+ ),
+ (
+ "SELECT t.foo FROM schema1.some_table AS t",
+ {Table("some_table", "schema1"): "id = 42"},
+ "SELECT t.foo FROM (SELECT * FROM schema1.some_table WHERE id = 42) AS t",
+ ),
+ (
+ "SELECT t.foo FROM schema1.some_table AS t",
+ {Table("some_table", "schema2"): "id = 42"},
+ "SELECT t.foo FROM schema1.some_table AS t",
+ ),
+ (
+ "SELECT t.foo FROM catalog1.schema1.some_table AS t",
+ {Table("some_table", "schema1", "catalog1"): "id = 42"},
+ "SELECT t.foo FROM (SELECT * FROM catalog1.schema1.some_table WHERE id = 42) AS t",
+ ),
+ (
+ "SELECT t.foo FROM catalog1.schema1.some_table AS t",
+ {Table("some_table", "schema1", "catalog2"): "id = 42"},
+ "SELECT t.foo FROM catalog1.schema1.some_table AS t",
+ ),
+ ],
+)
+def test_RLSAsSubquery(sql: str, rules: dict[Table, str], expected: str) -> None:
+ """
+ Test the `RLSAsSubquery` transformer.
+ """
+ statement = sqlglot.parse_one(sql)
+ transformer = RLSAsSubquery(rules)
+ assert str(statement.transform(transformer)) == expected
+
+
+@pytest.mark.parametrize(
+ "sql,rules,expected",
+ [
+ (
+ "SELECT t.foo FROM some_table AS t",
+ {Table("some_table"): "id = 42"},
+ "SELECT t.foo FROM some_table AS t WHERE id = 42",
+ ),
+ (
+ "SELECT t.foo FROM some_table AS t WHERE bar = 'baz'",
+ {Table("some_table"): "id = 42"},
+ "SELECT t.foo FROM some_table AS t WHERE id = 42 AND bar = 'baz'",
+ ),
+ ],
+)
+def test_RLSAsPredicate(sql: str, rules: dict[Table, str], expected: str) -> None:
+ """
+ Test the `RLSAsPredicate` transformer.
+ """
+ statement = sqlglot.parse_one(sql)
+ transformer = RLSAsPredicate(rules)
+ assert str(statement.transform(transformer)) == expected