[GitHub] [superset] etr2460 commented on a change in pull request #15403: feat: cancel db query on stop

[GitBox <gi...@apache.org>]

etr2460 commented on a change in pull request #15403:
URL: https://github.com/apache/superset/pull/15403#discussion_r666544702



##########
File path: superset/db_engine_specs/postgres.py
##########
@@ -296,3 +297,17 @@ def get_column_spec(  # type: ignore
         return super().get_column_spec(
             native_type, column_type_mappings=column_type_mappings
         )
+
+    @classmethod
+    def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any:
+        cursor.execute("SELECT current_user")
+        row = cursor.fetchone()
+        return row[0]
+
+    @classmethod
+    def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None:
+        cursor.execute(
+            "SELECT pg_terminate_backend(pid) "
+            "FROM pg_stat_activity "
+            "WHERE pid <> pg_backend_pid() and usename='%s'" % payload

Review comment:
       yeah, same question here. ideally we can only cancel this query

##########
File path: superset/db_engine_specs/base.py
##########
@@ -1304,6 +1304,28 @@ def get_column_spec(
             )
         return None
 
+    @classmethod
+    def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any:

Review comment:
       Can we improve the types here (e.g. better than `Any`)?

##########
File path: superset/db_engine_specs/mysql.py
##########
@@ -220,3 +221,13 @@ def get_column_spec(  # type: ignore
         return super().get_column_spec(
             native_type, column_type_mappings=column_type_mappings
         )
+
+    @classmethod
+    def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any:
+        cursor.execute("SELECT CONNECTION_ID()")
+        row = cursor.fetchone()
+        return row[0]
+
+    @classmethod
+    def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None:

Review comment:
       sae typing questions here

##########
File path: superset/sql_lab.py
##########
@@ -438,12 +449,18 @@ def execute_sql_statements(  # pylint: disable=too-many-arguments, too-many-loca
     with closing(engine.raw_connection()) as conn:
         # closing the connection closes the cursor as well
         cursor = conn.cursor()
+        cancel_query_payload = db_engine_spec.get_cancel_query_payload(cursor, query)
+        logger.info(cancel_query_payload)

Review comment:
       maybe remove the logger?

##########
File path: superset/sql_lab.py
##########
@@ -562,3 +582,23 @@ def execute_sql_statements(  # pylint: disable=too-many-arguments, too-many-loca
         return payload
 
     return None
+
+
+def cancel_query(query: Query, user_name: Optional[str] = None) -> None:
+    """Cancal a running query."""

Review comment:
       spelling nit: cancel
   
   also please update the docblock to match all the others with arguments

##########
File path: superset/db_engine_specs/base.py
##########
@@ -1304,6 +1304,28 @@ def get_column_spec(
             )
         return None
 
+    @classmethod
+    def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any:
+        """
+        Returns None if query can not be cancelled.
+        :param cursor: Cursor instance in which the query will be executed
+        :param query: Query instance
+        :return: Type of the payload can vary depends on databases
+        but must be jsonable. None if query can't be cancelled.
+        """
+        return None
+
+    @classmethod
+    def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None:

Review comment:
       same type question here




-- 
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-unsubscribe@superset.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org



---------------------------------------------------------------------
To unsubscribe, e-mail: notifications-unsubscribe@superset.apache.org
For additional commands, e-mail: notifications-help@superset.apache.org