feat: cancel db query on stop (#15403)

* feat: cancel db query on stop

* fix pylint

* Add unit tests

* Do not bind multiple times

* Stop only running queries

* Postgres to cancel only the required query

* Remove extra log

* Add docstring

* Better types, docstring and naming

* Use python3 format strings

* Update superset/sql_lab.py

Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>

* Add cancel_query_on_windows_unload option to database

* Return cancel_query as bool

Co-authored-by: Beto Dealmeida <roberto@dealmeida.net>

tests/integration_tests/db_engine_specs/mysql_tests.py

@@ -21,6 +21,7 @@ from sqlalchemy.dialects.mysql import DATE, NVARCHAR, TEXT, VARCHAR
 
 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.integration_tests.db_engine_specs.base_tests import (
     assert_generic_types,
@@ -238,3 +239,22 @@ class TestMySQLEngineSpecsDbEngineSpec(TestDbEngineSpec):
                 },
             )
         ]
+
+    @unittest.mock.patch("sqlalchemy.engine.Engine.connect")
+    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_id(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 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

tests/integration_tests/db_engine_specs/postgres_tests.py

@@ -23,6 +23,7 @@ from sqlalchemy.dialects import postgresql
 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.integration_tests.db_engine_specs.base_tests import (
     assert_generic_types,
@@ -443,6 +444,25 @@ psql: error: could not connect to server: Operation timed out
             )
         ]
 
+    @mock.patch("sqlalchemy.engine.Engine.connect")
+    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 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, 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():
     parameters = {

tests/integration_tests/db_engine_specs/snowflake_tests.py

@@ -15,12 +15,14 @@
 # specific language governing permissions and limitations
 # under the License.
 import json
+from unittest import mock
 
 from sqlalchemy import column
 
 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.integration_tests.db_engine_specs.base_tests import TestDbEngineSpec
 
 
@@ -99,3 +101,22 @@ class TestSnowflakeDbEngineSpec(TestDbEngineSpec):
                 },
             )
         ]
+
+    @mock.patch("sqlalchemy.engine.Engine.connect")
+    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_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, 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