From c6ec8e801f481e515f224906c527f3e91ed27751 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Sun, 27 Jun 2021 01:16:28 +0100 Subject: [PATCH 01/13] feat: cancel db query on stop --- .../src/SqlLab/components/SqlEditor.jsx | 8 ++++ superset/db_engine_specs/base.py | 23 ++++++++++ superset/db_engine_specs/mysql.py | 11 +++++ superset/db_engine_specs/postgres.py | 15 +++++++ superset/db_engine_specs/snowflake.py | 13 +++++- superset/sql_lab.py | 44 ++++++++++++++++++- superset/views/core.py | 1 + 7 files changed, 112 insertions(+), 3 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/SqlEditor.jsx b/superset-frontend/src/SqlLab/components/SqlEditor.jsx index 3883c43b2f916..8c50293e9cbf7 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor.jsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor.jsx @@ -194,6 +194,7 @@ class SqlEditor extends React.PureComponent { WINDOW_RESIZE_THROTTLE_MS, ); + this.onBeforeUnload = this.onBeforeUnload.bind(this); this.renderDropdown = this.renderDropdown.bind(this); } @@ -212,6 +213,7 @@ class SqlEditor extends React.PureComponent { this.setState({ height: this.getSqlEditorHeight() }); window.addEventListener('resize', this.handleWindowResize); + window.addEventListener('beforeunload', this.onBeforeUnload.bind(this)); // setup hotkeys const hotkeys = this.getHotkeyConfig(); @@ -222,6 +224,7 @@ class SqlEditor extends React.PureComponent { componentWillUnmount() { window.removeEventListener('resize', this.handleWindowResize); + window.removeEventListener('beforeunload', this.onBeforeUnload.bind(this)); } onResizeStart() { @@ -242,6 +245,11 @@ class SqlEditor extends React.PureComponent { } } + onBeforeUnload(event) { + event.preventDefault(); + this.stopQuery(); + } + onSqlChanged(sql) { this.setState({ sql }); this.setQueryEditorSqlWithDebounce(sql); diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 3d288b79c854a..031e64134c298 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1303,6 +1303,29 @@ def get_column_spec( ) return None + @classmethod + def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + """ + Returns None if query can not be cancelled. + :param cursor: Cursor instance in which the query will be executed + :param query: Query instance + :return: Type of the payload can vary depends on databases + but must be jsonable. None if query can't be cancelled. + """ + return None + + @classmethod + def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + """ + Cancels query in the underlying database. + The method is called only when payload is not None. + :param cursor: New cursor instance to the db of the query + :param query: Query instance + :param payload: Value returned by get_cancel_query_payload or set in + other life-cycle methods of the query + """ + pass + # schema for adding a database by providing parameters instead of the # full SQLAlchemy URI diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 4bb5979706b5b..d5a29771f4bd8 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -37,6 +37,7 @@ from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin from superset.errors import SupersetErrorType +from superset.models.sql_lab import Query from superset.utils import core as utils from superset.utils.core import ColumnSpec, GenericDataType @@ -207,3 +208,13 @@ def get_column_spec( # type: ignore return super().get_column_spec( native_type, column_type_mappings=column_type_mappings ) + + @classmethod + def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + cursor.execute("SELECT CONNECTION_ID()") + row = cursor.fetchone() + return row[0] + + @classmethod + def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + cursor.execute("KILL CONNECTION %d" % payload) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 03c955a67d81a..a92424cb58af5 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -40,6 +40,7 @@ from superset.db_engine_specs.base import BaseEngineSpec, BasicParametersMixin from superset.errors import SupersetErrorType from superset.exceptions import SupersetException +from superset.models.sql_lab import Query from superset.utils import core as utils from superset.utils.core import ColumnSpec, GenericDataType @@ -286,3 +287,17 @@ def get_column_spec( # type: ignore return super().get_column_spec( native_type, column_type_mappings=column_type_mappings ) + + @classmethod + def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + cursor.execute("SELECT current_user") + row = cursor.fetchone() + return row[0] + + @classmethod + def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + cursor.execute( + "SELECT pg_terminate_backend(pid) " + "FROM pg_stat_activity " + "WHERE pid <> pg_backend_pid() and usename='%s'" % payload + ) diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 7cce0581e9fa5..621c9ca36e005 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -16,12 +16,13 @@ # under the License. import json from datetime import datetime -from typing import Optional, TYPE_CHECKING +from typing import Any, Optional, TYPE_CHECKING from urllib import parse from sqlalchemy.engine.url import URL from superset.db_engine_specs.postgres import PostgresBaseEngineSpec +from superset.models.sql_lab import Query from superset.utils import core as utils if TYPE_CHECKING: @@ -99,3 +100,13 @@ def mutate_db_for_connection_test(database: "Database") -> None: engine_params["connect_args"] = connect_args extra["engine_params"] = engine_params database.extra = json.dumps(extra) + + @classmethod + def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + cursor.execute("SELECT CURRENT_SESSION()") + row = cursor.fetchone() + return row[0] + + @classmethod + def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + cursor.execute("SELECT SYSTEM$CANCEL_ALL_QUERIES(%s)" % payload) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index a40e132e8e3cc..5e7f832809402 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -73,6 +73,7 @@ def dummy_sql_query_mutator( SQL_QUERY_MUTATOR = config.get("SQL_QUERY_MUTATOR") or dummy_sql_query_mutator log_query = config["QUERY_LOGGER"] logger = logging.getLogger(__name__) +cancel_payload_key = "cancel_payload" class SqlLabException(Exception): @@ -83,6 +84,10 @@ class SqlLabSecurityException(SqlLabException): pass +class SqlLabQueryStoppedException(SqlLabException): + pass + + def handle_query_error( ex: Exception, query: Query, @@ -288,6 +293,12 @@ def execute_sql_statement( ) ) except Exception as ex: + # query is stopped in another thread/worker + # stopping rases expected exceptions which we should skip + session.refresh(query) + if query.status == QueryStatus.STOPPED: + raise SqlLabQueryStoppedException() + logger.error("Query %d: %s", query.id, type(ex), exc_info=True) logger.debug("Query %d: %s", query.id, ex) raise SqlLabException(db_engine_spec.extract_error_message(ex)) @@ -438,12 +449,18 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca with closing(engine.raw_connection()) as conn: # closing the connection closes the cursor as well cursor = conn.cursor() + cancel_query_payload = db_engine_spec.get_cancel_query_payload(cursor, query) + logger.info(cancel_query_payload) + if cancel_query_payload is not None: + query.set_extra_json_key(cancel_payload_key, cancel_query_payload) + session.commit() statement_count = len(statements) for i, statement in enumerate(statements): # Check if stopped - query = get_query(query_id, session) + session.refresh(query) if query.status == QueryStatus.STOPPED: - return None + payload.update({"status": query.status}) + return payload # For CTAS we create the table only on the last statement apply_ctas = query.select_as_cta and ( @@ -466,6 +483,9 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca log_params, apply_ctas, ) + except SqlLabQueryStoppedException: + payload.update({"status": QueryStatus.STOPPED}) + return payload except Exception as ex: # pylint: disable=broad-except msg = str(ex) prefix_message = ( @@ -562,3 +582,23 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca return payload return None + + +def cancel_query(query: Query, user_name: Optional[str] = None) -> None: + """Cancal a running query.""" + cancel_payload = query.extra.get(cancel_payload_key, None) + if cancel_payload is None: + return + + database = query.database + engine = database.get_sqla_engine( + schema=query.schema, + nullpool=True, + user_name=user_name, + source=QuerySource.SQL_LAB, + ) + db_engine_spec = database.db_engine_spec + + with closing(engine.raw_connection()) as conn: + with closing(conn.cursor()) as cursor: + db_engine_spec.cancel_query(cursor, query, cancel_payload) diff --git a/superset/views/core.py b/superset/views/core.py index f86f5c9eea295..a6f6aa04c530c 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -2264,6 +2264,7 @@ def stop_query(self) -> FlaskResponse: return self.json_response("OK") query.status = QueryStatus.STOPPED db.session.commit() + sql_lab.cancel_query(query, g.user.username if g.user else None) return self.json_response("OK") From 822e2f3168670c95284cf606f3ac47ec372684b0 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Mon, 28 Jun 2021 14:25:15 +0100 Subject: [PATCH 02/13] fix pylint --- superset/db_engine_specs/base.py | 1 - superset/sql_lab.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 031e64134c298..47dc34a97d2b2 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1324,7 +1324,6 @@ def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: :param payload: Value returned by get_cancel_query_payload or set in other life-cycle methods of the query """ - pass # schema for adding a database by providing parameters instead of the diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 5e7f832809402..284736128b8da 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -192,7 +192,7 @@ def get_sql_results( # pylint: disable=too-many-arguments return handle_query_error(ex, query, session) -# pylint: disable=too-many-arguments, too-many-locals +# pylint: disable=too-many-arguments, too-many-locals, too-many-statements def execute_sql_statement( sql_statement: str, query: Query, From 6d209378dca90d349b704d3ea05c7d17cd36ba52 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Sat, 3 Jul 2021 01:10:14 +0100 Subject: [PATCH 03/13] Add unit tests --- tests/db_engine_specs/mysql_tests.py | 14 ++++++++++++++ tests/db_engine_specs/postgres_tests.py | 17 +++++++++++++++++ tests/db_engine_specs/snowflake_tests.py | 15 +++++++++++++++ 3 files changed, 46 insertions(+) diff --git a/tests/db_engine_specs/mysql_tests.py b/tests/db_engine_specs/mysql_tests.py index 7f0755ac13082..f3eb6ef68034d 100644 --- a/tests/db_engine_specs/mysql_tests.py +++ b/tests/db_engine_specs/mysql_tests.py @@ -21,6 +21,7 @@ from superset.db_engine_specs.mysql import MySQLEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.models.sql_lab import Query from superset.utils.core import GenericDataType from tests.db_engine_specs.base_tests import assert_generic_types, TestDbEngineSpec @@ -235,3 +236,16 @@ def test_extract_errors(self): }, ) ] + + @unittest.mock.patch("sqlalchemy.engine.Engine.connect") + def test_get_cancel_query_payload(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + cursor_mock.fetchone.return_value = [123] + assert MySQLEngineSpec.get_cancel_query_payload(cursor_mock, query) == 123 + + @unittest.mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + assert MySQLEngineSpec.cancel_query(cursor_mock, query, 123) is None diff --git a/tests/db_engine_specs/postgres_tests.py b/tests/db_engine_specs/postgres_tests.py index 7f999f07dbd4b..4e15ae5dcd6c6 100644 --- a/tests/db_engine_specs/postgres_tests.py +++ b/tests/db_engine_specs/postgres_tests.py @@ -23,6 +23,7 @@ from superset.db_engine_specs import get_engine_specs from superset.db_engine_specs.postgres import PostgresEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType +from superset.models.sql_lab import Query from superset.utils.core import GenericDataType from tests.db_engine_specs.base_tests import assert_generic_types, TestDbEngineSpec from tests.fixtures.certificates import ssl_certificate @@ -440,6 +441,22 @@ def test_extract_errors(self): ) ] + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_get_cancel_query_payload(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + cursor_mock.fetchone.return_value = ["testuser"] + assert ( + PostgresEngineSpec.get_cancel_query_payload(cursor_mock, query) + == "testuser" + ) + + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + assert PostgresEngineSpec.cancel_query(cursor_mock, query, "testuser") is None + def test_base_parameters_mixin(): parameters = { diff --git a/tests/db_engine_specs/snowflake_tests.py b/tests/db_engine_specs/snowflake_tests.py index 14ca728df1abd..ea071bee45092 100644 --- a/tests/db_engine_specs/snowflake_tests.py +++ b/tests/db_engine_specs/snowflake_tests.py @@ -15,10 +15,12 @@ # specific language governing permissions and limitations # under the License. import json +from unittest import mock from superset.db_engine_specs.snowflake import SnowflakeEngineSpec from superset.errors import ErrorLevel, SupersetError, SupersetErrorType from superset.models.core import Database +from superset.models.sql_lab import Query from tests.db_engine_specs.base_tests import TestDbEngineSpec @@ -83,3 +85,16 @@ def test_extract_errors(self): }, ) ] + + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_get_cancel_query_payload(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + cursor_mock.fetchone.return_value = [123] + assert SnowflakeEngineSpec.get_cancel_query_payload(cursor_mock, query) == 123 + + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query(self, engine_mock): + query = Query() + cursor_mock = engine_mock.return_value.__enter__.return_value + assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, None) is None From efb86e51af1959de182faf486a6877164ebd933f Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Wed, 7 Jul 2021 01:24:22 +0100 Subject: [PATCH 04/13] Do not bind multiple times --- superset-frontend/src/SqlLab/components/SqlEditor.jsx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/SqlEditor.jsx b/superset-frontend/src/SqlLab/components/SqlEditor.jsx index 8c50293e9cbf7..19daf541a2fd9 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor.jsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor.jsx @@ -213,7 +213,7 @@ class SqlEditor extends React.PureComponent { this.setState({ height: this.getSqlEditorHeight() }); window.addEventListener('resize', this.handleWindowResize); - window.addEventListener('beforeunload', this.onBeforeUnload.bind(this)); + window.addEventListener('beforeunload', this.onBeforeUnload); // setup hotkeys const hotkeys = this.getHotkeyConfig(); @@ -224,7 +224,7 @@ class SqlEditor extends React.PureComponent { componentWillUnmount() { window.removeEventListener('resize', this.handleWindowResize); - window.removeEventListener('beforeunload', this.onBeforeUnload.bind(this)); + window.removeEventListener('beforeunload', this.onBeforeUnload); } onResizeStart() { From 54da5e42cb0bef9c2d9bc15a5ba1a0c6870ec4c3 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 02:22:12 +0100 Subject: [PATCH 05/13] Stop only running queries --- superset-frontend/src/SqlLab/components/SqlEditor.jsx | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/SqlEditor.jsx b/superset-frontend/src/SqlLab/components/SqlEditor.jsx index 19daf541a2fd9..f869e6bb1f5f7 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor.jsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor.jsx @@ -246,8 +246,10 @@ class SqlEditor extends React.PureComponent { } onBeforeUnload(event) { - event.preventDefault(); - this.stopQuery(); + if (this.props.latestQuery?.state === 'running') { + event.preventDefault(); + this.stopQuery(); + } } onSqlChanged(sql) { From 0f4947160ff0f51b76aa3b0f4d8fb79d6cbcd008 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 03:04:51 +0100 Subject: [PATCH 06/13] Postgres to cancel only the required query --- superset/db_engine_specs/postgres.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 4fd06132df1e3..1871fd6bc38c7 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -300,7 +300,7 @@ def get_column_spec( # type: ignore @classmethod def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: - cursor.execute("SELECT current_user") + cursor.execute("SELECT pg_backend_pid()") row = cursor.fetchone() return row[0] @@ -309,5 +309,5 @@ def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: cursor.execute( "SELECT pg_terminate_backend(pid) " "FROM pg_stat_activity " - "WHERE pid <> pg_backend_pid() and usename='%s'" % payload + "WHERE pid='%s'" % payload ) From ec3947757925981b2f07d4070dae8fcb81f69623 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 03:05:36 +0100 Subject: [PATCH 07/13] Remove extra log --- superset/sql_lab.py | 1 - 1 file changed, 1 deletion(-) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 284736128b8da..ea2fd7814023f 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -450,7 +450,6 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca # closing the connection closes the cursor as well cursor = conn.cursor() cancel_query_payload = db_engine_spec.get_cancel_query_payload(cursor, query) - logger.info(cancel_query_payload) if cancel_query_payload is not None: query.set_extra_json_key(cancel_payload_key, cancel_query_payload) session.commit() From 4706b4308ec4a82aa765aafaa70ad2057977b98b Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 03:17:48 +0100 Subject: [PATCH 08/13] Add docstring --- superset/sql_lab.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index ea2fd7814023f..339f666fa14b3 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -584,7 +584,13 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca def cancel_query(query: Query, user_name: Optional[str] = None) -> None: - """Cancal a running query.""" + """ + Cancel a running query. + + :param query: Query to cancel + :param user_name: Default username + :return: None + """ cancel_payload = query.extra.get(cancel_payload_key, None) if cancel_payload is None: return From ef8436b783a2baac74322ff2b98c0189edc332e6 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 04:23:53 +0100 Subject: [PATCH 09/13] Better types, docstring and naming --- superset/db_engine_specs/base.py | 18 +++++++++------- superset/db_engine_specs/mysql.py | 21 ++++++++++++++++--- superset/db_engine_specs/postgres.py | 21 ++++++++++++++++--- superset/db_engine_specs/snowflake.py | 21 ++++++++++++++++--- superset/sql_lab.py | 14 ++++++------- .../db_engine_specs/mysql_tests.py | 4 ++-- .../db_engine_specs/postgres_tests.py | 11 ++++------ .../db_engine_specs/snowflake_tests.py | 6 +++--- 8 files changed, 80 insertions(+), 36 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 8fef471b38835..0634c8470c492 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1305,24 +1305,26 @@ def get_column_spec( return None @classmethod - def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: """ - Returns None if query can not be cancelled. + Select identifiers from the database engine that uniquely identifies the + queries to cancel. The identifier is typically a session id, process id + or similar. + :param cursor: Cursor instance in which the query will be executed :param query: Query instance - :return: Type of the payload can vary depends on databases - but must be jsonable. None if query can't be cancelled. + :return: Query identifier """ return None @classmethod - def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: """ - Cancels query in the underlying database. - The method is called only when payload is not None. + Cancel query in the underlying database. + :param cursor: New cursor instance to the db of the query :param query: Query instance - :param payload: Value returned by get_cancel_query_payload or set in + :param cancel_query_id: Value returned by get_cancel_query_payload or set in other life-cycle methods of the query """ diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 01f41fe8b38c4..896edfa6f7150 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -223,11 +223,26 @@ def get_column_spec( # type: ignore ) @classmethod - def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: + """ + Get MySQL connection ID that will be used to cancel all other running + queries in the same connection. + + :param cursor: Cursor instance in which the query will be executed + :param query: Query instance + :return: MySQL Connection ID + """ cursor.execute("SELECT CONNECTION_ID()") row = cursor.fetchone() return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: - cursor.execute("KILL CONNECTION %d" % payload) + def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + """ + 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: MySQL Connection ID + """ + cursor.execute("KILL CONNECTION %s" % cancel_query_id) diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 1871fd6bc38c7..903d898bb938b 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -299,15 +299,30 @@ def get_column_spec( # type: ignore ) @classmethod - def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: + """ + Get Postgres 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: Postgres PID + """ cursor.execute("SELECT pg_backend_pid()") row = cursor.fetchone() return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: + def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + """ + 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: Postgres PID + """ cursor.execute( "SELECT pg_terminate_backend(pid) " "FROM pg_stat_activity " - "WHERE pid='%s'" % payload + "WHERE pid='%s'" % cancel_query_id ) diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index 2547171916ecd..a12da85939f09 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -131,11 +131,26 @@ def mutate_db_for_connection_test(database: "Database") -> None: database.extra = json.dumps(extra) @classmethod - def get_cancel_query_payload(cls, cursor: Any, query: Query) -> Any: + def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: + """ + Get Snowflake session ID 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: Snowflake Session ID + """ cursor.execute("SELECT CURRENT_SESSION()") row = cursor.fetchone() return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, payload: Any) -> None: - cursor.execute("SELECT SYSTEM$CANCEL_ALL_QUERIES(%s)" % payload) + def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + """ + 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: Snowflake Session ID + """ + cursor.execute("SELECT SYSTEM$CANCEL_ALL_QUERIES(%s)" % cancel_query_id) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 339f666fa14b3..32c8fe706c00e 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -73,7 +73,7 @@ def dummy_sql_query_mutator( SQL_QUERY_MUTATOR = config.get("SQL_QUERY_MUTATOR") or dummy_sql_query_mutator log_query = config["QUERY_LOGGER"] logger = logging.getLogger(__name__) -cancel_payload_key = "cancel_payload" +cancel_query_key = "cancel_query" class SqlLabException(Exception): @@ -449,9 +449,9 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca with closing(engine.raw_connection()) as conn: # closing the connection closes the cursor as well cursor = conn.cursor() - cancel_query_payload = db_engine_spec.get_cancel_query_payload(cursor, query) - if cancel_query_payload is not None: - query.set_extra_json_key(cancel_payload_key, cancel_query_payload) + cancel_query_id = db_engine_spec.get_cancel_query_id(cursor, query) + if cancel_query_id is not None: + query.set_extra_json_key(cancel_query_key, cancel_query_id) session.commit() statement_count = len(statements) for i, statement in enumerate(statements): @@ -591,8 +591,8 @@ def cancel_query(query: Query, user_name: Optional[str] = None) -> None: :param user_name: Default username :return: None """ - cancel_payload = query.extra.get(cancel_payload_key, None) - if cancel_payload is None: + cancel_query_id = query.extra.get(cancel_query_key, None) + if cancel_query_id is None: return database = query.database @@ -606,4 +606,4 @@ def cancel_query(query: Query, user_name: Optional[str] = None) -> None: with closing(engine.raw_connection()) as conn: with closing(conn.cursor()) as cursor: - db_engine_spec.cancel_query(cursor, query, cancel_payload) + db_engine_spec.cancel_query(cursor, query, cancel_query_id) diff --git a/tests/integration_tests/db_engine_specs/mysql_tests.py b/tests/integration_tests/db_engine_specs/mysql_tests.py index 6fc29078a7fe7..b5c2e53b996e7 100644 --- a/tests/integration_tests/db_engine_specs/mysql_tests.py +++ b/tests/integration_tests/db_engine_specs/mysql_tests.py @@ -241,11 +241,11 @@ def test_extract_errors(self): ] @unittest.mock.patch("sqlalchemy.engine.Engine.connect") - def test_get_cancel_query_payload(self, engine_mock): + def test_get_cancel_query_id(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value cursor_mock.fetchone.return_value = [123] - assert MySQLEngineSpec.get_cancel_query_payload(cursor_mock, query) == 123 + assert MySQLEngineSpec.get_cancel_query_id(cursor_mock, query) == 123 @unittest.mock.patch("sqlalchemy.engine.Engine.connect") def test_cancel_query(self, engine_mock): diff --git a/tests/integration_tests/db_engine_specs/postgres_tests.py b/tests/integration_tests/db_engine_specs/postgres_tests.py index da634121cc06c..874ab80d9792a 100644 --- a/tests/integration_tests/db_engine_specs/postgres_tests.py +++ b/tests/integration_tests/db_engine_specs/postgres_tests.py @@ -445,20 +445,17 @@ def test_extract_errors(self): ] @mock.patch("sqlalchemy.engine.Engine.connect") - def test_get_cancel_query_payload(self, engine_mock): + def test_get_cancel_query_id(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - cursor_mock.fetchone.return_value = ["testuser"] - assert ( - PostgresEngineSpec.get_cancel_query_payload(cursor_mock, query) - == "testuser" - ) + cursor_mock.fetchone.return_value = [123] + assert PostgresEngineSpec.get_cancel_query_id(cursor_mock, query) == 123 @mock.patch("sqlalchemy.engine.Engine.connect") def test_cancel_query(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - assert PostgresEngineSpec.cancel_query(cursor_mock, query, "testuser") is None + assert PostgresEngineSpec.cancel_query(cursor_mock, query, 123) is None def test_base_parameters_mixin(): diff --git a/tests/integration_tests/db_engine_specs/snowflake_tests.py b/tests/integration_tests/db_engine_specs/snowflake_tests.py index c4792c1279a1f..36cdb1e4c5596 100644 --- a/tests/integration_tests/db_engine_specs/snowflake_tests.py +++ b/tests/integration_tests/db_engine_specs/snowflake_tests.py @@ -103,14 +103,14 @@ def test_extract_errors(self): ] @mock.patch("sqlalchemy.engine.Engine.connect") - def test_get_cancel_query_payload(self, engine_mock): + def test_get_cancel_query_id(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value cursor_mock.fetchone.return_value = [123] - assert SnowflakeEngineSpec.get_cancel_query_payload(cursor_mock, query) == 123 + assert SnowflakeEngineSpec.get_cancel_query_id(cursor_mock, query) == 123 @mock.patch("sqlalchemy.engine.Engine.connect") def test_cancel_query(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, None) is None + assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, 123) is None From bfc8d60c26ad59cfcc0f3a4b9f460a4a6e8d8bbf Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 04:34:36 +0100 Subject: [PATCH 10/13] Use python3 format strings --- superset/db_engine_specs/mysql.py | 2 +- superset/db_engine_specs/postgres.py | 2 +- superset/db_engine_specs/snowflake.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 896edfa6f7150..0c862c747a694 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -245,4 +245,4 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: :param query: Query instance :param cancel_query_id: MySQL Connection ID """ - cursor.execute("KILL CONNECTION %s" % cancel_query_id) + cursor.execute(f"KILL CONNECTION {cancel_query_id}") diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index 903d898bb938b..e55db9ab89be8 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -324,5 +324,5 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: cursor.execute( "SELECT pg_terminate_backend(pid) " "FROM pg_stat_activity " - "WHERE pid='%s'" % cancel_query_id + f"WHERE pid='{cancel_query_id}'" ) diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index a12da85939f09..b6f1a0b676114 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -153,4 +153,4 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: :param query: Query instance :param cancel_query_id: Snowflake Session ID """ - cursor.execute("SELECT SYSTEM$CANCEL_ALL_QUERIES(%s)" % cancel_query_id) + cursor.execute(f"SELECT SYSTEM$CANCEL_ALL_QUERIES({cancel_query_id})") From e1fc4b3c34a7c7db7254b77b442adbdeb1bd4896 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Fri, 9 Jul 2021 16:47:50 +0100 Subject: [PATCH 11/13] Update superset/sql_lab.py Co-authored-by: Beto Dealmeida --- superset/sql_lab.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/superset/sql_lab.py b/superset/sql_lab.py index 32c8fe706c00e..a2f3096007393 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -294,7 +294,7 @@ def execute_sql_statement( ) except Exception as ex: # query is stopped in another thread/worker - # stopping rases expected exceptions which we should skip + # stopping raises expected exceptions which we should skip session.refresh(query) if query.status == QueryStatus.STOPPED: raise SqlLabQueryStoppedException() From e5223ecb03f3e48f2c9d0b43a851ed9eefb74ff6 Mon Sep 17 00:00:00 2001 From: Peter Kosztolanyi Date: Sun, 11 Jul 2021 01:40:16 +0100 Subject: [PATCH 12/13] Add cancel_query_on_windows_unload option to database --- .../src/SqlLab/components/SqlEditor.jsx | 5 ++++- .../src/SqlLab/reducers/sqlLab.js | 5 ++++- .../database/DatabaseModal/ExtraOptions.tsx | 18 ++++++++++++++++++ .../src/views/CRUD/data/database/types.ts | 1 + superset/databases/api.py | 1 + tests/integration_tests/databases/api_tests.py | 1 + 6 files changed, 29 insertions(+), 2 deletions(-) diff --git a/superset-frontend/src/SqlLab/components/SqlEditor.jsx b/superset-frontend/src/SqlLab/components/SqlEditor.jsx index f869e6bb1f5f7..540cfb77f3f41 100644 --- a/superset-frontend/src/SqlLab/components/SqlEditor.jsx +++ b/superset-frontend/src/SqlLab/components/SqlEditor.jsx @@ -246,7 +246,10 @@ class SqlEditor extends React.PureComponent { } onBeforeUnload(event) { - if (this.props.latestQuery?.state === 'running') { + if ( + this.props.database?.extra_json?.cancel_query_on_windows_unload && + this.props.latestQuery?.state === 'running' + ) { event.preventDefault(); this.stopQuery(); } diff --git a/superset-frontend/src/SqlLab/reducers/sqlLab.js b/superset-frontend/src/SqlLab/reducers/sqlLab.js index daa06a97c88da..9e423ba7c65b2 100644 --- a/superset-frontend/src/SqlLab/reducers/sqlLab.js +++ b/superset-frontend/src/SqlLab/reducers/sqlLab.js @@ -499,7 +499,10 @@ export default function sqlLabReducer(state = {}, action) { [actions.SET_DATABASES]() { const databases = {}; action.databases.forEach(db => { - databases[db.id] = db; + databases[db.id] = { + ...db, + extra_json: JSON.parse(db.extra || ''), + }; }); return { ...state, databases }; }, diff --git a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx index a2a49b2f707b9..80b4044202202 100644 --- a/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx +++ b/superset-frontend/src/views/CRUD/data/database/DatabaseModal/ExtraOptions.tsx @@ -294,6 +294,24 @@ const ExtraOptions = ({ /> + +
+ + +
+ Date: Sun, 11 Jul 2021 03:13:34 +0100 Subject: [PATCH 13/13] Return cancel_query as bool --- superset/db_engine_specs/base.py | 3 ++- superset/db_engine_specs/mysql.py | 10 ++++++++-- superset/db_engine_specs/postgres.py | 18 ++++++++++++------ superset/db_engine_specs/snowflake.py | 10 ++++++++-- superset/exceptions.py | 4 ++++ superset/sql_lab.py | 8 ++++---- superset/views/core.py | 6 +++++- .../db_engine_specs/mysql_tests.py | 8 +++++++- .../db_engine_specs/postgres_tests.py | 8 +++++++- .../db_engine_specs/snowflake_tests.py | 8 +++++++- 10 files changed, 64 insertions(+), 19 deletions(-) diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py index 0634c8470c492..77fdb1d410111 100644 --- a/superset/db_engine_specs/base.py +++ b/superset/db_engine_specs/base.py @@ -1318,7 +1318,7 @@ def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: return None @classmethod - def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> bool: """ Cancel query in the underlying database. @@ -1326,6 +1326,7 @@ def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: :param query: Query instance :param cancel_query_id: Value returned by get_cancel_query_payload or set in other life-cycle methods of the query + :return: True if query cancelled successfully, False otherwise """ diff --git a/superset/db_engine_specs/mysql.py b/superset/db_engine_specs/mysql.py index 0c862c747a694..a48678ef71c79 100644 --- a/superset/db_engine_specs/mysql.py +++ b/superset/db_engine_specs/mysql.py @@ -237,12 +237,18 @@ def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + 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: MySQL Connection ID + :return: True if query cancelled successfully, False otherwise """ - cursor.execute(f"KILL CONNECTION {cancel_query_id}") + try: + cursor.execute(f"KILL CONNECTION {cancel_query_id}") + except Exception: # pylint: disable=broad-except + return False + + return True diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py index e55db9ab89be8..fa8809e151014 100644 --- a/superset/db_engine_specs/postgres.py +++ b/superset/db_engine_specs/postgres.py @@ -313,16 +313,22 @@ def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + 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: Postgres PID + :return: True if query cancelled successfully, False otherwise """ - cursor.execute( - "SELECT pg_terminate_backend(pid) " - "FROM pg_stat_activity " - f"WHERE pid='{cancel_query_id}'" - ) + try: + cursor.execute( + "SELECT pg_terminate_backend(pid) " + "FROM pg_stat_activity " + f"WHERE pid='{cancel_query_id}'" + ) + except Exception: # pylint: disable=broad-except + return False + + return True diff --git a/superset/db_engine_specs/snowflake.py b/superset/db_engine_specs/snowflake.py index b6f1a0b676114..6dd85706562bd 100644 --- a/superset/db_engine_specs/snowflake.py +++ b/superset/db_engine_specs/snowflake.py @@ -145,12 +145,18 @@ def get_cancel_query_id(cls, cursor: Any, query: Query) -> Optional[str]: return row[0] @classmethod - def cancel_query(cls, cursor: Any, query: Query, cancel_query_id: str) -> None: + 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: Snowflake Session ID + :return: True if query cancelled successfully, False otherwise """ - cursor.execute(f"SELECT SYSTEM$CANCEL_ALL_QUERIES({cancel_query_id})") + try: + cursor.execute(f"SELECT SYSTEM$CANCEL_ALL_QUERIES({cancel_query_id})") + except Exception: # pylint: disable=broad-except + return False + + return True diff --git a/superset/exceptions.py b/superset/exceptions.py index 865187f64dc85..89ca41fbfcac0 100644 --- a/superset/exceptions.py +++ b/superset/exceptions.py @@ -210,3 +210,7 @@ def __init__(self, error: ValidationError): extra={"messages": error.messages}, ) super().__init__(error) + + +class SupersetCancelQueryException(SupersetException): + pass diff --git a/superset/sql_lab.py b/superset/sql_lab.py index a2f3096007393..1ebf10e599dcf 100644 --- a/superset/sql_lab.py +++ b/superset/sql_lab.py @@ -583,17 +583,17 @@ def execute_sql_statements( # pylint: disable=too-many-arguments, too-many-loca return None -def cancel_query(query: Query, user_name: Optional[str] = None) -> None: +def cancel_query(query: Query, user_name: Optional[str] = None) -> bool: """ Cancel a running query. :param query: Query to cancel :param user_name: Default username - :return: None + :return: True if query cancelled successfully, False otherwise """ cancel_query_id = query.extra.get(cancel_query_key, None) if cancel_query_id is None: - return + return False database = query.database engine = database.get_sqla_engine( @@ -606,4 +606,4 @@ def cancel_query(query: Query, user_name: Optional[str] = None) -> None: with closing(engine.raw_connection()) as conn: with closing(conn.cursor()) as cursor: - db_engine_spec.cancel_query(cursor, query, cancel_query_id) + return db_engine_spec.cancel_query(cursor, query, cancel_query_id) diff --git a/superset/views/core.py b/superset/views/core.py index cfe79ad139452..169cd373eea13 100755 --- a/superset/views/core.py +++ b/superset/views/core.py @@ -80,6 +80,7 @@ CertificateException, DatabaseNotFound, SerializationError, + SupersetCancelQueryException, SupersetErrorException, SupersetErrorsException, SupersetException, @@ -2335,9 +2336,12 @@ def stop_query(self) -> FlaskResponse: str(client_id), ) return self.json_response("OK") + + if not sql_lab.cancel_query(query, g.user.username if g.user else None): + raise SupersetCancelQueryException("Could not cancel query") + query.status = QueryStatus.STOPPED db.session.commit() - sql_lab.cancel_query(query, g.user.username if g.user else None) return self.json_response("OK") diff --git a/tests/integration_tests/db_engine_specs/mysql_tests.py b/tests/integration_tests/db_engine_specs/mysql_tests.py index b5c2e53b996e7..b069bba69047a 100644 --- a/tests/integration_tests/db_engine_specs/mysql_tests.py +++ b/tests/integration_tests/db_engine_specs/mysql_tests.py @@ -251,4 +251,10 @@ def test_get_cancel_query_id(self, engine_mock): def test_cancel_query(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - assert MySQLEngineSpec.cancel_query(cursor_mock, query, 123) is None + assert MySQLEngineSpec.cancel_query(cursor_mock, query, 123) is True + + @unittest.mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query_failed(self, engine_mock): + query = Query() + cursor_mock = engine_mock.raiseError.side_effect = Exception() + assert MySQLEngineSpec.cancel_query(cursor_mock, query, 123) is False diff --git a/tests/integration_tests/db_engine_specs/postgres_tests.py b/tests/integration_tests/db_engine_specs/postgres_tests.py index 874ab80d9792a..4d03c5810ab20 100644 --- a/tests/integration_tests/db_engine_specs/postgres_tests.py +++ b/tests/integration_tests/db_engine_specs/postgres_tests.py @@ -455,7 +455,13 @@ def test_get_cancel_query_id(self, engine_mock): def test_cancel_query(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - assert PostgresEngineSpec.cancel_query(cursor_mock, query, 123) is None + assert PostgresEngineSpec.cancel_query(cursor_mock, query, 123) is True + + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query_failed(self, engine_mock): + query = Query() + cursor_mock = engine_mock.raiseError.side_effect = Exception() + assert PostgresEngineSpec.cancel_query(cursor_mock, query, 123) is False def test_base_parameters_mixin(): diff --git a/tests/integration_tests/db_engine_specs/snowflake_tests.py b/tests/integration_tests/db_engine_specs/snowflake_tests.py index 36cdb1e4c5596..2e74e0e68d16f 100644 --- a/tests/integration_tests/db_engine_specs/snowflake_tests.py +++ b/tests/integration_tests/db_engine_specs/snowflake_tests.py @@ -113,4 +113,10 @@ def test_get_cancel_query_id(self, engine_mock): def test_cancel_query(self, engine_mock): query = Query() cursor_mock = engine_mock.return_value.__enter__.return_value - assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, 123) is None + assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, 123) is True + + @mock.patch("sqlalchemy.engine.Engine.connect") + def test_cancel_query_failed(self, engine_mock): + query = Query() + cursor_mock = engine_mock.raiseError.side_effect = Exception() + assert SnowflakeEngineSpec.cancel_query(cursor_mock, query, 123) is False