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 2026/05/08 17:49:48 UTC
(superset) branch rls-splice created (now af2d3babecc)
This is an automated email from the ASF dual-hosted git repository.
betodealmeida pushed a change to branch rls-splice
in repository https://gitbox.apache.org/repos/asf/superset.git
at af2d3babecc Improvements
This branch includes the following new commits:
new 3fe2b2505f3 feat: new splice RLSMethod
new af2d3babecc Improvements
The 2 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(superset) 01/02: feat: new splice RLSMethod
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
betodealmeida pushed a commit to branch rls-splice
in repository https://gitbox.apache.org/repos/asf/superset.git
commit 3fe2b2505f3d35009d7f7958b33da78933b24f75
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Fri May 8 13:33:42 2026 -0400
feat: new splice RLSMethod
---
superset/db_engine_specs/base.py | 22 ++-
superset/sql/parse.py | 117 +++++++++++-
superset/sql/rls_splice.py | 263 ++++++++++++++++++++++++++
superset/utils/rls.py | 25 ++-
tests/unit_tests/db_engine_specs/test_base.py | 43 ++++-
tests/unit_tests/sql/parse_tests.py | 159 ++++++++++++++++
6 files changed, 615 insertions(+), 14 deletions(-)
diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 4d26ca8517a..78fc2e78afa 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -557,6 +557,14 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
# if True, database will be listed as option in the upload file form
supports_file_upload = True
+ # Optional override for the RLS method used by ``get_rls_method``. When set,
+ # the engine spec opts into a specific strategy regardless of the
+ # ``allows_subqueries`` / ``allows_alias_in_select`` defaults. Use
+ # ``RLSMethod.AS_PREDICATE_SPLICE`` for engines whose sqlglot dialect can
+ # parse but not faithfully regenerate the SQL — splice mode rewrites the
+ # original query string instead of round-tripping through the generator.
+ rls_method: RLSMethod | None = None
+
# Is the DB engine spec able to change the default schema? This requires implementing # noqa: E501
# a custom `adjust_engine_params` method.
supports_dynamic_schema = False
@@ -623,10 +631,18 @@ class BaseEngineSpec: # pylint: disable=too-many-public-methods
"""
Returns the RLS method to be used for this engine.
- 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.
+ There are three ways to insert RLS: replacing the table with a subquery
+ that has the RLS (safest, but not supported in all databases), appending
+ the RLS to the ``WHERE`` clause via AST transformation, or splicing the
+ RLS into the original SQL string (preserves dialect-specific syntax that
+ the sqlglot generator would otherwise transpile).
+
+ Engine specs can opt into a specific strategy by setting the class-level
+ ``rls_method`` attribute; otherwise the choice falls back to subquery
+ when supported, and predicate otherwise.
"""
+ if cls.rls_method is not None:
+ return cls.rls_method
return (
RLSMethod.AS_SUBQUERY
if cls.allows_subqueries and cls.allows_alias_in_select
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
index bb3ef5e1c4b..a45943ea463 100644
--- a/superset/sql/parse.py
+++ b/superset/sql/parse.py
@@ -142,6 +142,7 @@ class RLSMethod(enum.Enum):
AS_PREDICATE = enum.auto()
AS_SUBQUERY = enum.auto()
+ AS_PREDICATE_SPLICE = enum.auto()
class RLSTransformer:
@@ -355,6 +356,7 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
statement: str | None = None,
engine: str = "base",
ast: InternalRepresentation | None = None,
+ source: str | None = None,
):
if ast:
self._parsed = ast
@@ -365,6 +367,16 @@ class BaseSQLStatement(Generic[InternalRepresentation]):
self.engine = engine
self.tables = self._extract_tables_from_statement(self._parsed, self.engine)
+ # Original SQL substring for this statement, when known. Used by the
+ # splice-mode RLS path which rewrites this string instead of regenerating
+ # SQL from the AST. ``None`` means the statement was constructed from an
+ # AST without an associated source string (splice mode falls back).
+ self._source_sql: str | None = source if source is not None else statement
+ # Verbatim SQL to return from ``format()``. Set by string-rewriting
+ # operations (e.g. splice-mode RLS) that produce a final SQL string and
+ # need to bypass the dialect generator. Cleared by AST-mutating methods
+ # since those invalidate this cached text.
+ self._raw_sql: str | None = None
@classmethod
def split_script(
@@ -559,9 +571,10 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
statement: str | None = None,
engine: str = "base",
ast: exp.Expression | None = None,
+ source: str | None = None,
):
self._dialect = SQLGLOT_DIALECTS.get(engine)
- super().__init__(statement, engine, ast)
+ super().__init__(statement, engine, ast, source)
@classmethod
def _parse(cls, script: str, engine: str) -> list[exp.Expression]:
@@ -626,10 +639,55 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
script: str,
engine: str,
) -> list[SQLStatement]:
+ asts = [ast for ast in cls._parse(script, engine) if ast]
+ sources = cls._split_source(script, engine, len(asts))
return [
- cls(ast=ast, engine=engine) for ast in cls._parse(script, engine) if ast
+ cls(ast=ast, engine=engine, source=source)
+ for ast, source in zip(asts, sources, strict=False)
]
+ @classmethod
+ def _split_source(
+ cls,
+ script: str,
+ engine: str,
+ expected_count: int,
+ ) -> list[str | None]:
+ """
+ Slice ``script`` into per-statement substrings using top-level semicolon
+ positions from the tokenizer. Returns a list of length ``expected_count``;
+ any entry is ``None`` if the slicing didn't yield a usable substring.
+
+ The returned substrings preserve the original byte content of the script
+ for each statement — necessary for splice-mode RLS, which rewrites the
+ original SQL rather than regenerating from the AST.
+ """
+ none_result: list[str | None] = [None] * expected_count
+ dialect = SQLGLOT_DIALECTS.get(engine)
+ try:
+ tokens = list(Dialect.get_or_raise(dialect).tokenize(script))
+ except sqlglot.errors.SqlglotError:
+ return none_result
+
+ # Top-level semicolon offsets (depth 0).
+ boundaries: list[int] = []
+ depth = 0
+ for tok in tokens:
+ if tok.token_type == sqlglot.tokens.TokenType.L_PAREN:
+ depth += 1
+ elif tok.token_type == sqlglot.tokens.TokenType.R_PAREN:
+ depth -= 1
+ elif tok.token_type == sqlglot.tokens.TokenType.SEMICOLON and depth == 0:
+ boundaries.append(tok.start)
+
+ starts = [0, *(b + 1 for b in boundaries)]
+ ends = [*boundaries, len(script)]
+ sources = [script[s:e].strip() for s, e in zip(starts, ends, strict=False)]
+ sources = [s for s in sources if s]
+ if len(sources) != expected_count:
+ return none_result
+ return list(sources)
+
@classmethod
def _parse_statement(
cls,
@@ -722,7 +780,13 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
def format(self, comments: bool = True) -> str:
"""
Pretty-format the SQL statement.
+
+ When a string-rewriting operation (e.g. splice-mode RLS) has cached a
+ verbatim result in ``_raw_sql``, return it as-is — the whole point of
+ those operations is to avoid the dialect generator round-trip.
"""
+ if self._raw_sql is not None:
+ return self._raw_sql
return Dialect.get_or_raise(self._dialect).generate(
self._parsed,
copy=True,
@@ -808,6 +872,8 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
"""
Modify the `LIMIT` or `TOP` value of the SQL statement inplace.
"""
+ # AST mutation invalidates any cached verbatim SQL (e.g. from splice).
+ self._raw_sql = None
if method == LimitMethod.FORCE_LIMIT:
self._parsed.args["limit"] = exp.Limit(
expression=exp.Literal(this=str(limit), is_string=False)
@@ -902,7 +968,7 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
self,
catalog: str | None,
schema: str | None,
- predicates: dict[Table, list[exp.Expression]],
+ predicates: dict[Table, list[exp.Expression]] | dict[Table, list[str]],
method: RLSMethod,
) -> None:
"""
@@ -910,11 +976,18 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
:param catalog: The default catalog for non-qualified table names
:param schema: The default schema for non-qualified table names
+ :param predicates: Mapping of fully qualified ``Table`` to predicates.
+ For ``AS_PREDICATE`` and ``AS_SUBQUERY`` the predicates are sqlglot
+ expressions. For ``AS_PREDICATE_SPLICE`` they are raw SQL strings.
:param method: The method to use for applying the rules.
"""
if not predicates:
return
+ if method == RLSMethod.AS_PREDICATE_SPLICE:
+ self._apply_rls_splice(catalog, schema, predicates)
+ return
+
transformers = {
RLSMethod.AS_PREDICATE: RLSAsPredicateTransformer,
RLSMethod.AS_SUBQUERY: RLSAsSubqueryTransformer,
@@ -925,6 +998,44 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
transformer = transformers[method](catalog, schema, predicates)
self._parsed = self._parsed.transform(transformer)
+ def _apply_rls_splice(
+ self,
+ catalog: str | None,
+ schema: str | None,
+ predicates: dict[Table, list[exp.Expression]] | dict[Table, list[str]],
+ ) -> None:
+ """
+ Apply RLS via text splicing on the original SQL.
+
+ Requires the source SQL substring to be available. Raises ``ValueError``
+ if it isn't — the caller must ensure the statement was constructed from
+ a source string (the standard ``SQLScript`` path does this).
+ """
+ from superset.sql.rls_splice import apply_rls_splice
+
+ if self._source_sql is None:
+ raise ValueError(
+ "Splice-mode RLS requires the source SQL string; "
+ "this SQLStatement was constructed without one."
+ )
+
+ # Splice operates on raw predicate strings; coerce expressions if needed.
+ string_predicates: dict[Table, list[str]] = {
+ table: [
+ pred if isinstance(pred, str) else pred.sql(dialect=self._dialect)
+ for pred in preds
+ ]
+ for table, preds in predicates.items()
+ }
+ spliced = apply_rls_splice(
+ self._source_sql,
+ catalog,
+ schema,
+ string_predicates,
+ dialect=self._dialect,
+ )
+ self._raw_sql = spliced
+
class KQLSplitState(enum.Enum):
"""
diff --git a/superset/sql/rls_splice.py b/superset/sql/rls_splice.py
new file mode 100644
index 00000000000..0166a13fe49
--- /dev/null
+++ b/superset/sql/rls_splice.py
@@ -0,0 +1,263 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+"""
+RLS predicate injection via text splicing.
+
+Instead of round-tripping through sqlglot's generator (which transpiles
+dialect-specific functions like ``LAST_DAY`` into something else), this approach:
+
+ 1. Parses the SQL with sqlglot — only to understand structure (scope tree).
+ 2. Uses sqlglot's tokenizer to get byte-accurate positions for every token
+ in the original SQL string.
+ 3. For each ``SELECT`` scope that references a table with an RLS predicate,
+ finds the exact byte offset to inject at — either the end of an existing
+ ``WHERE`` clause, or just before ``GROUP BY`` / ``ORDER BY`` / ``HAVING``
+ / ``LIMIT`` / the closing paren of a subquery.
+ 4. Splices the predicate text directly into the original string at that
+ offset — never calling ``.sql()``, so the generator never runs.
+
+Result: everything outside the splice points is the original SQL, byte for
+byte. Dialect-specific functions, comments, and formatting are all preserved
+exactly.
+
+Known limitations:
+ - SQL that fails to parse under the chosen dialect raises a ``ParseError``.
+ A thin dialect subclass is still required for parsing — but only for
+ parsing, not generation.
+ - Predicate strings are spliced in as raw SQL. They must come from a trusted
+ source (the RLS config), not user input.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import sqlglot
+from sqlglot import exp
+from sqlglot.optimizer.scope import traverse_scope
+from sqlglot.tokens import Token, TokenType
+
+if TYPE_CHECKING:
+ from superset.sql.parse import Table
+
+
+# Token types that end a WHERE clause / FROM section at the current paren depth,
+# indicating where a new predicate must be inserted just before.
+_CLAUSE_ENDS = {
+ TokenType.GROUP_BY,
+ TokenType.HAVING,
+ TokenType.ORDER_BY,
+ TokenType.LIMIT,
+ TokenType.FETCH,
+ TokenType.UNION,
+ TokenType.INTERSECT,
+ TokenType.EXCEPT,
+}
+
+
+def _before_whitespace(sql: str, offset: int) -> int:
+ """Back up past any whitespace immediately before *offset*."""
+ while offset > 0 and sql[offset - 1] in (" ", "\t", "\n", "\r"):
+ offset -= 1
+ return offset
+
+
+def _table_from_node(
+ node: exp.Table,
+ catalog: str | None,
+ schema: str | None,
+) -> Table:
+ """
+ Build a fully qualified ``Table`` from a sqlglot ``exp.Table`` node, defaulting
+ unqualified parts to the supplied catalog/schema.
+ """
+ # Imported lazily to avoid a circular import with ``superset.sql.parse``.
+ from superset.sql.parse import Table
+
+ return Table(
+ table=node.name,
+ schema=node.db if node.db else schema,
+ catalog=node.catalog if node.catalog else catalog,
+ )
+
+
+def apply_rls_splice(
+ sql: str,
+ catalog: str | None,
+ schema: str | None,
+ predicates: dict[Table, list[str]],
+ dialect: str | None = None,
+) -> str:
+ """
+ Inject RLS predicates into ``sql`` by splicing text at the right positions.
+
+ :param sql: The original SQL query. Returned unchanged except at splice points.
+ :param catalog: The default catalog for non-qualified table names.
+ :param schema: The default schema for non-qualified table names.
+ :param predicates: Mapping of ``Table`` to predicate SQL strings. Each entry
+ maps a fully qualified table to one or more raw predicate strings to
+ ``AND`` together when that table is referenced in a SELECT scope.
+ :param dialect: The sqlglot dialect used for *parsing only* — to understand
+ scope structure and locate token positions. The generator is never
+ called, so this does not affect output formatting.
+ :return: The query with RLS predicates injected into every relevant SELECT
+ scope.
+ """
+ if not predicates or not any(predicates.values()):
+ return sql
+
+ resolved_dialect = sqlglot.Dialect.get_or_raise(dialect)
+ tokens = list(resolved_dialect.tokenize(sql))
+ tree = sqlglot.parse_one(sql, dialect=dialect)
+
+ 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)
+
+ # Apply splices in reverse offset order so earlier positions stay valid.
+ splices.sort(key=lambda item: item[0], reverse=True)
+ result = sql
+ for offset, text in splices:
+ result = result[:offset] + text + result[offset:]
+ return result
+
+
+def _splice_for_scope(
+ sql: str,
+ tokens: list[Token],
+ scope: object,
+ predicates: dict[Table, list[str]],
+ catalog: str | None,
+ schema: str | None,
+) -> tuple[int, str] | None:
+ """
+ Compute the (offset, text) splice for a single SELECT scope, or ``None`` if
+ the scope has no matching predicates or no usable anchor.
+ """
+ scope_preds = _collect_scope_predicates(scope, predicates, catalog, schema)
+ if not scope_preds:
+ return None
+
+ # 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
+
+ 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)
+
+
+def _collect_scope_predicates(
+ scope: object,
+ predicates: dict[Table, list[str]],
+ catalog: str | None,
+ schema: str | None,
+) -> list[str]:
+ """
+ Collect the predicates that apply to direct Table sources in ``scope``,
+ deduped while preserving order.
+ """
+ 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
+
+
+def _find_splice_point(
+ sql: str,
+ tokens: list[Token],
+ anchor: int,
+ has_where: bool,
+ pred_sql: str,
+) -> tuple[int, str] | None:
+ """
+ Scan tokens forward from ``anchor``, tracking paren depth, to find where to
+ insert the RLS predicate for a single scope.
+ """
+ depth = 0
+ 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:
+ # Closing paren of our subquery — insert just before it.
+ offset = _before_whitespace(sql, tok.start)
+ text = f" AND {pred_sql}" if has_where else f" WHERE {pred_sql}"
+ return (offset, text)
+ 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 not has_where and tok.token_type in _CLAUSE_ENDS:
+ # Insert WHERE before this clause keyword.
+ return (_before_whitespace(sql, tok.start), f" WHERE {pred_sql}")
+
+ # 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)
+
+
+def _find_after_where(
+ sql: str,
+ tokens: list[Token],
+ where_index: int,
+ pred_sql: str,
+) -> tuple[int, str] | None:
+ """
+ Given the index of a ``WHERE`` token in ``tokens``, find the offset just
+ after the WHERE clause body where ``AND <pred>`` should be inserted.
+ """
+ depth = 0
+ prev_end = tokens[where_index].end
+ for tok in tokens[where_index + 1 :]:
+ if tok.token_type == TokenType.L_PAREN:
+ depth += 1
+ elif tok.token_type == TokenType.R_PAREN:
+ if depth == 0:
+ return (_before_whitespace(sql, tok.start), f" AND {pred_sql}")
+ depth -= 1
+ elif depth == 0 and tok.token_type in _CLAUSE_ENDS:
+ return (_before_whitespace(sql, tok.start), f" AND {pred_sql}")
+ prev_end = tok.end
+ return (prev_end + 1, f" AND {pred_sql}")
diff --git a/superset/utils/rls.py b/superset/utils/rls.py
index 7e6cdf2aee7..456b589e365 100644
--- a/superset/utils/rls.py
+++ b/superset/utils/rls.py
@@ -22,7 +22,7 @@ from typing import Any, TYPE_CHECKING
from sqlalchemy import and_, or_
from superset import db
-from superset.sql.parse import Table
+from superset.sql.parse import RLSMethod, Table
if TYPE_CHECKING:
from superset.models.core import Database
@@ -40,17 +40,23 @@ def apply_rls(
:returns: True if any RLS predicates were actually applied, False otherwise.
"""
- # 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.
+ # There are three ways to insert RLS:
+ # - replace the table with a subquery containing the RLS (safest, but not
+ # supported in all databases)
+ # - append the RLS to the ``WHERE`` clause via AST transformation
+ # - splice the RLS into the original SQL string (preserves dialect-specific
+ # syntax that the sqlglot generator would otherwise transpile)
method = database.db_engine_spec.get_rls_method()
- # collect all RLS predicates for all tables in the query
+ # In splice mode predicates stay as raw SQL strings and are inserted verbatim
+ # into the source query — re-parsing them would force a generator round-trip
+ # later and defeat the purpose.
+ use_splice = method == RLSMethod.AS_PREDICATE_SPLICE
predicates: dict[Table, list[Any]] = {}
for table in parsed_statement.tables:
table = table.qualify(catalog=catalog, schema=schema)
- predicates[table] = [
- parsed_statement.parse_predicate(predicate)
+ raw_predicates = [
+ predicate
for predicate in get_predicates_for_table(
table,
database,
@@ -58,6 +64,11 @@ def apply_rls(
)
if predicate
]
+ predicates[table] = (
+ raw_predicates
+ if use_splice
+ else [parsed_statement.parse_predicate(p) for p in raw_predicates]
+ )
has_predicates = any(predicates.values())
parsed_statement.apply_rls(catalog, schema, predicates, method)
diff --git a/tests/unit_tests/db_engine_specs/test_base.py b/tests/unit_tests/db_engine_specs/test_base.py
index 5eae41458de..1b58809b3b2 100644
--- a/tests/unit_tests/db_engine_specs/test_base.py
+++ b/tests/unit_tests/db_engine_specs/test_base.py
@@ -36,7 +36,7 @@ from sqlalchemy.sql import sqltypes
from superset.db_engine_specs.base import BaseEngineSpec, convert_inspector_columns
from superset.errors import ErrorLevel, SupersetError, SupersetErrorType
from superset.exceptions import OAuth2RedirectError
-from superset.sql.parse import Table
+from superset.sql.parse import RLSMethod, Table
from superset.superset_typing import (
OAuth2ClientConfig,
OAuth2State,
@@ -1283,3 +1283,44 @@ def test_start_oauth2_dance_falls_back_to_url_for(mocker: MockerFixture) -> None
error = exc_info.value.error
assert error.extra["redirect_uri"] == fallback_uri
+
+
+def test_get_rls_method_default_subquery() -> None:
+ """
+ By default, an engine that supports subqueries and aliases-in-select
+ uses the safer subquery RLS strategy.
+ """
+
+ class _Spec(BaseEngineSpec):
+ allows_subqueries = True
+ allows_alias_in_select = True
+
+ assert _Spec.get_rls_method() == RLSMethod.AS_SUBQUERY
+
+
+def test_get_rls_method_default_predicate_when_no_subqueries() -> None:
+ """
+ Engines without subquery / alias-in-select support fall back to the
+ AST predicate strategy.
+ """
+
+ class _Spec(BaseEngineSpec):
+ allows_subqueries = False
+ allows_alias_in_select = True
+
+ assert _Spec.get_rls_method() == RLSMethod.AS_PREDICATE
+
+
+def test_get_rls_method_class_attribute_override() -> None:
+ """
+ Setting ``rls_method`` on an engine spec opts the engine into a specific
+ strategy regardless of the subquery/alias defaults — used by engines whose
+ sqlglot dialect can parse but not faithfully regenerate SQL.
+ """
+
+ class _SpliceSpec(BaseEngineSpec):
+ allows_subqueries = True
+ allows_alias_in_select = True
+ rls_method = RLSMethod.AS_PREDICATE_SPLICE
+
+ assert _SpliceSpec.get_rls_method() == RLSMethod.AS_PREDICATE_SPLICE
diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py
index 78b00f4487d..7835840b04a 100644
--- a/tests/unit_tests/sql/parse_tests.py
+++ b/tests/unit_tests/sql/parse_tests.py
@@ -2548,6 +2548,165 @@ def test_rls_predicate_transformer(
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",
+ ),
+ # 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",
+ ),
+ # 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",
+ ),
+ # 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 "
+ "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"},
+ "SELECT o.id FROM some_table o JOIN locations l "
+ "ON o.loc_id = l.id WHERE 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",
+ },
+ "SELECT * FROM some_table JOIN events "
+ "ON some_table.id = events.order_id "
+ "WHERE tenant_id = 42 AND user_id = 99",
+ ),
+ # 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"},
+ "SELECT x FROM (SELECT LAST_DAY(d) AS x FROM some_table "
+ "WHERE 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"},
+ "WITH cte AS (SELECT LAST_DAY(d) FROM some_table "
+ "WHERE 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"},
+ "SELECT id, LAST_DAY(created_at) FROM some_table "
+ "WHERE region = 'US' AND tenant_id = 42",
+ ),
+ # 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"},
+ "SELECT LAST_DAY(created_at) -- last day of month\n"
+ "FROM some_table\n"
+ "WHERE region = 'US' AND tenant_id = 42",
+ ),
+ # 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"},
+ "SELECT t.foo FROM schema2.some_table AS t",
+ ),
+ ],
+)
+def test_rls_predicate_splice(
+ sql: str,
+ rules: dict[Table, str],
+ expected: str,
+) -> None:
+ """
+ Test the splice-mode RLS via ``RLSMethod.AS_PREDICATE_SPLICE``.
+
+ Splice mode rewrites the original SQL string instead of re-rendering the
+ AST through the dialect generator, so byte-level fidelity (including
+ dialect-specific functions, comments, and whitespace) is preserved.
+ """
+ statement = SQLStatement(sql)
+ statement.apply_rls(
+ "catalog1",
+ "schema1",
+ {k: [v] for k, v in rules.items()},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ assert statement.format() == expected
+
+
+def test_rls_predicate_splice_requires_source() -> None:
+ """
+ Splice mode requires the original SQL substring; constructing a statement
+ purely from an AST should make splice mode raise.
+ """
+ ast = parse_one("SELECT * FROM some_table")
+ statement = SQLStatement(ast=ast, engine="postgresql")
+ with pytest.raises(ValueError, match="Splice-mode RLS requires the source SQL"):
+ statement.apply_rls(
+ "catalog1",
+ "schema1",
+ {Table("some_table", "schema1", "catalog1"): ["id = 42"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+
+
+def test_rls_predicate_splice_preserves_dialect_function() -> None:
+ """
+ Splice mode must NOT round-trip through the sqlglot generator. ``LAST_DAY``
+ on the postgres dialect would otherwise be transpiled by the generator.
+ """
+ sql = "SELECT LAST_DAY(d) FROM some_table"
+ statement = SQLStatement(sql, engine="postgresql")
+ statement.apply_rls(
+ None,
+ None,
+ {Table("some_table"): ["tenant_id = 42"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ assert "LAST_DAY(d)" in statement.format()
+
+
+def test_rls_predicate_splice_string_predicates_skip_parse() -> None:
+ """
+ Splice mode accepts predicate strings directly — no ``parse_predicate`` is
+ needed at the call site.
+ """
+ sql = "SELECT * FROM some_table"
+ statement = SQLStatement(sql, engine="postgresql")
+ statement.apply_rls(
+ None,
+ None,
+ {Table("some_table"): ["tenant_id = 42 AND active"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ assert statement.format() == (
+ "SELECT * FROM some_table WHERE tenant_id = 42 AND active"
+ )
+
+
@pytest.mark.parametrize(
"sql, table, expected",
[
(superset) 02/02: Improvements
Posted by be...@apache.org.
This is an automated email from the ASF dual-hosted git repository.
betodealmeida pushed a commit to branch rls-splice
in repository https://gitbox.apache.org/repos/asf/superset.git
commit af2d3babecc150dfebe39ae72213675bf99ddb1f
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Fri May 8 13:49:37 2026 -0400
Improvements
---
superset/sql/parse.py | 7 ++-
superset/sql/rls_splice.py | 42 ++++++++++++++--
tests/unit_tests/sql/parse_tests.py | 95 +++++++++++++++++++++++++++++++++++++
3 files changed, 139 insertions(+), 5 deletions(-)
diff --git a/superset/sql/parse.py b/superset/sql/parse.py
index a45943ea463..d944dce5a2f 100644
--- a/superset/sql/parse.py
+++ b/superset/sql/parse.py
@@ -873,7 +873,12 @@ class SQLStatement(BaseSQLStatement[exp.Expression]):
Modify the `LIMIT` or `TOP` value of the SQL statement inplace.
"""
# AST mutation invalidates any cached verbatim SQL (e.g. from splice).
- self._raw_sql = None
+ # If we already have a rewritten SQL string, re-parse it first so further
+ # AST mutations (like LIMIT injection) preserve prior text-based rewrites.
+ if self._raw_sql is not None:
+ self._parsed = self._parse_statement(self._raw_sql, self.engine)
+ self._source_sql = self._raw_sql
+ self._raw_sql = None
if method == LimitMethod.FORCE_LIMIT:
self._parsed.args["limit"] = exp.Limit(
expression=exp.Literal(this=str(limit), is_string=False)
diff --git a/superset/sql/rls_splice.py b/superset/sql/rls_splice.py
index 0166a13fe49..48576f7ac71 100644
--- a/superset/sql/rls_splice.py
+++ b/superset/sql/rls_splice.py
@@ -62,8 +62,15 @@ _CLAUSE_ENDS = {
TokenType.GROUP_BY,
TokenType.HAVING,
TokenType.ORDER_BY,
+ TokenType.WINDOW,
+ TokenType.QUALIFY,
TokenType.LIMIT,
TokenType.FETCH,
+ TokenType.CLUSTER_BY,
+ TokenType.DISTRIBUTE_BY,
+ TokenType.SORT_BY,
+ TokenType.CONNECT_BY,
+ TokenType.START_WITH,
TokenType.UNION,
TokenType.INTERSECT,
TokenType.EXCEPT,
@@ -77,6 +84,33 @@ def _before_whitespace(sql: str, offset: int) -> int:
return offset
+def _before_trivia(sql: str, offset: int) -> int:
+ """
+ Back up past whitespace and adjacent comments immediately before *offset*.
+
+ This ensures insertion points land before inline/block comments that appear
+ between `FROM`/`WHERE` and the next clause keyword.
+ """
+ while True:
+ offset = _before_whitespace(sql, offset)
+
+ # Inline comment ending at offset, eg: "... -- comment\nGROUP BY".
+ line_start = sql.rfind("\n", 0, offset) + 1
+ inline_comment_start = sql.rfind("--", line_start, offset)
+ if inline_comment_start != -1:
+ offset = inline_comment_start
+ continue
+
+ # Block comment ending at offset, eg: "... /* comment */GROUP BY".
+ if offset >= 2 and sql[offset - 2 : offset] == "*/":
+ block_comment_start = sql.rfind("/*", 0, offset - 2)
+ if block_comment_start != -1:
+ offset = block_comment_start
+ continue
+
+ return offset
+
+
def _table_from_node(
node: exp.Table,
catalog: str | None,
@@ -217,7 +251,7 @@ 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_whitespace(sql, tok.start)
+ offset = _before_trivia(sql, tok.start)
text = f" AND {pred_sql}" if has_where else f" WHERE {pred_sql}"
return (offset, text)
depth -= 1
@@ -231,7 +265,7 @@ def _find_splice_point(
if not has_where and tok.token_type in _CLAUSE_ENDS:
# Insert WHERE before this clause keyword.
- return (_before_whitespace(sql, tok.start), f" WHERE {pred_sql}")
+ return (_before_trivia(sql, tok.start), f" WHERE {pred_sql}")
# No clause boundary found — append at end of SQL.
text = f" AND {pred_sql}" if has_where else f" WHERE {pred_sql}"
@@ -255,9 +289,9 @@ def _find_after_where(
depth += 1
elif tok.token_type == TokenType.R_PAREN:
if depth == 0:
- return (_before_whitespace(sql, tok.start), f" AND {pred_sql}")
+ return (_before_trivia(sql, tok.start), f" AND {pred_sql}")
depth -= 1
elif depth == 0 and tok.token_type in _CLAUSE_ENDS:
- return (_before_whitespace(sql, tok.start), f" AND {pred_sql}")
+ return (_before_trivia(sql, tok.start), f" AND {pred_sql}")
prev_end = tok.end
return (prev_end + 1, f" AND {pred_sql}")
diff --git a/tests/unit_tests/sql/parse_tests.py b/tests/unit_tests/sql/parse_tests.py
index 7835840b04a..79a9cec6953 100644
--- a/tests/unit_tests/sql/parse_tests.py
+++ b/tests/unit_tests/sql/parse_tests.py
@@ -1704,6 +1704,21 @@ def test_set_limit_value(
assert statement.format() == expected
+def test_set_limit_value_after_splice_reparses_from_raw_sql() -> None:
+ """
+ When a statement has cached verbatim SQL from splice-mode rewrites, setting
+ limit should reparse that SQL before mutating the AST.
+ """
+ statement = SQLStatement("SELECT * FROM some_table", "postgresql")
+ statement._raw_sql = "SELECT * FROM some_table WHERE tenant_id = 42"
+
+ statement.set_limit_value(10, LimitMethod.FORCE_LIMIT)
+ formatted = statement.format()
+
+ assert "tenant_id = 42" in formatted
+ assert "LIMIT 10" in formatted
+
+
@pytest.mark.parametrize(
"kql, limit, expected",
[
@@ -2707,6 +2722,86 @@ def test_rls_predicate_splice_string_predicates_skip_parse() -> None:
)
+@pytest.mark.parametrize(
+ "sql, expected",
+ [
+ (
+ "SELECT * FROM some_table -- hi\nGROUP BY id",
+ "SELECT * FROM some_table WHERE 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",
+ ),
+ ],
+)
+def test_rls_predicate_splice_inserts_before_comments(sql: str, expected: str) -> None:
+ """
+ Splice mode should insert predicates before comments that precede the next
+ clause boundary, so comments do not swallow the injected SQL.
+ """
+ statement = SQLStatement(sql, engine="postgresql")
+ statement.apply_rls(
+ None,
+ None,
+ {Table("some_table"): ["tenant_id = 42"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ assert statement.format() == expected
+
+
+@pytest.mark.parametrize(
+ "sql, engine, expected",
+ [
+ (
+ "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 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 WINDOW w AS (PARTITION BY id)",
+ ),
+ ],
+)
+def test_rls_predicate_splice_handles_additional_clause_boundaries(
+ sql: str,
+ engine: str,
+ expected: str,
+) -> None:
+ """
+ Splice mode should insert WHERE before clause types that can legally follow
+ FROM/WHERE (for example QUALIFY and WINDOW).
+ """
+ statement = SQLStatement(sql, engine=engine)
+ statement.apply_rls(
+ None,
+ None,
+ {Table("some_table"): ["tenant_id = 42"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ assert statement.format() == expected
+
+
+def test_rls_predicate_splice_then_limit_keeps_rls() -> None:
+ """
+ LIMIT rewrites after splice-mode RLS should retain injected predicates.
+ """
+ statement = SQLStatement("SELECT * FROM some_table", engine="postgresql")
+ statement.apply_rls(
+ None,
+ None,
+ {Table("some_table"): ["tenant_id = 42"]},
+ RLSMethod.AS_PREDICATE_SPLICE,
+ )
+ statement.set_limit_value(101, LimitMethod.FORCE_LIMIT)
+
+ formatted = statement.format()
+ assert "tenant_id = 42" in formatted
+ assert "LIMIT 101" in formatted
+
+
@pytest.mark.parametrize(
"sql, table, expected",
[