def test_insert_rls(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,
) -> 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(token_list=statement,
database_id=1,
default_schema="my_schema")).strip() == expected.strip())
def test_add_table_name(rls: str, table: str, expected: str) -> None:
condition = sqlparse.parse(rls)[0]
add_table_name(condition, table)
assert str(condition) == expected