mirror of
https://github.com/apache/superset.git
synced 2026-07-27 00:52:33 +00:00
Increase parity
This commit is contained in:
Beto Dealmeida
@@ -76,6 +76,22 @@ _CLAUSE_ENDS = {
|
||||
TokenType.EXCEPT,
|
||||
}
|
||||
|
||||
_JOIN_STARTS = {
|
||||
TokenType.JOIN,
|
||||
TokenType.STRAIGHT_JOIN,
|
||||
TokenType.JOIN_MARKER,
|
||||
}
|
||||
|
||||
|
||||
def _splice_priority(text: str) -> int:
|
||||
"""
|
||||
Priority for applying splices at the same offset.
|
||||
|
||||
Insert full SQL fragments (WHERE/ON/predicates) before closing parens so
|
||||
wrapping splices like ``pred AND (existing)`` compose correctly.
|
||||
"""
|
||||
return 1 if text != ")" else 0
|
||||
|
||||
|
||||
def _before_whitespace(sql: str, offset: int) -> int:
|
||||
"""Back up past any whitespace immediately before *offset*."""
|
||||
@@ -161,83 +177,151 @@ def apply_rls_splice(
|
||||
|
||||
splices: list[tuple[int, str]] = []
|
||||
for scope in traverse_scope(tree):
|
||||
splice = _splice_for_scope(sql, tokens, scope, predicates, catalog, schema)
|
||||
if splice is not None:
|
||||
splices.append(splice)
|
||||
splices.extend(
|
||||
_splices_for_scope(
|
||||
sql,
|
||||
tokens,
|
||||
scope,
|
||||
predicates,
|
||||
catalog,
|
||||
schema,
|
||||
dialect,
|
||||
)
|
||||
)
|
||||
|
||||
# Apply splices in reverse offset order so earlier positions stay valid.
|
||||
splices.sort(key=lambda item: item[0], reverse=True)
|
||||
# For equal offsets, apply predicate/WHERE/ON inserts before ")" inserts.
|
||||
splices.sort(key=lambda item: (item[0], _splice_priority(item[1])), reverse=True)
|
||||
result = sql
|
||||
for offset, text in splices:
|
||||
result = result[:offset] + text + result[offset:]
|
||||
return result
|
||||
|
||||
|
||||
def _splice_for_scope(
|
||||
def _splices_for_scope(
|
||||
sql: str,
|
||||
tokens: list[Token],
|
||||
scope: object,
|
||||
predicates: dict[Table, list[str]],
|
||||
catalog: str | None,
|
||||
schema: str | None,
|
||||
) -> tuple[int, str] | None:
|
||||
dialect: str | None,
|
||||
) -> list[tuple[int, str]]:
|
||||
"""
|
||||
Compute the (offset, text) splice for a single SELECT scope, or ``None`` if
|
||||
the scope has no matching predicates or no usable anchor.
|
||||
Compute all splices for a single SELECT scope.
|
||||
|
||||
This mirrors ``RLSAsPredicateTransformer`` semantics:
|
||||
- predicates for FROM tables are applied to the SELECT WHERE clause as
|
||||
``pred AND (existing_where)``
|
||||
- predicates for JOIN tables are applied to each JOIN ON clause as
|
||||
``pred AND (existing_on)`` (or ``ON pred`` when ON is absent)
|
||||
"""
|
||||
scope_preds = _collect_scope_predicates(scope, predicates, catalog, schema)
|
||||
if not scope_preds:
|
||||
return None
|
||||
from_predicates: list[str] = []
|
||||
from_table_ends: list[int] = []
|
||||
join_splices: list[tuple[int, str]] = []
|
||||
|
||||
# Anchor: rightmost character position among the table-name identifiers
|
||||
# directly owned by this scope. Used to skip past tokens that belong to
|
||||
# earlier parts of the query (projections, JOIN ON clauses, etc.).
|
||||
table_ends = [
|
||||
ident._meta["end"]
|
||||
for source in scope.sources.values() # type: ignore[attr-defined]
|
||||
if isinstance(source, exp.Table)
|
||||
for ident in [source.find(exp.Identifier)]
|
||||
if ident and getattr(ident, "_meta", None)
|
||||
]
|
||||
if not table_ends:
|
||||
return None
|
||||
for source in scope.sources.values(): # type: ignore[attr-defined]
|
||||
source_type, table_end, pred_sql = _classify_source_predicate(
|
||||
source,
|
||||
predicates,
|
||||
catalog,
|
||||
schema,
|
||||
dialect,
|
||||
)
|
||||
if source_type == "none" or table_end is None or pred_sql is None:
|
||||
continue
|
||||
|
||||
has_where = scope.expression.args.get("where") is not None # type: ignore[attr-defined]
|
||||
pred_sql = " AND ".join(scope_preds)
|
||||
return _find_splice_point(sql, tokens, max(table_ends), has_where, pred_sql)
|
||||
if source_type == "from":
|
||||
from_predicates.append(pred_sql)
|
||||
from_table_ends.append(table_end)
|
||||
elif source_type == "join":
|
||||
join_splice = _find_join_splice(sql, tokens, table_end, pred_sql)
|
||||
if join_splice:
|
||||
join_splices.extend(join_splice)
|
||||
|
||||
if not from_predicates:
|
||||
return join_splices
|
||||
|
||||
combined_predicates = " AND ".join(dict.fromkeys(from_predicates))
|
||||
from_splice = _find_where_splice(
|
||||
sql,
|
||||
tokens,
|
||||
max(from_table_ends),
|
||||
combined_predicates,
|
||||
)
|
||||
return [*join_splices, *from_splice]
|
||||
|
||||
|
||||
def _collect_scope_predicates(
|
||||
scope: object,
|
||||
def _table_end(source: exp.Table) -> int | None:
|
||||
ident = source.find(exp.Identifier)
|
||||
if ident and getattr(ident, "_meta", None):
|
||||
return ident._meta["end"]
|
||||
return None
|
||||
|
||||
|
||||
def _classify_source_predicate(
|
||||
source: object,
|
||||
predicates: dict[Table, list[str]],
|
||||
catalog: str | None,
|
||||
schema: str | None,
|
||||
) -> list[str]:
|
||||
dialect: str | None,
|
||||
) -> tuple[str, int | None, str | None]:
|
||||
"""
|
||||
Collect the predicates that apply to direct Table sources in ``scope``,
|
||||
deduped while preserving order.
|
||||
Return source kind (from/join/none), table end offset, and predicate SQL.
|
||||
"""
|
||||
scope_preds: list[str] = []
|
||||
for source in scope.sources.values(): # type: ignore[attr-defined]
|
||||
if not isinstance(source, exp.Table):
|
||||
continue
|
||||
table = _table_from_node(source, catalog, schema)
|
||||
for predicate in predicates.get(table, []):
|
||||
if predicate and predicate not in scope_preds:
|
||||
scope_preds.append(predicate)
|
||||
return scope_preds
|
||||
if not isinstance(source, exp.Table):
|
||||
return ("none", None, None)
|
||||
|
||||
table = _table_from_node(source, catalog, schema)
|
||||
table_predicates = [
|
||||
_qualify_predicate(predicate, source, dialect)
|
||||
for predicate in predicates.get(table, [])
|
||||
if predicate
|
||||
]
|
||||
if not table_predicates:
|
||||
return ("none", None, None)
|
||||
|
||||
table_end = _table_end(source)
|
||||
if table_end is None:
|
||||
return ("none", None, None)
|
||||
|
||||
pred_sql = " AND ".join(dict.fromkeys(table_predicates))
|
||||
if isinstance(source.parent, exp.From):
|
||||
return ("from", table_end, pred_sql)
|
||||
if isinstance(source.parent, exp.Join):
|
||||
return ("join", table_end, pred_sql)
|
||||
return ("none", None, None)
|
||||
|
||||
|
||||
def _find_splice_point(
|
||||
sql: str,
|
||||
def _qualify_predicate(
|
||||
predicate: str,
|
||||
table_node: exp.Table,
|
||||
dialect: str | None,
|
||||
) -> str:
|
||||
"""
|
||||
Qualify predicate columns with the table alias/name, mirroring
|
||||
``RLSAsPredicateTransformer``.
|
||||
"""
|
||||
parsed = sqlglot.parse_one(predicate, dialect=dialect)
|
||||
table = table_node.alias_or_name
|
||||
table_expr = exp.to_identifier(table)
|
||||
for column in parsed.find_all(exp.Column):
|
||||
column.set("table", table_expr.copy())
|
||||
return parsed.sql(dialect=dialect)
|
||||
|
||||
|
||||
def _scan_until_scope_boundary(
|
||||
tokens: list[Token],
|
||||
anchor: int,
|
||||
has_where: bool,
|
||||
pred_sql: str,
|
||||
) -> tuple[int, str] | None:
|
||||
*,
|
||||
stop_at_join: bool,
|
||||
) -> tuple[str, int | None]:
|
||||
"""
|
||||
Scan tokens forward from ``anchor``, tracking paren depth, to find where to
|
||||
insert the RLS predicate for a single scope.
|
||||
Scan tokens forward from ``anchor`` until a clause/scope boundary.
|
||||
|
||||
Returns ``("where", index)`` when a WHERE token is found at depth 0,
|
||||
``("boundary", index)`` for a non-WHERE boundary token, and
|
||||
``("eof", None)`` when no boundary token is found.
|
||||
"""
|
||||
depth = 0
|
||||
for i, tok in enumerate(tokens):
|
||||
@@ -250,48 +334,147 @@ def _find_splice_point(
|
||||
|
||||
if tok.token_type == TokenType.R_PAREN:
|
||||
if depth == 0:
|
||||
# Closing paren of our subquery — insert just before it.
|
||||
offset = _before_trivia(sql, tok.start)
|
||||
text = f" AND {pred_sql}" if has_where else f" WHERE {pred_sql}"
|
||||
return (offset, text)
|
||||
return ("boundary", i)
|
||||
depth -= 1
|
||||
continue
|
||||
|
||||
if depth > 0:
|
||||
continue
|
||||
|
||||
if has_where and tok.token_type == TokenType.WHERE:
|
||||
return _find_after_where(sql, tokens, i, pred_sql)
|
||||
if tok.token_type == TokenType.WHERE:
|
||||
return ("where", i)
|
||||
|
||||
if not has_where and tok.token_type in _CLAUSE_ENDS:
|
||||
# Insert WHERE before this clause keyword.
|
||||
return (_before_trivia(sql, tok.start), f" WHERE {pred_sql}")
|
||||
if tok.token_type in _CLAUSE_ENDS or (
|
||||
stop_at_join and tok.token_type in _JOIN_STARTS
|
||||
):
|
||||
return ("boundary", i)
|
||||
|
||||
# No clause boundary found — append at end of SQL.
|
||||
text = f" AND {pred_sql}" if has_where else f" WHERE {pred_sql}"
|
||||
return (len(sql), text)
|
||||
return ("eof", None)
|
||||
|
||||
|
||||
def _find_after_where(
|
||||
def _find_condition_end(
|
||||
sql: str,
|
||||
tokens: list[Token],
|
||||
where_index: int,
|
||||
pred_sql: str,
|
||||
) -> tuple[int, str] | None:
|
||||
start_index: int,
|
||||
*,
|
||||
stop_at_join: bool,
|
||||
) -> int:
|
||||
"""
|
||||
Given the index of a ``WHERE`` token in ``tokens``, find the offset just
|
||||
after the WHERE clause body where ``AND <pred>`` should be inserted.
|
||||
Find the end offset for a WHERE/ON condition body.
|
||||
"""
|
||||
depth = 0
|
||||
prev_end = tokens[where_index].end
|
||||
for tok in tokens[where_index + 1 :]:
|
||||
prev_end = tokens[start_index].end
|
||||
for tok in tokens[start_index + 1 :]:
|
||||
if tok.token_type == TokenType.L_PAREN:
|
||||
depth += 1
|
||||
elif tok.token_type == TokenType.R_PAREN:
|
||||
if depth == 0:
|
||||
return (_before_trivia(sql, tok.start), f" AND {pred_sql}")
|
||||
return _before_trivia(sql, tok.start)
|
||||
depth -= 1
|
||||
elif depth == 0 and tok.token_type in _CLAUSE_ENDS:
|
||||
return (_before_trivia(sql, tok.start), f" AND {pred_sql}")
|
||||
elif depth == 0 and (
|
||||
(stop_at_join and tok.token_type == TokenType.WHERE)
|
||||
or tok.token_type in _CLAUSE_ENDS
|
||||
or (stop_at_join and tok.token_type in _JOIN_STARTS)
|
||||
):
|
||||
return _before_trivia(sql, tok.start)
|
||||
prev_end = tok.end
|
||||
return (prev_end + 1, f" AND {pred_sql}")
|
||||
return prev_end + 1
|
||||
|
||||
|
||||
def _find_where_splice(
|
||||
sql: str,
|
||||
tokens: list[Token],
|
||||
anchor: int,
|
||||
pred_sql: str,
|
||||
) -> list[tuple[int, str]]:
|
||||
"""
|
||||
Build splices for adding predicate semantics to the SELECT WHERE clause:
|
||||
``pred`` when absent, ``pred AND (existing)`` when present.
|
||||
"""
|
||||
kind, idx = _scan_until_scope_boundary(tokens, anchor, stop_at_join=False)
|
||||
if kind == "where" and idx is not None:
|
||||
if idx + 1 >= len(tokens):
|
||||
return [(tokens[idx].end + 1, f" {pred_sql}")]
|
||||
body_start = tokens[idx + 1].start
|
||||
body_end = _find_condition_end(sql, tokens, idx, stop_at_join=False)
|
||||
return [
|
||||
(body_start, f"{pred_sql} AND ("),
|
||||
(body_end, ")"),
|
||||
]
|
||||
|
||||
if kind == "boundary" and idx is not None:
|
||||
return [(_before_trivia(sql, tokens[idx].start), f" WHERE {pred_sql}")]
|
||||
|
||||
return [(len(sql), f" WHERE {pred_sql}")]
|
||||
|
||||
|
||||
def _find_join_splice(
|
||||
sql: str,
|
||||
tokens: list[Token],
|
||||
anchor: int,
|
||||
pred_sql: str,
|
||||
) -> list[tuple[int, str]]:
|
||||
"""
|
||||
Build splices for adding predicate semantics to a JOIN clause:
|
||||
``ON pred`` when ON absent, ``ON pred AND (existing_on)`` when present.
|
||||
"""
|
||||
on_index, boundary_index = _scan_join_clause(tokens, anchor)
|
||||
|
||||
if on_index is not None:
|
||||
if on_index + 1 >= len(tokens):
|
||||
return [(tokens[on_index].end + 1, f" {pred_sql}")]
|
||||
body_start = tokens[on_index + 1].start
|
||||
body_end = _find_condition_end(sql, tokens, on_index, stop_at_join=True)
|
||||
return [
|
||||
(body_start, f"{pred_sql} AND ("),
|
||||
(body_end, ")"),
|
||||
]
|
||||
|
||||
if boundary_index is not None:
|
||||
return [(_before_trivia(sql, tokens[boundary_index].start), f" ON {pred_sql}")]
|
||||
|
||||
return [(len(sql), f" ON {pred_sql}")]
|
||||
|
||||
|
||||
def _scan_join_clause(
|
||||
tokens: list[Token],
|
||||
anchor: int,
|
||||
) -> tuple[int | None, int | None]:
|
||||
"""
|
||||
Find ON and boundary token indexes for a JOIN segment.
|
||||
"""
|
||||
depth = 0
|
||||
on_index: int | None = None
|
||||
boundary_index: int | None = None
|
||||
|
||||
for i, tok in enumerate(tokens):
|
||||
if tok.start <= anchor:
|
||||
continue
|
||||
|
||||
if tok.token_type == TokenType.L_PAREN:
|
||||
depth += 1
|
||||
continue
|
||||
|
||||
if tok.token_type == TokenType.R_PAREN:
|
||||
if depth == 0:
|
||||
boundary_index = i
|
||||
break
|
||||
depth -= 1
|
||||
continue
|
||||
|
||||
if depth > 0:
|
||||
continue
|
||||
|
||||
if tok.token_type == TokenType.ON and on_index is None:
|
||||
on_index = i
|
||||
continue
|
||||
|
||||
if tok.token_type == TokenType.WHERE:
|
||||
boundary_index = i
|
||||
break
|
||||
|
||||
if tok.token_type in _JOIN_STARTS or tok.token_type in _CLAUSE_ENDS:
|
||||
boundary_index = i
|
||||
break
|
||||
|
||||
return on_index, boundary_index
|
||||
|
||||
@@ -3102,89 +3102,142 @@ def test_rls_predicate_transformer(
|
||||
assert statement.format() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sql, rules, expected",
|
||||
[
|
||||
(
|
||||
"SELECT t.foo FROM some_table AS t",
|
||||
{Table("some_table", "schema1", "catalog1"): "t.id = 42"},
|
||||
"SELECT t.foo FROM some_table AS t WHERE t.id = 42",
|
||||
),
|
||||
(
|
||||
"SELECT t.foo FROM some_table AS t WHERE bar = 'baz' OR foo = 'qux'",
|
||||
{Table("some_table", "schema1", "catalog1"): "t.id = 42"},
|
||||
"SELECT t.foo FROM some_table AS t WHERE t.id = 42 "
|
||||
"AND (bar = 'baz' OR foo = 'qux')",
|
||||
),
|
||||
(
|
||||
"SELECT * FROM table JOIN other_table ON table.id = other_table.id",
|
||||
{Table("other_table", "schema1", "catalog1"): "other_table.id = 42"},
|
||||
"SELECT * FROM table JOIN other_table ON other_table.id = 42 "
|
||||
"AND (table.id = other_table.id)",
|
||||
),
|
||||
(
|
||||
"SELECT * FROM table JOIN other_table",
|
||||
{Table("other_table", "schema1", "catalog1"): "other_table.id = 42"},
|
||||
"SELECT * FROM table JOIN other_table ON other_table.id = 42",
|
||||
),
|
||||
(
|
||||
"SELECT * FROM table JOIN other_table ON table.id = other_table.id "
|
||||
"WHERE 1=1",
|
||||
{Table("other_table", "schema1", "catalog1"): "other_table.id = 42"},
|
||||
"SELECT * FROM table JOIN other_table ON other_table.id = 42 "
|
||||
"AND (table.id = other_table.id) WHERE 1=1",
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_rls_predicate_splice_semantics_match_predicate(
|
||||
sql: str,
|
||||
rules: dict[Table, str],
|
||||
expected: str,
|
||||
) -> None:
|
||||
"""
|
||||
Splice mode should preserve predicate-mode semantics for boolean grouping
|
||||
and JOIN-vs-WHERE placement.
|
||||
"""
|
||||
statement = SQLStatement(sql)
|
||||
statement.apply_rls(
|
||||
"catalog1",
|
||||
"schema1",
|
||||
{k: [v] for k, v in rules.items()},
|
||||
RLSMethod.AS_PREDICATE_SPLICE,
|
||||
)
|
||||
assert statement.format() == expected
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"sql, rules, expected",
|
||||
[
|
||||
# Simple — no WHERE clause to extend.
|
||||
(
|
||||
"SELECT LAST_DAY(d) FROM some_table",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE tenant_id = 42",
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE some_table.tenant_id = 42",
|
||||
),
|
||||
# Append to an existing WHERE clause.
|
||||
(
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE status = 'open'",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE status = 'open' "
|
||||
"AND tenant_id = 42",
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table "
|
||||
"WHERE some_table.tenant_id = 42 AND (status = 'open')",
|
||||
),
|
||||
# WHERE precedes GROUP BY: predicate goes before GROUP BY.
|
||||
(
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE status = 'open' GROUP BY d",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE status = 'open' "
|
||||
"AND tenant_id = 42 GROUP BY d",
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table "
|
||||
"WHERE some_table.tenant_id = 42 AND (status = 'open') GROUP BY d",
|
||||
),
|
||||
# No WHERE, but GROUP BY and ORDER BY are present.
|
||||
(
|
||||
"SELECT LAST_DAY(d) FROM some_table GROUP BY d ORDER BY d",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE tenant_id = 42 "
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE some_table.tenant_id = 42 "
|
||||
"GROUP BY d ORDER BY d",
|
||||
),
|
||||
# JOIN — predicate scoped to one of the tables.
|
||||
(
|
||||
"SELECT o.id FROM some_table o JOIN locations l ON o.loc_id = l.id",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "o.tenant_id = 42"},
|
||||
"SELECT o.id FROM some_table o JOIN locations l "
|
||||
"ON o.loc_id = l.id WHERE tenant_id = 42",
|
||||
"ON o.loc_id = l.id WHERE o.tenant_id = 42",
|
||||
),
|
||||
# JOIN — different predicate per table, both spliced into one WHERE.
|
||||
(
|
||||
"SELECT * FROM some_table JOIN events ON some_table.id = events.order_id",
|
||||
{
|
||||
Table("some_table", "schema1", "catalog1"): "tenant_id = 42",
|
||||
Table("events", "schema1", "catalog1"): "user_id = 99",
|
||||
Table("events", "schema1", "catalog1"): "events.user_id = 99",
|
||||
Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42",
|
||||
},
|
||||
"SELECT * FROM some_table JOIN events "
|
||||
"ON some_table.id = events.order_id "
|
||||
"WHERE tenant_id = 42 AND user_id = 99",
|
||||
"ON events.user_id = 99 AND (some_table.id = events.order_id) "
|
||||
"WHERE some_table.tenant_id = 42",
|
||||
),
|
||||
# Subquery in FROM — splice into the inner SELECT.
|
||||
(
|
||||
"SELECT x FROM (SELECT LAST_DAY(d) AS x FROM some_table) sub",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT x FROM (SELECT LAST_DAY(d) AS x FROM some_table "
|
||||
"WHERE tenant_id = 42) sub",
|
||||
"WHERE some_table.tenant_id = 42) sub",
|
||||
),
|
||||
# CTE — splice into the CTE body.
|
||||
(
|
||||
"WITH cte AS (SELECT LAST_DAY(d) FROM some_table) SELECT * FROM cte",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"WITH cte AS (SELECT LAST_DAY(d) FROM some_table "
|
||||
"WHERE tenant_id = 42) SELECT * FROM cte",
|
||||
"WHERE some_table.tenant_id = 42) SELECT * FROM cte",
|
||||
),
|
||||
# Dialect-specific function (LAST_DAY) preserved verbatim.
|
||||
(
|
||||
"SELECT id, LAST_DAY(created_at) FROM some_table WHERE region = 'US'",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT id, LAST_DAY(created_at) FROM some_table "
|
||||
"WHERE region = 'US' AND tenant_id = 42",
|
||||
"WHERE some_table.tenant_id = 42 AND (region = 'US')",
|
||||
),
|
||||
# Multiline + inline comment preserved exactly.
|
||||
(
|
||||
"SELECT LAST_DAY(created_at) -- last day of month\n"
|
||||
"FROM some_table\n"
|
||||
"WHERE region = 'US'",
|
||||
{Table("some_table", "schema1", "catalog1"): "tenant_id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "some_table.tenant_id = 42"},
|
||||
"SELECT LAST_DAY(created_at) -- last day of month\n"
|
||||
"FROM some_table\n"
|
||||
"WHERE region = 'US' AND tenant_id = 42",
|
||||
"WHERE some_table.tenant_id = 42 AND (region = 'US')",
|
||||
),
|
||||
# Schema-qualified table name (no default schema match) — no predicate.
|
||||
(
|
||||
"SELECT t.foo FROM schema2.some_table AS t",
|
||||
{Table("some_table", "schema1", "catalog1"): "id = 42"},
|
||||
{Table("some_table", "schema1", "catalog1"): "t.id = 42"},
|
||||
"SELECT t.foo FROM schema2.some_table AS t",
|
||||
),
|
||||
],
|
||||
@@ -3237,10 +3290,12 @@ def test_rls_predicate_splice_preserves_dialect_function() -> None:
|
||||
statement.apply_rls(
|
||||
None,
|
||||
None,
|
||||
{Table("some_table"): ["tenant_id = 42"]},
|
||||
{Table("some_table"): ["some_table.tenant_id = 42"]},
|
||||
RLSMethod.AS_PREDICATE_SPLICE,
|
||||
)
|
||||
assert "LAST_DAY(d)" in statement.format()
|
||||
assert statement.format() == (
|
||||
"SELECT LAST_DAY(d) FROM some_table WHERE some_table.tenant_id = 42"
|
||||
)
|
||||
|
||||
|
||||
def test_rls_predicate_splice_string_predicates_skip_parse() -> None:
|
||||
@@ -3253,11 +3308,11 @@ def test_rls_predicate_splice_string_predicates_skip_parse() -> None:
|
||||
statement.apply_rls(
|
||||
None,
|
||||
None,
|
||||
{Table("some_table"): ["tenant_id = 42 AND active"]},
|
||||
{Table("some_table"): ["some_table.tenant_id = 42 AND some_table.active"]},
|
||||
RLSMethod.AS_PREDICATE_SPLICE,
|
||||
)
|
||||
assert statement.format() == (
|
||||
"SELECT * FROM some_table WHERE tenant_id = 42 AND active"
|
||||
"SELECT * FROM some_table WHERE some_table.tenant_id = 42 AND some_table.active"
|
||||
)
|
||||
|
||||
|
||||
@@ -3266,11 +3321,13 @@ def test_rls_predicate_splice_string_predicates_skip_parse() -> None:
|
||||
[
|
||||
(
|
||||
"SELECT * FROM some_table -- hi\nGROUP BY id",
|
||||
"SELECT * FROM some_table WHERE tenant_id = 42 -- hi\nGROUP BY id",
|
||||
"SELECT * FROM some_table WHERE some_table.tenant_id = 42 "
|
||||
"-- hi\nGROUP BY id",
|
||||
),
|
||||
(
|
||||
"SELECT * FROM some_table /* inline */ GROUP BY id",
|
||||
"SELECT * FROM some_table WHERE tenant_id = 42 /* inline */ GROUP BY id",
|
||||
"SELECT * FROM some_table WHERE some_table.tenant_id = 42 "
|
||||
"/* inline */ GROUP BY id",
|
||||
),
|
||||
],
|
||||
)
|
||||
@@ -3283,7 +3340,7 @@ def test_rls_predicate_splice_inserts_before_comments(sql: str, expected: str) -
|
||||
statement.apply_rls(
|
||||
None,
|
||||
None,
|
||||
{Table("some_table"): ["tenant_id = 42"]},
|
||||
{Table("some_table"): ["some_table.tenant_id = 42"]},
|
||||
RLSMethod.AS_PREDICATE_SPLICE,
|
||||
)
|
||||
assert statement.format() == expected
|
||||
@@ -3296,13 +3353,14 @@ def test_rls_predicate_splice_inserts_before_comments(sql: str, expected: str) -
|
||||
"SELECT * FROM some_table QUALIFY row_number() OVER "
|
||||
"(PARTITION BY id ORDER BY ts DESC) = 1",
|
||||
"snowflake",
|
||||
"SELECT * FROM some_table WHERE tenant_id = 42 "
|
||||
"SELECT * FROM some_table WHERE some_table.tenant_id = 42 "
|
||||
"QUALIFY row_number() OVER (PARTITION BY id ORDER BY ts DESC) = 1",
|
||||
),
|
||||
(
|
||||
"SELECT sum(v) OVER () FROM some_table WINDOW w AS (PARTITION BY id)",
|
||||
"postgresql",
|
||||
"SELECT sum(v) OVER () FROM some_table WHERE tenant_id = 42 "
|
||||
"SELECT sum(v) OVER () FROM some_table "
|
||||
"WHERE some_table.tenant_id = 42 "
|
||||
"WINDOW w AS (PARTITION BY id)",
|
||||
),
|
||||
],
|
||||
@@ -3320,7 +3378,7 @@ def test_rls_predicate_splice_handles_additional_clause_boundaries(
|
||||
statement.apply_rls(
|
||||
None,
|
||||
None,
|
||||
{Table("some_table"): ["tenant_id = 42"]},
|
||||
{Table("some_table"): ["some_table.tenant_id = 42"]},
|
||||
RLSMethod.AS_PREDICATE_SPLICE,
|
||||
)
|
||||
assert statement.format() == expected
|
||||
@@ -3340,7 +3398,7 @@ def test_rls_predicate_splice_then_limit_keeps_rls() -> None:
|
||||
statement.set_limit_value(101, LimitMethod.FORCE_LIMIT)
|
||||
|
||||
formatted = statement.format()
|
||||
assert "tenant_id = 42" in formatted
|
||||
assert "some_table.tenant_id = 42" in formatted
|
||||
assert "LIMIT 101" in formatted
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user