From 52597ad0357a6d3296849c62fe388cbf3c1f2edd Mon Sep 17 00:00:00 2001 From: Beto Dealmeida Date: Wed, 8 Jun 2022 18:08:05 -0700 Subject: [PATCH] feat: safer insert RLS --- superset/connectors/sqla/utils.py | 4 +- superset/sql_lab.py | 17 ++- superset/sql_parse.py | 99 +++++++++++++++++- tests/unit_tests/sql_lab_test.py | 6 +- tests/unit_tests/sql_parse_tests.py | 154 +++++++++++++++++++++++++++- 5 files changed, 269 insertions(+), 11 deletions(-) diff --git a/superset/connectors/sqla/utils.py b/superset/connectors/sqla/utils.py index 1786c5bf17169..288865b8ac6b3 100644 --- a/superset/connectors/sqla/utils.py +++ b/superset/connectors/sqla/utils.py @@ -44,7 +44,7 @@ ) from superset.models.core import Database from superset.result_set import SupersetResultSet -from superset.sql_parse import has_table_query, insert_rls, ParsedQuery +from superset.sql_parse import has_table_query, insert_rls_in_predicate, ParsedQuery from superset.superset_typing import ResultSetColumnType from superset.utils.memoized import memoized @@ -174,7 +174,7 @@ def validate_adhoc_subquery( level=ErrorLevel.ERROR, ) ) - statement = insert_rls(statement, database_id, default_schema) + statement = insert_rls_in_predicate(statement, database_id, default_schema) statements.append(statement) return ";\n".join(str(statement) for statement in statements) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 785d16327f7f2..94dcb6dd42a5e 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -47,7 +47,12 @@ from superset.models.core import Database from superset.models.sql_lab import Query from superset.result_set import SupersetResultSet -from superset.sql_parse import CtasMethod, insert_rls, ParsedQuery +from superset.sql_parse import ( + CtasMethod, + insert_rls_as_subquery, + insert_rls_in_predicate, + ParsedQuery, +) from superset.sqllab.limiting_factor import LimitingFactor from superset.utils.celery import session_scope from superset.utils.core import ( @@ -201,6 +206,16 @@ def execute_sql_statement( # pylint: disable=too-many-arguments,too-many-statem parsed_query = ParsedQuery(sql_statement) if is_feature_enabled("RLS_IN_SQLLAB"): + # There are two ways to insert RLS: either replacing the table with a subquery + # that has the RLS, or appending the RLS to the ``WHERE`` clause. The former is + # safer, but not supported in all databases. + insert_rls = ( + insert_rls_as_subquery + if database.db_engine_spec.allows_subqueries + and database.db_engine_spec.allows_alias_in_select + else insert_rls_in_predicate + ) + # Insert any applicable RLS predicates parsed_query = ParsedQuery( str( diff --git a/superset/sql_parse.py b/superset/sql_parse.py index b585810f785a2..4d7b71a33be25 100644 --- a/superset/sql_parse.py +++ b/superset/sql_parse.py @@ -42,6 +42,7 @@ Punctuation, String, Whitespace, + Wildcard, ) from sqlparse.utils import imt @@ -597,7 +598,7 @@ def get_rls_for_table( return rls -def insert_rls( +def insert_rls_as_subquery( token_list: TokenList, database_id: int, default_schema: Optional[str], @@ -605,6 +606,98 @@ def insert_rls( ) -> TokenList: """ Update a statement inplace applying any associated RLS predicates. + + The RLS predicate is applied as subquery replacing the original table: + + before: SELECT * FROM some_table WHERE 1=1 + after: SELECT * FROM ( + SELECT * FROM some_table WHERE some_table.id=42 + ) AS some_table + WHERE 1=1 + + This method is safer than ``insert_rls_in_predicate``, but doesn't work in all + databases. + """ + rls: Optional[TokenList] = None + state = InsertRLSState.SCANNING + for token in token_list.tokens: + + # Recurse into child token list + if isinstance(token, TokenList): + i = token_list.tokens.index(token) + token_list.tokens[i] = insert_rls_as_subquery( + token, database_id, default_schema + ) + + # Found a source keyword (FROM/JOIN) + if imt(token, m=[(Keyword, "FROM"), (Keyword, "JOIN")]): + state = InsertRLSState.SEEN_SOURCE + + # Found identifier/keyword after FROM/JOIN, test for table + elif state == InsertRLSState.SEEN_SOURCE and ( + isinstance(token, Identifier) or token.ttype == Keyword + ): + rls = get_rls_for_table(token, database_id, default_schema, username) + if rls: + # replace table with subquery + subquery_alias = ( + token.tokens[-1].value + if isinstance(token, Identifier) + else token.value + ) + i = token_list.tokens.index(token) + token_list.tokens[i] = Identifier( + [ + Parenthesis( + [ + Token(Punctuation, "("), + Token(DML, "SELECT"), + Token(Whitespace, " "), + Token(Wildcard, "*"), + Token(Whitespace, " "), + Token(Keyword, "FROM"), + Token(Whitespace, " "), + token, + Token(Whitespace, " "), + Where( + [ + Token(Keyword, "WHERE"), + Token(Whitespace, " "), + rls, + ] + ), + Token(Punctuation, ")"), + ] + ), + Token(Whitespace, " "), + Token(Keyword, "AS"), + Token(Whitespace, " "), + Identifier([Token(Name, subquery_alias)]), + ] + ) + state = InsertRLSState.SCANNING + + # Found nothing, leaving source + elif state == InsertRLSState.SEEN_SOURCE and token.ttype != Whitespace: + state = InsertRLSState.SCANNING + + return token_list + + +def insert_rls_in_predicate( + token_list: TokenList, + database_id: int, + default_schema: Optional[str], + username: Optional[str] = None, +) -> TokenList: + """ + Update a statement inplace applying any associated RLS predicates. + + The RLS predicate is ``AND``ed to any existing predicates: + + before: SELECT * FROM some_table WHERE 1=1 + after: SELECT * FROM some_table WHERE ( 1=1) AND some_table.id=42 + """ rls: Optional[TokenList] = None state = InsertRLSState.SCANNING @@ -613,7 +706,9 @@ def insert_rls( # Recurse into child token list if isinstance(token, TokenList): i = token_list.tokens.index(token) - token_list.tokens[i] = insert_rls(token, database_id, default_schema) + token_list.tokens[i] = insert_rls_in_predicate( + token, database_id, default_schema + ) # Found a source keyword (FROM/JOIN) if imt(token, m=[(Keyword, "FROM"), (Keyword, "JOIN")]): diff --git a/tests/unit_tests/sql_lab_test.py b/tests/unit_tests/sql_lab_test.py index 9950fb9fedda5..1af5c9a6c5923 100644 --- a/tests/unit_tests/sql_lab_test.py +++ b/tests/unit_tests/sql_lab_test.py @@ -88,7 +88,7 @@ def test_execute_sql_statement_with_rls( cursor = mocker.MagicMock() SupersetResultSet = mocker.patch("superset.sql_lab.SupersetResultSet") mocker.patch( - "superset.sql_lab.insert_rls", + "superset.sql_lab.insert_rls_in_predicate", return_value=sqlparse.parse("SELECT * FROM sales WHERE organization_id=42")[0], ) mocker.patch("superset.sql_lab.is_feature_enabled", return_value=True) @@ -115,13 +115,13 @@ def test_execute_sql_statement_with_rls( SupersetResultSet.assert_called_with([(42,)], cursor.description, db_engine_spec) -def test_sql_lab_insert_rls( +def test_sql_lab_insert_rls_in_predicate( mocker: MockerFixture, session: Session, app_context: None, ) -> None: """ - Integration test for `insert_rls`. + Integration test for `insert_rls_in_predicate`. """ from flask_appbuilder.security.sqla.models import Role, User diff --git a/tests/unit_tests/sql_parse_tests.py b/tests/unit_tests/sql_parse_tests.py index 1d2c788496af0..3abd9af76a6b4 100644 --- a/tests/unit_tests/sql_parse_tests.py +++ b/tests/unit_tests/sql_parse_tests.py @@ -32,7 +32,8 @@ extract_table_references, get_rls_for_table, has_table_query, - insert_rls, + insert_rls_as_subquery, + insert_rls_in_predicate, ParsedQuery, sanitize_clause, strip_comments_from_sql, @@ -1221,6 +1222,151 @@ def test_has_table_query(sql: str, expected: bool) -> None: assert has_table_query(statement) == expected +@pytest.mark.parametrize( + "sql,table,rls,expected", + [ + # Basic test + ( + "SELECT * FROM some_table WHERE 1=1", + "some_table", + "id=42", + ( + "SELECT * FROM (SELECT * FROM some_table WHERE some_table.id=42) " + "AS some_table WHERE 1=1" + ), + ), + # Here "table" is a reserved word; since sqlparse is too aggressive when + # characterizing reserved words we need to support them even when not quoted. + ( + "SELECT * FROM table WHERE 1=1", + "table", + "id=42", + "SELECT * FROM (SELECT * FROM table WHERE table.id=42) AS table WHERE 1=1", + ), + # RLS is only applied to queries reading from the associated table + ( + "SELECT * FROM table WHERE 1=1", + "other_table", + "id=42", + "SELECT * FROM table WHERE 1=1", + ), + ( + "SELECT * FROM other_table WHERE 1=1", + "table", + "id=42", + "SELECT * FROM other_table WHERE 1=1", + ), + # JOINs are supported + ( + "SELECT * FROM table JOIN other_table ON table.id = other_table.id", + "other_table", + "id=42", + ( + "SELECT * FROM table JOIN " + "(SELECT * FROM other_table WHERE other_table.id=42) AS other_table " + "ON table.id = other_table.id" + ), + ), + # Subqueries + ( + "SELECT * FROM (SELECT * FROM other_table)", + "other_table", + "id=42", + ( + "SELECT * FROM (SELECT * FROM (" + "SELECT * FROM other_table WHERE other_table.id=42" + ") AS other_table)" + ), + ), + # UNION + ( + "SELECT * FROM table UNION ALL SELECT * FROM other_table", + "table", + "id=42", + ( + "SELECT * FROM (SELECT * FROM table WHERE table.id=42) AS table " + "UNION ALL SELECT * FROM other_table" + ), + ), + ( + "SELECT * FROM table UNION ALL SELECT * FROM other_table", + "other_table", + "id=42", + ( + "SELECT * FROM table UNION ALL SELECT * FROM (" + "SELECT * FROM other_table WHERE other_table.id=42) AS other_table" + ), + ), + # When comparing fully qualified table names (eg, schema.table) to simple names + # (eg, table) we are also conservative, assuming the schema is the same, since + # we don't have information on the default schema. + ( + "SELECT * FROM schema.table_name", + "table_name", + "id=42", + ( + "SELECT * FROM (SELECT * FROM schema.table_name " + "WHERE table_name.id=42) AS table_name" + ), + ), + ( + "SELECT * FROM schema.table_name", + "schema.table_name", + "id=42", + ( + "SELECT * FROM (SELECT * FROM schema.table_name " + "WHERE schema.table_name.id=42) AS table_name" + ), + ), + ( + "SELECT * FROM table_name", + "schema.table_name", + "id=42", + ( + "SELECT * FROM (SELECT * FROM table_name WHERE " + "schema.table_name.id=42) AS table_name" + ), + ), + ], +) +def test_insert_rls_as_subquery( + mocker: MockerFixture, sql: str, table: str, rls: str, expected: str +) -> None: + """ + Insert into a statement a given RLS condition associated with a table. + """ + condition = sqlparse.parse(rls)[0] + add_table_name(condition, table) + + # pylint: disable=unused-argument + def get_rls_for_table( + candidate: Token, + database_id: int, + default_schema: str, + username: Optional[str] = None, + ) -> Optional[TokenList]: + """ + Return the RLS ``condition`` if ``candidate`` matches ``table``. + """ + # compare ignoring schema + for left, right in zip(str(candidate).split(".")[::-1], table.split(".")[::-1]): + if left != right: + return None + return condition + + mocker.patch("superset.sql_parse.get_rls_for_table", new=get_rls_for_table) + + statement = sqlparse.parse(sql)[0] + assert ( + str( + insert_rls_as_subquery( + token_list=statement, database_id=1, default_schema="my_schema" + ) + ).strip() + == expected.strip() + ) + + @pytest.mark.parametrize( "sql,table,rls,expected", [ @@ -1395,7 +1541,7 @@ def test_has_table_query(sql: str, expected: bool) -> None: ), ], ) -def test_insert_rls( +def test_insert_rls_in_predicate( mocker: MockerFixture, sql: str, table: str, rls: str, expected: str ) -> None: """ @@ -1425,7 +1571,9 @@ def get_rls_for_table( statement = sqlparse.parse(sql)[0] assert ( str( - insert_rls(token_list=statement, database_id=1, default_schema="my_schema") + insert_rls_in_predicate( + token_list=statement, database_id=1, default_schema="my_schema" + ) ).strip() == expected.strip() )