diff --git a/superset/sql/rls_splice.py b/superset/sql/rls_splice.py index 48576f7ac71..2a26a9920d9 100644 --- a/superset/sql/rls_splice.py +++ b/superset/sql/rls_splice.py @@ -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 `` 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 diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py index f5a9ad574dd..553f6cef3cb 100644 --- a/tests/unit_tests/sql/parse_tests.py +++ b/tests/unit_tests/sql/parse_tests.py @@ -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