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 <venkata_cred@Venkatas-MacBook-Pro.local>
Co-authored-by: Elizabeth Thompson <eschutho@gmail.com>
This commit is contained in:
yourssvk 2022-06-29 23:42:59 +05:30 committed by GitHub
parent cadd259788
commit 90d486a643
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 41 additions and 1 deletions

View File

@ -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