def setUp(self):
session = db.session
# Create roles
security_manager.add_role(self.NAME_AB_ROLE)
security_manager.add_role(self.NAME_Q_ROLE)
gamma_user = security_manager.find_user(username="******")
gamma_user.roles.append(security_manager.find_role(self.NAME_AB_ROLE))
gamma_user.roles.append(security_manager.find_role(self.NAME_Q_ROLE))
self.create_user_with_roles("NoRlsRoleUser", ["Gamma"])
session.commit()
# Create regular RowLevelSecurityFilter (energy_usage, unicode_test)
self.rls_entry1 = RowLevelSecurityFilter()
self.rls_entry1.tables.extend(
session.query(SqlaTable).filter(
SqlaTable.table_name.in_(["energy_usage",
"unicode_test"])).all())
self.rls_entry1.filter_type = "Regular"
self.rls_entry1.clause = "value > {{ cache_key_wrapper(1) }}"
self.rls_entry1.group_key = None
self.rls_entry1.roles.append(security_manager.find_role("Gamma"))
self.rls_entry1.roles.append(security_manager.find_role("Alpha"))
db.session.add(self.rls_entry1)
# Create regular RowLevelSecurityFilter (birth_names name starts with A or B)
self.rls_entry2 = RowLevelSecurityFilter()
self.rls_entry2.tables.extend(
session.query(SqlaTable).filter(
SqlaTable.table_name.in_(["birth_names"])).all())
self.rls_entry2.filter_type = "Regular"
self.rls_entry2.clause = "name like 'A%' or name like 'B%'"
self.rls_entry2.group_key = "name"
self.rls_entry2.roles.append(security_manager.find_role("NameAB"))
db.session.add(self.rls_entry2)
# Create Regular RowLevelSecurityFilter (birth_names name starts with Q)
self.rls_entry3 = RowLevelSecurityFilter()
self.rls_entry3.tables.extend(
session.query(SqlaTable).filter(
SqlaTable.table_name.in_(["birth_names"])).all())
self.rls_entry3.filter_type = "Regular"
self.rls_entry3.clause = "name like 'Q%'"
self.rls_entry3.group_key = "name"
self.rls_entry3.roles.append(security_manager.find_role("NameQ"))
db.session.add(self.rls_entry3)
# Create Base RowLevelSecurityFilter (birth_names boys)
self.rls_entry4 = RowLevelSecurityFilter()
self.rls_entry4.tables.extend(
session.query(SqlaTable).filter(
SqlaTable.table_name.in_(["birth_names"])).all())
self.rls_entry4.filter_type = "Base"
self.rls_entry4.clause = "gender = 'boy'"
self.rls_entry4.group_key = "gender"
self.rls_entry4.roles.append(security_manager.find_role("Admin"))
db.session.add(self.rls_entry4)
db.session.commit()
def setUp(self):
session = db.session
# Create the RowLevelSecurityFilter
self.rls_entry = RowLevelSecurityFilter()
self.rls_entry.table = (session.query(SqlaTable).filter_by(
table_name="birth_names").first())
self.rls_entry.clause = "gender = 'male'"
self.rls_entry.roles.append(
security_manager.find_role("Gamma")
) # db.session.query(Role).filter_by(name="Gamma").first())
db.session.add(self.rls_entry)
db.session.commit()
def setUp(self):
# Create the RowLevelSecurityFilter
self.rls_entry = RowLevelSecurityFilter()
self.rls_entry.tables.extend(
db.session.query(SqlaTable).filter(
SqlaTable.table_name.in_(["energy_usage",
"unicode_test"])).all())
self.rls_entry.clause = "value > 1"
self.rls_entry.roles.append(
security_manager.find_role("Gamma")
) # db.session.query(Role).filter_by(name="Gamma").first())
self.rls_entry.roles.append(security_manager.find_role("Alpha"))
db.session.add(self.rls_entry)
db.session.commit()
def test_sql_lab_insert_rls(
mocker: MockerFixture,
session: Session,
app_context: None,
) -> None:
"""
Integration test for `insert_rls`.
"""
from flask_appbuilder.security.sqla.models import Role, User
from superset.connectors.sqla.models import RowLevelSecurityFilter, SqlaTable
from superset.models.core import Database
from superset.models.sql_lab import Query
from superset.security.manager import SupersetSecurityManager
from superset.sql_lab import execute_sql_statement
from superset.utils.core import RowLevelSecurityFilterType
engine = session.connection().engine
Query.metadata.create_all(engine) # pylint: disable=no-member
connection = engine.raw_connection()
connection.execute("CREATE TABLE t (c INTEGER)")
for i in range(10):
connection.execute("INSERT INTO t VALUES (?)", (i, ))
cursor = connection.cursor()
query = Query(
sql="SELECT c FROM t",
client_id="abcde",
database=Database(database_name="test_db", sqlalchemy_uri="sqlite://"),
schema=None,
limit=5,
select_as_cta_used=False,
)
session.add(query)
session.commit()
admin = User(
first_name="Alice",
last_name="Doe",
email="*****@*****.**",
username="******",
roles=[Role(name="Admin")],
)
# first without RLS
with override_user(admin):
superset_result_set = execute_sql_statement(
sql_statement=query.sql,
query=query,
session=session,
cursor=cursor,
log_params=None,
apply_ctas=False,
)
assert (superset_result_set.to_pandas_df().to_markdown() == """
| | c |
|---:|----:|
| 0 | 0 |
| 1 | 1 |
| 2 | 2 |
| 3 | 3 |
| 4 | 4 |""".strip())
assert query.executed_sql == "SELECT c FROM t\nLIMIT 6"
# now with RLS
rls = RowLevelSecurityFilter(
filter_type=RowLevelSecurityFilterType.REGULAR,
tables=[SqlaTable(database_id=1, schema=None, table_name="t")],
roles=[admin.roles[0]],
group_key=None,
clause="c > 5",
)
session.add(rls)
session.flush()
mocker.patch.object(SupersetSecurityManager,
"find_user",
return_value=admin)
mocker.patch("superset.sql_lab.is_feature_enabled", return_value=True)
with override_user(admin):
superset_result_set = execute_sql_statement(
sql_statement=query.sql,
query=query,
session=session,
cursor=cursor,
log_params=None,
apply_ctas=False,
)
assert (superset_result_set.to_pandas_df().to_markdown() == """
| | c |
|---:|----:|
| 0 | 6 |
| 1 | 7 |
| 2 | 8 |
| 3 | 9 |""".strip())
assert query.executed_sql == "SELECT c FROM t WHERE (t.c > 5)\nLIMIT 6"
