From 9a26a211d46b746357734e9dce48eb8b0d72d1f4 Mon Sep 17 00:00:00 2001
From: yourssvk <yourssvk@gmail.com>
Date: Wed, 29 Jun 2022 23:42:59 +0530
Subject: [PATCH] =?UTF-8?q?fix:=20SQL=20Lab=20cancel=20query=20in=20Redshi?=
 =?UTF-8?q?ft=20database=20connection=20does=20not=20wo=E2=80=A6=20(#16326?=
 =?UTF-8?q?)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix: SQL Lab cancel query in Redshift database connection does not work #16325

Co-authored-by: Venkata Krishnan Somasundaram <venkata_cred@Venkatas-MacBook-Pro.local>
Co-authored-by: Elizabeth Thompson <eschutho@gmail.com>
(cherry picked from commit 90d486a643e59b44278d2cef64652151eb6248f6)
---
 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 6b79007526a1..7e2717d77610 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 @@ def _mutate_label(label: str) -> str:
         :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