betodealmeida commented on code in PR #33542:
URL: https://github.com/apache/superset/pull/33542#discussion_r2116083368
##########
superset/sql_lab.py:
##########
@@ -197,101 +197,131 @@ def get_sql_results( # pylint:
disable=too-many-arguments
return handle_query_error(ex, query)
-def execute_sql_statement( # pylint: disable=too-many-statements,
too-many-locals # noqa: C901
- sql_statement: str,
- query: Query,
- cursor: Any,
- log_params: Optional[dict[str, Any]],
- apply_ctas: bool = False,
-) -> SupersetResultSet:
- """Executes a single SQL statement"""
- database: Database = query.database
- db_engine_spec = database.db_engine_spec
+def apply_rls(query: Query, parsed_statement: BaseSQLStatement[Any]) -> None:
+ """
+ Modify statement inplace to ensure RLS rules are applied.
+ """
+ # we need the default schema to fully qualify the table names
+ default_schema = query.database.get_default_schema_for_query(query)
+
+ # 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.
+ method = (
+ RLSMethod.AS_SUBQUERY
+ if query.database.db_engine_spec.allows_subqueries
+ and query.database.db_engine_spec.allows_alias_in_select
+ else RLSMethod.AS_PREDICATE
+ )
- parsed_query = ParsedQuery(sql_statement, engine=db_engine_spec.engine)
- 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
- )
+ # collect all RLS predicates
+ predicates: dict[Table, list[Any]] = defaultdict(list)
+ for table in parsed_statement.tables:
+ if table_predicates := get_predicates_for_table(
+ table,
+ query.database,
+ query.catalog,
+ default_schema,
+ ):
+ predicates[table].extend(
+ parsed_statement.parse_predicate(predicate)
+ for predicate in table_predicates
+ )
- # Insert any applicable RLS predicates
- parsed_query = ParsedQuery(
- str(
- insert_rls(
- parsed_query._parsed[0], # pylint:
disable=protected-access
- database.id,
- query.schema,
- )
- ),
- engine=db_engine_spec.engine,
- )
+ parsed_statement.apply_rls(query.catalog, default_schema, predicates,
method)
- sql = parsed_query.stripped()
- # This is a test to see if the query is being
- # limited by either the dropdown or the sql.
- # We are testing to see if more rows exist than the limit.
- increased_limit = None if query.limit is None else query.limit + 1
+def get_predicates_for_table(
+ table: Table,
+ database: Database,
+ catalog: str,
+ schema: str,
+) -> list[str]:
+ """
+ Get the RLS predicates for a table.
- if not database.allow_dml:
- errors = []
- try:
- parsed_statement = SQLStatement(
- statement=sql_statement,
- engine=db_engine_spec.engine,
- )
- disallowed = parsed_statement.is_mutating()
- except SupersetParseError as ex:
- # if we fail to parse the query, disallow by default
- disallowed = True
- errors.append(ex.error)
-
- if disallowed:
- errors.append(
- SupersetError(
- message=__(
- "This database does not allow for DDL/DML, and the
query "
- "could not be parsed to confirm it is a read-only
query. Please " # noqa: E501
- "contact your administrator for more assistance."
- ),
- error_type=SupersetErrorType.DML_NOT_ALLOWED_ERROR,
- level=ErrorLevel.ERROR,
- )
+ This is used to inject RLS rules into SQL statements run in SQL Lab.
+ """
+ dataset = (
+ db.session.query(SqlaTable)
+ .filter(
+ and_(
+ SqlaTable.database_id == database.id,
+ SqlaTable.catalog == table.catalog or catalog,
+ SqlaTable.schema == table.schema or schema,
+ SqlaTable.table_name == table.table,
)
- raise SupersetErrorsException(errors)
-
- original_sql = sql
- if apply_ctas:
- if not query.tmp_table_name:
- start_dttm = datetime.fromtimestamp(query.start_time)
- query.tmp_table_name = (
-
f"tmp_{query.user_id}_table_{start_dttm.strftime('%Y_%m_%d_%H_%M_%S')}"
+ )
+ .one_or_none()
+ )
+ if not dataset:
+ return []
+
+ return [
+ str(
+ and_(*filters).compile(
+ dialect=database.get_dialect(),
+ compile_kwargs={"literal_binds": True},
)
- sql = parsed_query.as_create_table(
- query.tmp_table_name,
- schema_name=query.tmp_schema_name,
- method=query.ctas_method,
)
- query.select_as_cta_used = True
+ for filters in dataset.get_sqla_row_level_filters()
+ ]
+
+S = TypeVar("S", bound=BaseSQLStatement[Any])
+
+
+def apply_ctas(query: Query, parsed_statement: S) -> S:
+ """
+ Apply CTAS/CVAS.
+ """
+ if not query.tmp_table_name:
+ start_dttm = datetime.fromtimestamp(query.start_time)
+ prefix = f"tmp_{query.user_id}_table"
+ query.tmp_table_name =
start_dttm.strftime(f"{prefix}_%Y_%m_%d_%H_%M_%S")
+
+ catalog = (
+ query.catalog
+ if query.database.db_engine_spec.supports_cross_catalog_queries
+ else None
Review Comment:
There are a few DBs that support catalogs but not x-catalog queries —
Postgres and Hive are two that I can think of. The thing with these DBs is that
queries like these are invalid:
```sql
CREATE TABLE my_catalog.my_schema.my_table AS SELECT (...)
```
So even if they are configured with a catalog (for example, Postgres is
always configured with a catalog, which it calls a "database"), we don't want
to use the default catalog. For these databases we want to explicitly set the
catalog to null in the `Table` object, so that the query looks like this:
```sql
CREATE TABLE my_schema.my_table AS SELECT (...)
```
--
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]
