From 8e93d0326031f54ab8cbe696e371714973a90f03 Mon Sep 17 00:00:00 2001 From: Venkata Krishnan Somasundaram Date: Wed, 18 Aug 2021 20:57:25 +0530 Subject: [PATCH 1/5] fix: SQL Lab cancel query in Redshift database connection does not work #16325 --- 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..feb21a31b27a 100644 --- a/superset/db_engine_specs/redshift.py +++ b/superset/db_engine_specs/redshift.py @@ -15,13 +15,17 @@ # specific language governing permissions and limitations # under the License. import re -from typing import Any, Dict, Pattern, Tuple +import logging +from typing import Any, Dict, Pattern, Tuple, Optional 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 From f6106f5db2e8a3e2761a9fe3774b7c847e1e0236 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 24 Jun 2022 16:55:50 -0700 Subject: [PATCH 2/5] lint --- superset/db_engine_specs/redshift.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/db_engine_specs/redshift.py b/superset/db_engine_specs/redshift.py index feb21a31b27a..f0f44516ec10 100644 --- a/superset/db_engine_specs/redshift.py +++ b/superset/db_engine_specs/redshift.py @@ -131,7 +131,7 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: :return: True if query cancelled successfully, False otherwise """ try: - logger.info("Killing Redshift PID:%s",str(cancel_query_id)) + logger.info("Killing Redshift PID:%s", str(cancel_query_id)) cursor.execute( "SELECT pg_cancel_backend(procpid) " "FROM pg_stat_activity " From fe77cd1619b27031e2defb7cdceedf5bac0265ad Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 24 Jun 2022 19:51:22 -0700 Subject: [PATCH 3/5] Lint --- superset/db_engine_specs/redshift.py | 1 + 1 file changed, 1 insertion(+) diff --git a/superset/db_engine_specs/redshift.py b/superset/db_engine_specs/redshift.py index f0f44516ec10..bfebdd49033d 100644 --- a/superset/db_engine_specs/redshift.py +++ b/superset/db_engine_specs/redshift.py @@ -14,6 +14,7 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. +import logging import re import logging from typing import Any, Dict, Pattern, Tuple, Optional From a7d9be9cf43c5a30047134a0e03af88613442ab8 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Fri, 24 Jun 2022 19:51:39 -0700 Subject: [PATCH 4/5] Lint --- superset/db_engine_specs/redshift.py | 1 - 1 file changed, 1 deletion(-) diff --git a/superset/db_engine_specs/redshift.py b/superset/db_engine_specs/redshift.py index bfebdd49033d..884e06d494f2 100644 --- a/superset/db_engine_specs/redshift.py +++ b/superset/db_engine_specs/redshift.py @@ -16,7 +16,6 @@ # under the License. import logging import re -import logging from typing import Any, Dict, Pattern, Tuple, Optional from flask_babel import gettext as __ From 66d3dce649dbb8f017a9b7759e7d71b12149d423 Mon Sep 17 00:00:00 2001 From: Elizabeth Thompson Date: Mon, 27 Jun 2022 16:21:08 -0700 Subject: [PATCH 5/5] lint --- superset/db_engine_specs/redshift.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/db_engine_specs/redshift.py b/superset/db_engine_specs/redshift.py index 884e06d494f2..7e2717d77610 100644 --- a/superset/db_engine_specs/redshift.py +++ b/superset/db_engine_specs/redshift.py @@ -16,7 +16,7 @@ # under the License. import logging import re -from typing import Any, Dict, Pattern, Tuple, Optional +from typing import Any, Dict, Optional, Pattern, Tuple from flask_babel import gettext as __