feat: cancel db query on stop

superset-frontend/src/SqlLab/components/SqlEditor.jsx

@@ -194,6 +194,7 @@ class SqlEditor extends React.PureComponent {
       WINDOW_RESIZE_THROTTLE_MS,
     );
 
+    this.onBeforeUnload = this.onBeforeUnload.bind(this);
     this.renderDropdown = this.renderDropdown.bind(this);
   }
 
@@ -212,6 +213,7 @@ class SqlEditor extends React.PureComponent {
     this.setState({ height: this.getSqlEditorHeight() });
 
     window.addEventListener('resize', this.handleWindowResize);
+    window.addEventListener('beforeunload', this.onBeforeUnload.bind(this));
 
     // setup hotkeys
     const hotkeys = this.getHotkeyConfig();
@@ -222,6 +224,7 @@ class SqlEditor extends React.PureComponent {
 
   componentWillUnmount() {
     window.removeEventListener('resize', this.handleWindowResize);
+    window.removeEventListener('beforeunload', this.onBeforeUnload.bind(this));
   }
 
   onResizeStart() {
@@ -242,6 +245,11 @@ class SqlEditor extends React.PureComponent {
     }
   }
 
+  onBeforeUnload(event) {
+    event.preventDefault();
+    this.stopQuery();
+  }
+
   onSqlChanged(sql) {
     this.setState({ sql });
     this.setQueryEditorSqlWithDebounce(sql);

superset/db_engine_specs/base.py

@@ -1303,6 +1303,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:
+        """
+        Cancels query in the underlying database.
+        The method is called only when payload is not None.
+        :param cursor: New cursor instance to the db of the query
+        :param query: Query instance
+        :param payload: Value returned by get_cancel_query_payload or set in
+        other life-cycle methods of the query
+        """
+
 
 # schema for adding a database by providing parameters instead of the
 # full SQLAlchemy URI

superset/db_engine_specs/mysql.py

@@ -37,6 +37,7 @@
 
 from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin
 from superset.errors import SupersetErrorType
+from superset.models.sql_lab import Query
 from superset.utils import core as utils
 from superset.utils.core import ColumnSpec, GenericDataType
 
@@ -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:
+        cursor.execute("KILL CONNECTION %d" % payload)

superset/db_engine_specs/postgres.py

@@ -40,6 +40,7 @@
 from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin
 from superset.errors import SupersetErrorType
 from superset.exceptions import SupersetException
+from superset.models.sql_lab import Query
 from superset.utils import core as utils
 from superset.utils.core import ColumnSpec, GenericDataType
 
@@ -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
+        )

superset/db_engine_specs/snowflake.py

@@ -25,6 +25,7 @@
 
 from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
 from superset.errors import SupersetErrorType
+from superset.models.sql_lab import Query
 from superset.utils import core as utils
 
 if TYPE_CHECKING:
@@ -128,3 +129,13 @@ def mutate_db_for_connection_test(database: "Database") -> None:
         engine_params["connect_args"] = connect_args
         extra["engine_params"] = engine_params
         database.extra = json.dumps(extra)
+
+    @classmethod
+    def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any:
+        cursor.execute("SELECT CURRENT_SESSION()")
+        row = cursor.fetchone()
+        return row[0]
+
+    @classmethod
+    def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None:
+        cursor.execute("SELECT SYSTEM$CANCEL_ALL_QUERIES(%s)" % payload)

superset/sql_lab.py

@@ -73,6 +73,7 @@ def dummy_sql_query_mutator(
 SQL_QUERY_MUTATOR = config.get("SQL_QUERY_MUTATOR") or dummy_sql_query_mutator
 log_query = config["QUERY_LOGGER"]
 logger = logging.getLogger(__name__)
+cancel_payload_key = "cancel_payload"
 
 
 class SqlLabException(Exception):
@@ -83,6 +84,10 @@ class SqlLabSecurityException(SqlLabException):
     pass
 
 
+class SqlLabQueryStoppedException(SqlLabException):
+    pass
+
+
 def handle_query_error(
     ex: Exception,
     query: Query,
@@ -187,7 +192,7 @@ def get_sql_results(  # pylint: disable=too-many-arguments
             return handle_query_error(ex, query, session)
 
 
-# pylint: disable=too-many-arguments, too-many-locals
+# pylint: disable=too-many-arguments, too-many-locals, too-many-statements
 def execute_sql_statement(
     sql_statement: str,
     query: Query,
@@ -288,6 +293,12 @@ def execute_sql_statement(
             )
         )
     except Exception as ex:
+        # query is stopped in another thread/worker
+        # stopping rases expected exceptions which we should skip
+        session.refresh(query)
+        if query.status == QueryStatus.STOPPED:
+            raise SqlLabQueryStoppedException()
+
         logger.error("Query %d: %s", query.id, type(ex), exc_info=True)
         logger.debug("Query %d: %s", query.id, ex)
         raise SqlLabException(db_engine_spec.extract_error_message(ex))
@@ -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)
+        if cancel_query_payload is not None:
+            query.set_extra_json_key(cancel_payload_key, cancel_query_payload)
+            session.commit()
         statement_count = len(statements)
         for i, statement in enumerate(statements):
             # Check if stopped
-            query = get_query(query_id, session)
+            session.refresh(query)
             if query.status == QueryStatus.STOPPED:
-                return None
+                payload.update({"status": query.status})
+                return payload
 
             # For CTAS we create the table only on the last statement
             apply_ctas = query.select_as_cta and (
@@ -466,6 +483,9 @@ def execute_sql_statements(  # pylint: disable=too-many-arguments, too-many-loca
                     log_params,
                     apply_ctas,
                 )
+            except SqlLabQueryStoppedException:
+                payload.update({"status": QueryStatus.STOPPED})
+                return payload
             except Exception as ex:  # pylint: disable=broad-except
                 msg = str(ex)
                 prefix_message = (
@@ -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."""
+    cancel_payload = query.extra.get(cancel_payload_key, None)
+    if cancel_payload is None:
+        return
+
+    database = query.database
+    engine = database.get_sqla_engine(
+        schema=query.schema,
+        nullpool=True,
+        user_name=user_name,
+        source=QuerySource.SQL_LAB,
+    )
+    db_engine_spec = database.db_engine_spec
+
+    with closing(engine.raw_connection()) as conn:
+        with closing(conn.cursor()) as cursor:
+            db_engine_spec.cancel_query(cursor, query, cancel_payload)

superset/views/core.py

@@ -2307,6 +2307,7 @@ def stop_query(self) -> FlaskResponse:
             return self.json_response("OK")
         query.status = QueryStatus.STOPPED
         db.session.commit()
+        sql_lab.cancel_query(query, g.user.username if g.user else None)
 
         return self.json_response("OK")