villebro commented on code in PR #19837:
URL: https://github.com/apache/superset/pull/19837#discussion_r970323718


##########
superset/connectors/sqla/models.py:
##########
@@ -922,7 +922,9 @@ def mutate_query_from_config(self, sql: str) -> str:
 
         Typically adds comments to the query with context"""
         sql_query_mutator = config["SQL_QUERY_MUTATOR"]
-        if sql_query_mutator:
+        mutate_after_split = config["MUTATE_AFTER_SPLIT"]
+        if sql_query_mutator and not mutate_after_split:
+            username = utils.get_username()

Review Comment:
   This variable appears to be redundant as it's not referenced anywhere below
   ```suggestion
   ```



##########
superset/models/core.py:
##########
@@ -438,12 +442,29 @@ def _log_query(sql: str) -> None:
         with closing(engine.raw_connection()) as conn:
             cursor = conn.cursor()
             for sql_ in sqls[:-1]:
+                if mutate_after_split:
+                    sql_ = sql_query_mutator(
+                        sql_,
+                        user_name=username,
+                        security_manager=security_manager,
+                        database=None
+                    )
                 _log_query(sql_)
                 self.db_engine_spec.execute(cursor, sql_)
                 cursor.fetchall()
-
-            _log_query(sqls[-1])
-            self.db_engine_spec.execute(cursor, sqls[-1])
+            
+            if mutate_after_split:
+                last_sql = sql_query_mutator(
+                    sqls[-1],
+                    user_name=username,
+                    security_manager=security_manager,
+                    database=None
+                )
+                _log_query(last_sql)
+                self.db_engine_spec.execute(cursor, last_sql)
+            else:
+                _log_query(sqls[-1])
+                self.db_engine_spec.execute(cursor, sqls[-1])

Review Comment:
   It appears `sqls[-1]` is needed in both branches of the if statement. This 
should simplify it somewhat
   ```suggestion
               last_sql = sqls[-1]
               if mutate_after_split:
                   last_sql = sql_query_mutator(
                       last_sql,
                       user_name=username,
                       security_manager=security_manager,
                       database=None
                   )
                   _log_query(last_sql)
                   self.db_engine_spec.execute(cursor, last_sql)
               else:
                   _log_query(last_sql)
                   self.db_engine_spec.execute(cursor, last_sql)
   ```



##########
superset/models/core.py:
##########
@@ -415,7 +415,11 @@ def get_df(  # pylint: disable=too-many-locals
         mutator: Optional[Callable[[pd.DataFrame], None]] = None,
     ) -> pd.DataFrame:
         sqls = self.db_engine_spec.parse_sql(sql)
+
         engine = self.get_sqla_engine(schema)
+        username = utils.get_username() or username

Review Comment:
   This variable doesn't appear to be defined in this scope?
   ```suggestion
           username = utils.get_username()
   ```



-- 
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]

Reply via email to