[superset] branch master updated: fix: SQL Lab cancel query in Redshift database connection does not wo… (#16326)

[el...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

elizabeth pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 90d486a643 fix: SQL Lab cancel query in Redshift database connection does not wo… (#16326)
90d486a643 is described below

commit 90d486a643e59b44278d2cef64652151eb6248f6
Author: yourssvk <yo...@gmail.com>
AuthorDate: Wed Jun 29 23:42:59 2022 +0530

    fix: SQL Lab cancel query in Redshift database connection does not wo… (#16326)
    
    * fix: SQL Lab cancel query in Redshift database connection does not work #16325
    
    
    Co-authored-by: Venkata Krishnan Somasundaram <ve...@Venkatas-MacBook-Pro.local>
    Co-authored-by: Elizabeth Thompson <es...@gmail.com>
---
 superset/db_engine_specs/redshift.py | 42 +++++++++++++++++++++++++++++++++++-
 1 file changed, 41 insertions(+), 1 deletion(-)

diff --git a/superset/db_engine_specs/redshift.py b/superset/db_engine_specs/redshift.py
index 6b79007526..7e2717d776 100644
--- a/superset/db_engine_specs/redshift.py
+++ b/superset/db_engine_specs/redshift.py
@@ -14,14 +14,18 @@
 # KIND, either express or implied.  See the License for the
 # specific language governing permissions and limitations
 # under the License.
+import logging
 import re
-from typing import Any, Dict, Pattern, Tuple
+from typing import Any, Dict, Optional, Pattern, Tuple
 
 from flask_babel import gettext as __
 
 from superset.db_engine_specs.base import BasicParametersMixin
 from superset.db_engine_specs.postgres import PostgresBaseEngineSpec
 from superset.errors import SupersetErrorType
+from superset.models.sql_lab import Query
+
+logger = logging.getLogger()
 
 # Regular expressions to catch custom errors
 CONNECTION_ACCESS_DENIED_REGEX = re.compile(
@@ -101,3 +105,39 @@ class RedshiftEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
         :return: Conditionally mutated label
         """
         return label.lower()
+
+    @classmethod
+    def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]:
+        """
+        Get Redshift PID that will be used to cancel all other running
+        queries in the same session.
+
+        :param cursor: Cursor instance in which the query will be executed
+        :param query: Query instance
+        :return: Redshift PID
+        """
+        cursor.execute("SELECT pg_backend_pid()")
+        row = cursor.fetchone()
+        return row[0]
+
+    @classmethod
+    def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool:
+        """
+        Cancel query in the underlying database.
+
+        :param cursor: New cursor instance to the db of the query
+        :param query: Query instance
+        :param cancel_query_id: Redshift PID
+        :return: True if query cancelled successfully, False otherwise
+        """
+        try:
+            logger.info("Killing Redshift PID:%s", str(cancel_query_id))
+            cursor.execute(
+                "SELECT pg_cancel_backend(procpid) "
+                "FROM pg_stat_activity "
+                f"WHERE procpid='{cancel_query_id}'"
+            )
+            cursor.close()
+        except Exception:  # pylint: disable=broad-except
+            return False
+        return True