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/23 18:50:50 UTC
(superset) branch rls-sqlglot updated (988c628181 -> 46c38a6aeb)
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 988c628181 WIP
new 46c38a6aeb 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 (988c628181)
\
N -- N -- N refs/heads/rls-sqlglot (46c38a6aeb)
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_parse.py | 10 +++---
tests/unit_tests/sql_parse_tests.py | 72 ++++++++++++++++++++++++++++++-------
2 files changed, 63 insertions(+), 19 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 46c38a6aeb3978c86a53ad88c82e299fa6caf1e7
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_parse.py | 162 ++++++++++++++++++
tests/unit_tests/sql_parse_tests.py | 331 ++++++++++++++++++++++++++----------
4 files changed, 414 insertions(+), 134 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_parse.py b/superset/sql_parse.py
index f32647042b..46c6f78e29 100644
--- a/superset/sql_parse.py
+++ b/superset/sql_parse.py
@@ -437,10 +437,94 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
"""
raise NotImplementedError()
+ def apply_rls(
+ self,
+ database: Database,
+ 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.
@@ -529,6 +613,84 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
for eq in set_item.find_all(exp.EQ)
}
+ def apply_rls(
+ self,
+ database: Database,
+ catalog: str | None,
+ schema: str | None,
+ ) -> exp.Expression:
+ """
+ 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
+ """
+ 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
+
+ transformer = (
+ RLSAsSubquery(rules)
+ if (
+ database.db_engine_spec.allows_subqueries
+ and database.db_engine_spec.allows_alias_in_select
+ )
+ else RLSAsPredicate(rules)
+ )
+
+ return statement.transform(transformer)
+
+ 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)
+
class KQLSplitState(enum.Enum):
"""
diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py
index 3b80b7e01d..4e53baaf7c 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:
@@ -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