Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix: SQL Lab cancel query in Redshift database connection does not wo… #16326

Merged
merged 5 commits into from
Jun 29, 2022
Merged
Changes from 2 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
42 changes: 41 additions & 1 deletion superset/db_engine_specs/redshift.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,13 +15,17 @@
# specific language governing permissions and limitations
# under the License.
import re
eschutho marked this conversation as resolved.
Show resolved Hide resolved
from typing import Any, Dict, Pattern, Tuple
import logging
eschutho marked this conversation as resolved.
Show resolved Hide resolved
from typing import Any, Dict, Pattern, Tuple, Optional
eschutho marked this conversation as resolved.
Show resolved Hide resolved

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(
Expand Down Expand Up @@ -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