rusackas commented on code in PR #40662:
URL: https://github.com/apache/superset/pull/40662#discussion_r3384699102
##########
superset/commands/sql_lab/estimate.py:
##########
@@ -69,6 +77,49 @@ def validate(self) -> None:
)
security_manager.raise_for_access(database=self._database)
+ def _apply_sql_security(self, sql: str) -> str:
+ """Run the disallowed-function/table, DML and RLS controls against the
+ SQL to be estimated, mirroring ``sql_lab.execute_sql_statements``.
+
+ Returns the SQL with RLS predicates injected (when ``RLS_IN_SQLLAB`` is
+ enabled), so the cost estimate reflects the same constrained query the
+ user would actually be allowed to run.
+ """
+ db_engine_spec = self._database.db_engine_spec
+ parsed_script = SQLScript(sql, engine=db_engine_spec.engine)
+
+ disallowed_functions = app.config["DISALLOWED_SQL_FUNCTIONS"].get(
+ db_engine_spec.engine,
+ set(),
+ )
+ if disallowed_functions and parsed_script.check_functions_present(
+ disallowed_functions
+ ):
+ raise SupersetDisallowedSQLFunctionException(disallowed_functions)
+
+ disallowed_tables = app.config["DISALLOWED_SQL_TABLES"].get(
+ db_engine_spec.engine,
+ set(),
+ )
+ if disallowed_tables and
parsed_script.check_tables_present(disallowed_tables):
+ found_tables = set()
+ for statement in parsed_script.statements:
+ present = {table.table.lower() for table in statement.tables}
+ for table in disallowed_tables:
+ if table.lower() in present:
+ found_tables.add(table)
+ raise SupersetDisallowedSQLTableException(found_tables or
disallowed_tables)
+
+ if parsed_script.has_mutation() and not self._database.allow_dml:
+ raise SupersetDMLNotAllowedException()
+
+ if is_feature_enabled("RLS_IN_SQLLAB"):
+ for statement in parsed_script.statements:
+ apply_rls(self._database, self._catalog, self._schema,
statement)
+ return parsed_script.format()
Review Comment:
Good catch — fixed. Before injecting RLS I now resolve the database default
catalog/schema the same way the execution path does (`catalog = self._catalog
or self._database.get_default_catalog()`, then `schema = self._schema or
self._database.get_default_schema(catalog)`, mirroring
`SQLExecutor._prepare_scripts` and `sql_lab.execute_sql_statements`), instead
of passing the raw `""`/`None`. That way unqualified table references match
datasets under the default schema and the estimate enforces the same RLS
predicates the real query would. Added unit tests pinning both the
default-resolution and explicit-override cases.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
