betodealmeida commented on code in PR #27943: URL: https://github.com/apache/superset/pull/27943#discussion_r1575286982
########## superset/models/core.py: ########## @@ -572,7 +587,30 @@ def quote_identifier(self) -> Callable[[str], str]: def get_reserved_words(self) -> set[str]: return self.get_dialect().preparer.reserved_words - def get_df( # pylint: disable=too-many-locals + def mutate_sql_based_on_config(self, sql_: str, is_splitted: bool = False) -> str: + """ + Mutates the SQL query based on the app configuration. + + Two config params here affect the behavior of the SQL query mutator: + - `SQL_QUERY_MUTATOR`: A user-provided function that mutates the SQL query. + - `MUTATE_AFTER_SPLIT`: If True, the SQL query mutator is only called after the + sql is broken down into smaller queries. If False, the SQL query mutator applies + on the group of queries as a whole. Here the called passes the context + as to whether the SQL is split or already. + """ + sql_mutator = config["SQL_QUERY_MUTATOR"] + if sql_mutator and ( + (is_splitted and config["MUTATE_AFTER_SPLIT"]) + or (not is_splitted and not config["MUTATE_AFTER_SPLIT"]) + ): Review Comment: Small suggestion to simplify the logic: ```suggestion if sql_mutator and (is_splitted == config["MUTATE_AFTER_SPLIT"]): ``` I'd also rename `is_splitted` to `is_split` ([ref](https://english.stackexchange.com/questions/467873/the-past-participle-of-split-split-or-splitted#467898)). ########## superset/models/core.py: ########## @@ -603,42 +632,30 @@ def _log_query(sql: str) -> None: with self.get_raw_connection(schema=schema) as conn: cursor = conn.cursor() - for sql_ in sqls[:-1]: - if mutate_after_split: - sql_ = sql_query_mutator( - sql_, - security_manager=security_manager, - database=None, - ) + df = None + for i, sql_ in enumerate(sqls): + sql_ = self.mutate_sql_based_on_config(sql_, is_splitted=True) _log_query(sql_) - self.db_engine_spec.execute(cursor, sql_, self) - cursor.fetchall() - - if mutate_after_split: - last_sql = sql_query_mutator( - sqls[-1], - security_manager=security_manager, - database=None, - ) - _log_query(last_sql) - self.db_engine_spec.execute(cursor, last_sql, self) - else: - _log_query(sqls[-1]) - self.db_engine_spec.execute(cursor, sqls[-1], self) - - data = self.db_engine_spec.fetch_data(cursor) - result_set = SupersetResultSet( - data, cursor.description, self.db_engine_spec - ) - df = result_set.to_pandas_df() + with event_logger.log_context( + action="execute_sql", + database=self, + object_ref=__name__, + ): + self.db_engine_spec.execute(cursor, sql_, self) + if i < len(sqls) - 1: + # If it's not the last, we don't keep the results + cursor.fetchall() Review Comment: It would be nice in the future if we captured all results and surfaced them in the UI, it's a feature that similar tools have. ########## superset/utils/log.py: ########## @@ -47,10 +47,12 @@ def collect_request_payload() -> dict[str, Any]: payload: dict[str, Any] = { "path": request.path, - **request.form.to_dict(), - # url search params can overwrite POST body - **request.args.to_dict(), } + payload.update(**request.form.to_dict()) + payload.update(**request.args.to_dict()) Review Comment: ```suggestion ``` -- 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: notifications-unsubscr...@superset.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org --------------------------------------------------------------------- To unsubscribe, e-mail: notifications-unsubscr...@superset.apache.org For additional commands, e-mail: notifications-h...@superset.apache.org