| # Licensed to the Apache Software Foundation (ASF) under one |
| # or more contributor license agreements. See the NOTICE file |
| # distributed with this work for additional information |
| # regarding copyright ownership. The ASF licenses this file |
| # to you under the Apache License, Version 2.0 (the |
| # "License"); you may not use this file except in compliance |
| # with the License. You may obtain a copy of the License at |
| # |
| # http://www.apache.org/licenses/LICENSE-2.0 |
| # |
| # Unless required by applicable law or agreed to in writing, |
| # software distributed under the License is distributed on an |
| # "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY |
| # KIND, either express or implied. See the License for the |
| # specific language governing permissions and limitations |
| # under the License. |
| # isort:skip_file |
| """Unit tests for Sql Lab""" |
| import json |
| from datetime import datetime, timedelta |
| from random import random |
| from unittest import mock |
| |
| from parameterized import parameterized |
| import prison |
| import pytest |
| |
| from superset import db, security_manager |
| from superset.connectors.sqla.models import SqlaTable |
| from superset.db_engine_specs import BaseEngineSpec |
| from superset.errors import ErrorLevel, SupersetErrorType |
| from superset.models.sql_lab import Query, SavedQuery |
| from superset.result_set import SupersetResultSet |
| from superset.sql_lab import execute_sql_statements, SqlLabException |
| from superset.sql_parse import CtasMethod |
| from superset.utils.core import ( |
| datetime_to_epoch, |
| get_example_database, |
| get_main_database, |
| ) |
| |
| from .base_tests import SupersetTestCase |
| from .conftest import CTAS_SCHEMA_NAME |
| |
| QUERY_1 = "SELECT * FROM birth_names LIMIT 1" |
| QUERY_2 = "SELECT * FROM NO_TABLE" |
| QUERY_3 = "SELECT * FROM birth_names LIMIT 10" |
| |
| |
| class TestSqlLab(SupersetTestCase): |
| """Testings for Sql Lab""" |
| |
| def run_some_queries(self): |
| db.session.query(Query).delete() |
| db.session.commit() |
| self.run_sql(QUERY_1, client_id="client_id_1", user_name="admin") |
| self.run_sql(QUERY_2, client_id="client_id_3", user_name="admin") |
| self.run_sql(QUERY_3, client_id="client_id_2", user_name="gamma_sqllab") |
| self.logout() |
| |
| def tearDown(self): |
| self.logout() |
| db.session.query(Query).delete() |
| db.session.commit() |
| db.session.close() |
| |
| def test_sql_json(self): |
| self.login("admin") |
| |
| data = self.run_sql("SELECT * FROM birth_names LIMIT 10", "1") |
| self.assertLess(0, len(data["data"])) |
| |
| data = self.run_sql("SELECT * FROM unexistant_table", "2") |
| assert ( |
| data["errors"][0]["error_type"] == SupersetErrorType.GENERIC_DB_ENGINE_ERROR |
| ) |
| assert data["errors"][0]["level"] == ErrorLevel.ERROR |
| assert data["errors"][0]["extra"] == { |
| "issue_codes": [ |
| { |
| "code": 1002, |
| "message": "Issue 1002 - The database returned an unexpected error.", |
| } |
| ] |
| } |
| |
| def test_sql_json_to_saved_query_info(self): |
| """ |
| SQLLab: Test SQLLab query execution info propagation to saved queries |
| """ |
| from freezegun import freeze_time |
| |
| self.login("admin") |
| |
| sql_statement = "SELECT * FROM birth_names LIMIT 10" |
| examples_db_id = get_example_database().id |
| saved_query = SavedQuery(db_id=examples_db_id, sql=sql_statement) |
| db.session.add(saved_query) |
| db.session.commit() |
| |
| with freeze_time("2020-01-01T00:00:00Z"): |
| self.run_sql(sql_statement, "1") |
| saved_query_ = ( |
| db.session.query(SavedQuery) |
| .filter( |
| SavedQuery.db_id == examples_db_id, SavedQuery.sql == sql_statement |
| ) |
| .one_or_none() |
| ) |
| assert saved_query_.rows is not None |
| assert saved_query_.last_run == datetime.now() |
| |
| # Rollback changes |
| db.session.delete(saved_query_) |
| db.session.commit() |
| |
| @parameterized.expand([CtasMethod.TABLE, CtasMethod.VIEW]) |
| def test_sql_json_cta_dynamic_db(self, ctas_method): |
| examples_db = get_example_database() |
| if examples_db.backend == "sqlite": |
| # sqlite doesn't support database creation |
| return |
| |
| with mock.patch( |
| "superset.views.core.get_cta_schema_name", |
| lambda d, u, s, sql: f"{u.username}_database", |
| ): |
| old_allow_ctas = examples_db.allow_ctas |
| examples_db.allow_ctas = True # enable cta |
| |
| self.login("admin") |
| tmp_table_name = f"test_target_{ctas_method.lower()}" |
| self.run_sql( |
| "SELECT * FROM birth_names", |
| "1", |
| database_name="examples", |
| tmp_table_name=tmp_table_name, |
| select_as_cta=True, |
| ctas_method=ctas_method, |
| ) |
| |
| # assertions |
| db.session.commit() |
| examples_db = get_example_database() |
| engine = examples_db.get_sqla_engine() |
| data = engine.execute( |
| f"SELECT * FROM admin_database.{tmp_table_name}" |
| ).fetchall() |
| self.assertEqual( |
| 100, len(data) |
| ) # SQL_MAX_ROW not applied due to the SQLLAB_CTAS_NO_LIMIT set to True |
| |
| # cleanup |
| engine.execute(f"DROP {ctas_method} admin_database.{tmp_table_name}") |
| examples_db.allow_ctas = old_allow_ctas |
| db.session.commit() |
| |
| def test_multi_sql(self): |
| self.login("admin") |
| |
| multi_sql = """ |
| SELECT * FROM birth_names LIMIT 1; |
| SELECT * FROM birth_names LIMIT 2; |
| """ |
| data = self.run_sql(multi_sql, "2234") |
| self.assertLess(0, len(data["data"])) |
| |
| def test_explain(self): |
| self.login("admin") |
| |
| data = self.run_sql("EXPLAIN SELECT * FROM birth_names", "1") |
| self.assertLess(0, len(data["data"])) |
| |
| def test_sql_json_has_access(self): |
| examples_db = get_example_database() |
| examples_db_permission_view = security_manager.add_permission_view_menu( |
| "database_access", examples_db.perm |
| ) |
| astronaut = security_manager.add_role("ExampleDBAccess") |
| security_manager.add_permission_role(astronaut, examples_db_permission_view) |
| # Gamma user, with sqllab and db permission |
| self.create_user_with_roles("Gagarin", ["ExampleDBAccess", "Gamma", "sql_lab"]) |
| |
| data = self.run_sql(QUERY_1, "1", user_name="Gagarin") |
| db.session.query(Query).delete() |
| db.session.commit() |
| self.assertLess(0, len(data["data"])) |
| |
| def test_sql_json_schema_access(self): |
| examples_db = get_example_database() |
| db_backend = examples_db.backend |
| if db_backend == "sqlite": |
| # sqlite doesn't support database creation |
| return |
| |
| sqllab_test_db_schema_permission_view = security_manager.add_permission_view_menu( |
| "schema_access", f"[{examples_db.name}].[{CTAS_SCHEMA_NAME}]" |
| ) |
| schema_perm_role = security_manager.add_role("SchemaPermission") |
| security_manager.add_permission_role( |
| schema_perm_role, sqllab_test_db_schema_permission_view |
| ) |
| self.create_user_with_roles( |
| "SchemaUser", ["SchemaPermission", "Gamma", "sql_lab"] |
| ) |
| |
| examples_db.get_sqla_engine().execute( |
| f"CREATE TABLE IF NOT EXISTS {CTAS_SCHEMA_NAME}.test_table AS SELECT 1 as c1, 2 as c2" |
| ) |
| |
| data = self.run_sql( |
| f"SELECT * FROM {CTAS_SCHEMA_NAME}.test_table", "3", user_name="SchemaUser" |
| ) |
| self.assertEqual(1, len(data["data"])) |
| |
| data = self.run_sql( |
| f"SELECT * FROM {CTAS_SCHEMA_NAME}.test_table", |
| "4", |
| user_name="SchemaUser", |
| schema=CTAS_SCHEMA_NAME, |
| ) |
| self.assertEqual(1, len(data["data"])) |
| |
| # postgres needs a schema as a part of the table name. |
| if db_backend == "mysql": |
| data = self.run_sql( |
| "SELECT * FROM test_table", |
| "5", |
| user_name="SchemaUser", |
| schema=CTAS_SCHEMA_NAME, |
| ) |
| self.assertEqual(1, len(data["data"])) |
| |
| db.session.query(Query).delete() |
| get_example_database().get_sqla_engine().execute( |
| f"DROP TABLE IF EXISTS {CTAS_SCHEMA_NAME}.test_table" |
| ) |
| db.session.commit() |
| |
| def test_queries_endpoint(self): |
| self.run_some_queries() |
| |
| # Not logged in, should error out |
| resp = self.client.get("/superset/queries/0") |
| # Redirects to the login page |
| self.assertEqual(403, resp.status_code) |
| |
| # Admin sees queries |
| self.login("admin") |
| data = self.get_json_resp("/superset/queries/0") |
| self.assertEqual(2, len(data)) |
| data = self.get_json_resp("/superset/queries/0.0") |
| self.assertEqual(2, len(data)) |
| |
| # Run 2 more queries |
| self.run_sql("SELECT * FROM birth_names LIMIT 1", client_id="client_id_4") |
| self.run_sql("SELECT * FROM birth_names LIMIT 2", client_id="client_id_5") |
| self.login("admin") |
| data = self.get_json_resp("/superset/queries/0") |
| self.assertEqual(4, len(data)) |
| |
| now = datetime.now() + timedelta(days=1) |
| query = ( |
| db.session.query(Query) |
| .filter_by(sql="SELECT * FROM birth_names LIMIT 1") |
| .first() |
| ) |
| query.changed_on = now |
| db.session.commit() |
| |
| data = self.get_json_resp( |
| "/superset/queries/{}".format(float(datetime_to_epoch(now)) - 1000) |
| ) |
| self.assertEqual(1, len(data)) |
| |
| self.logout() |
| resp = self.client.get("/superset/queries/0") |
| # Redirects to the login page |
| self.assertEqual(403, resp.status_code) |
| |
| def test_search_query_on_db_id(self): |
| self.run_some_queries() |
| self.login("admin") |
| examples_dbid = get_example_database().id |
| |
| # Test search queries on database Id |
| data = self.get_json_resp( |
| f"/superset/search_queries?database_id={examples_dbid}" |
| ) |
| self.assertEqual(3, len(data)) |
| db_ids = [k["dbId"] for k in data] |
| self.assertEqual([examples_dbid for i in range(3)], db_ids) |
| |
| resp = self.get_resp("/superset/search_queries?database_id=-1") |
| data = json.loads(resp) |
| self.assertEqual(0, len(data)) |
| |
| def test_search_query_on_user(self): |
| self.run_some_queries() |
| self.login("admin") |
| |
| # Test search queries on user Id |
| user_id = security_manager.find_user("admin").id |
| data = self.get_json_resp("/superset/search_queries?user_id={}".format(user_id)) |
| self.assertEqual(2, len(data)) |
| user_ids = {k["userId"] for k in data} |
| self.assertEqual(set([user_id]), user_ids) |
| |
| user_id = security_manager.find_user("gamma_sqllab").id |
| resp = self.get_resp("/superset/search_queries?user_id={}".format(user_id)) |
| data = json.loads(resp) |
| self.assertEqual(1, len(data)) |
| self.assertEqual(data[0]["userId"], user_id) |
| |
| def test_search_query_on_status(self): |
| self.run_some_queries() |
| self.login("admin") |
| # Test search queries on status |
| resp = self.get_resp("/superset/search_queries?status=success") |
| data = json.loads(resp) |
| self.assertEqual(2, len(data)) |
| states = [k["state"] for k in data] |
| self.assertEqual(["success", "success"], states) |
| |
| resp = self.get_resp("/superset/search_queries?status=failed") |
| data = json.loads(resp) |
| self.assertEqual(1, len(data)) |
| self.assertEqual(data[0]["state"], "failed") |
| |
| def test_search_query_on_text(self): |
| self.run_some_queries() |
| self.login("admin") |
| url = "/superset/search_queries?search_text=birth" |
| data = self.get_json_resp(url) |
| self.assertEqual(2, len(data)) |
| self.assertIn("birth", data[0]["sql"]) |
| |
| def test_search_query_on_time(self): |
| self.run_some_queries() |
| self.login("admin") |
| first_query_time = ( |
| db.session.query(Query).filter_by(sql=QUERY_1).one() |
| ).start_time |
| second_query_time = ( |
| db.session.query(Query).filter_by(sql=QUERY_3).one() |
| ).start_time |
| # Test search queries on time filter |
| from_time = "from={}".format(int(first_query_time)) |
| to_time = "to={}".format(int(second_query_time)) |
| params = [from_time, to_time] |
| resp = self.get_resp("/superset/search_queries?" + "&".join(params)) |
| data = json.loads(resp) |
| self.assertEqual(2, len(data)) |
| |
| def test_search_query_only_owned(self) -> None: |
| """ |
| Test a search query with a user that does not have can_access_all_queries. |
| """ |
| # Test search_queries for Alpha user |
| self.run_some_queries() |
| self.login("gamma_sqllab") |
| |
| user_id = security_manager.find_user("gamma_sqllab").id |
| data = self.get_json_resp("/superset/search_queries") |
| |
| self.assertEqual(1, len(data)) |
| user_ids = {k["userId"] for k in data} |
| self.assertEqual(set([user_id]), user_ids) |
| |
| def test_alias_duplicate(self): |
| self.run_sql( |
| "SELECT name as col, gender as col FROM birth_names LIMIT 10", |
| client_id="2e2df3", |
| user_name="admin", |
| raise_on_error=True, |
| ) |
| |
| def test_ps_conversion_no_dict(self): |
| cols = [["string_col", "string"], ["int_col", "int"], ["float_col", "float"]] |
| data = [["a", 4, 4.0]] |
| results = SupersetResultSet(data, cols, BaseEngineSpec) |
| |
| self.assertEqual(len(data), results.size) |
| self.assertEqual(len(cols), len(results.columns)) |
| |
| def test_pa_conversion_tuple(self): |
| cols = ["string_col", "int_col", "list_col", "float_col"] |
| data = [("Text", 111, [123], 1.0)] |
| results = SupersetResultSet(data, cols, BaseEngineSpec) |
| |
| self.assertEqual(len(data), results.size) |
| self.assertEqual(len(cols), len(results.columns)) |
| |
| def test_pa_conversion_dict(self): |
| cols = ["string_col", "dict_col", "int_col"] |
| data = [["a", {"c1": 1, "c2": 2, "c3": 3}, 4]] |
| results = SupersetResultSet(data, cols, BaseEngineSpec) |
| |
| self.assertEqual(len(data), results.size) |
| self.assertEqual(len(cols), len(results.columns)) |
| |
| def test_sqllab_viz(self): |
| self.login("admin") |
| examples_dbid = get_example_database().id |
| payload = { |
| "chartType": "dist_bar", |
| "datasourceName": f"test_viz_flow_table_{random()}", |
| "schema": "superset", |
| "columns": [ |
| {"is_date": False, "type": "STRING", "name": f"viz_type_{random()}"}, |
| {"is_date": False, "type": "OBJECT", "name": f"ccount_{random()}"}, |
| ], |
| "sql": """\ |
| SELECT * |
| FROM birth_names |
| LIMIT 10""", |
| "dbId": examples_dbid, |
| } |
| data = {"data": json.dumps(payload)} |
| resp = self.get_json_resp("/superset/sqllab_viz/", data=data) |
| self.assertIn("table_id", resp) |
| |
| # ensure owner is set correctly |
| table_id = resp["table_id"] |
| table = db.session.query(SqlaTable).filter_by(id=table_id).one() |
| self.assertEqual([owner.username for owner in table.owners], ["admin"]) |
| view_menu = security_manager.find_view_menu(table.get_perm()) |
| assert view_menu is not None |
| |
| # Cleanup |
| db.session.delete(table) |
| db.session.commit() |
| |
| def test_sqllab_viz_bad_payload(self): |
| self.login("admin") |
| payload = { |
| "chartType": "dist_bar", |
| "schema": "superset", |
| "columns": [ |
| {"is_date": False, "type": "STRING", "name": f"viz_type_{random()}"}, |
| {"is_date": False, "type": "OBJECT", "name": f"ccount_{random()}"}, |
| ], |
| "sql": """\ |
| SELECT * |
| FROM birth_names |
| LIMIT 10""", |
| } |
| data = {"data": json.dumps(payload)} |
| url = "/superset/sqllab_viz/" |
| response = self.client.post(url, data=data, follow_redirects=True) |
| assert response.status_code == 400 |
| |
| def test_sqllab_table_viz(self): |
| self.login("admin") |
| examples_db = get_example_database() |
| examples_db.get_sqla_engine().execute( |
| "DROP TABLE IF EXISTS test_sqllab_table_viz" |
| ) |
| examples_db.get_sqla_engine().execute( |
| "CREATE TABLE test_sqllab_table_viz AS SELECT 2 as col" |
| ) |
| examples_dbid = examples_db.id |
| |
| payload = { |
| "datasourceName": "test_sqllab_table_viz", |
| "columns": [], |
| "dbId": examples_dbid, |
| } |
| |
| data = {"data": json.dumps(payload)} |
| resp = self.get_json_resp("/superset/get_or_create_table/", data=data) |
| self.assertIn("table_id", resp) |
| |
| # ensure owner is set correctly |
| table_id = resp["table_id"] |
| table = db.session.query(SqlaTable).filter_by(id=table_id).one() |
| self.assertEqual([owner.username for owner in table.owners], ["admin"]) |
| db.session.delete(table) |
| get_example_database().get_sqla_engine().execute( |
| "DROP TABLE test_sqllab_table_viz" |
| ) |
| db.session.commit() |
| |
| def test_sql_limit(self): |
| self.login("admin") |
| test_limit = 1 |
| data = self.run_sql("SELECT * FROM birth_names", client_id="sql_limit_1") |
| self.assertGreater(len(data["data"]), test_limit) |
| data = self.run_sql( |
| "SELECT * FROM birth_names", client_id="sql_limit_2", query_limit=test_limit |
| ) |
| self.assertEqual(len(data["data"]), test_limit) |
| data = self.run_sql( |
| "SELECT * FROM birth_names LIMIT {}".format(test_limit), |
| client_id="sql_limit_3", |
| query_limit=test_limit + 1, |
| ) |
| self.assertEqual(len(data["data"]), test_limit) |
| data = self.run_sql( |
| "SELECT * FROM birth_names LIMIT {}".format(test_limit + 1), |
| client_id="sql_limit_4", |
| query_limit=test_limit, |
| ) |
| self.assertEqual(len(data["data"]), test_limit) |
| |
| def test_query_api_filter(self) -> None: |
| """ |
| Test query api without can_only_access_owned_queries perm added to |
| Admin and make sure all queries show up. |
| """ |
| self.run_some_queries() |
| self.login(username="admin") |
| |
| url = "/api/v1/query/" |
| data = self.get_json_resp(url) |
| admin = security_manager.find_user("admin") |
| gamma_sqllab = security_manager.find_user("gamma_sqllab") |
| self.assertEqual(3, len(data["result"])) |
| user_queries = [result.get("user").get("username") for result in data["result"]] |
| assert admin.username in user_queries |
| assert gamma_sqllab.username in user_queries |
| |
| def test_query_api_can_access_all_queries(self) -> None: |
| """ |
| Test query api with can_access_all_queries perm added to |
| gamma and make sure all queries show up. |
| """ |
| session = db.session |
| |
| # Add all_query_access perm to Gamma user |
| all_queries_view = security_manager.find_permission_view_menu( |
| "all_query_access", "all_query_access" |
| ) |
| |
| security_manager.add_permission_role( |
| security_manager.find_role("gamma_sqllab"), all_queries_view |
| ) |
| session.commit() |
| |
| # Test search_queries for Admin user |
| self.run_some_queries() |
| self.login("gamma_sqllab") |
| url = "/api/v1/query/" |
| data = self.get_json_resp(url) |
| self.assertEqual(3, len(data["result"])) |
| |
| # Remove all_query_access from gamma sqllab |
| all_queries_view = security_manager.find_permission_view_menu( |
| "all_query_access", "all_query_access" |
| ) |
| security_manager.del_permission_role( |
| security_manager.find_role("gamma_sqllab"), all_queries_view |
| ) |
| |
| session.commit() |
| |
| def test_query_admin_can_access_all_queries(self) -> None: |
| """ |
| Test query api with all_query_access perm added to |
| Admin and make sure only Admin queries show up. This is the default |
| """ |
| # Test search_queries for Admin user |
| self.run_some_queries() |
| self.login("admin") |
| |
| url = "/api/v1/query/" |
| data = self.get_json_resp(url) |
| self.assertEqual(3, len(data["result"])) |
| |
| def test_api_database(self): |
| self.login("admin") |
| self.create_fake_db() |
| get_example_database() |
| get_main_database() |
| |
| arguments = { |
| "keys": [], |
| "filters": [{"col": "expose_in_sqllab", "opr": "eq", "value": True}], |
| "order_column": "database_name", |
| "order_direction": "asc", |
| "page": 0, |
| "page_size": -1, |
| } |
| url = f"api/v1/database/?q={prison.dumps(arguments)}" |
| |
| self.assertEqual( |
| {"examples", "fake_db_100", "main"}, |
| {r.get("database_name") for r in self.get_json_resp(url)["result"]}, |
| ) |
| self.delete_fake_db() |
| |
| @mock.patch.dict( |
| "superset.extensions.feature_flag_manager._feature_flags", |
| {"ENABLE_TEMPLATE_PROCESSING": True}, |
| clear=True, |
| ) |
| def test_sql_json_parameter_error(self): |
| self.login("admin") |
| |
| data = self.run_sql( |
| "SELECT * FROM birth_names WHERE state = '{{ state }}' LIMIT 10", |
| "1", |
| template_params=json.dumps({"state": "CA"}), |
| ) |
| assert data["status"] == "success" |
| |
| data = self.run_sql( |
| "SELECT * FROM birth_names WHERE state = '{{ stat }}' LIMIT 10", |
| "2", |
| template_params=json.dumps({"state": "CA"}), |
| ) |
| assert data["errors"][0]["error_type"] == "MISSING_TEMPLATE_PARAMS_ERROR" |
| assert data["errors"][0]["extra"] == { |
| "issue_codes": [ |
| { |
| "code": 1006, |
| "message": "Issue 1006 - One or more parameters specified in the query are missing.", |
| } |
| ], |
| "template_parameters": {"state": "CA"}, |
| "undefined_parameters": ["stat"], |
| } |
| |
| @mock.patch("superset.sql_lab.get_query") |
| @mock.patch("superset.sql_lab.execute_sql_statement") |
| def test_execute_sql_statements(self, mock_execute_sql_statement, mock_get_query): |
| sql = """ |
| -- comment |
| SET @value = 42; |
| SELECT @value AS foo; |
| -- comment |
| """ |
| mock_session = mock.MagicMock() |
| mock_query = mock.MagicMock() |
| mock_query.database.allow_run_async = False |
| mock_cursor = mock.MagicMock() |
| mock_query.database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value = ( |
| mock_cursor |
| ) |
| mock_query.database.db_engine_spec.run_multiple_statements_as_one = False |
| mock_get_query.return_value = mock_query |
| |
| execute_sql_statements( |
| query_id=1, |
| rendered_query=sql, |
| return_results=True, |
| store_results=False, |
| user_name="admin", |
| session=mock_session, |
| start_time=None, |
| expand_data=False, |
| log_params=None, |
| ) |
| mock_execute_sql_statement.assert_has_calls( |
| [ |
| mock.call( |
| "SET @value = 42", |
| mock_query, |
| "admin", |
| mock_session, |
| mock_cursor, |
| None, |
| False, |
| ), |
| mock.call( |
| "SELECT @value AS foo", |
| mock_query, |
| "admin", |
| mock_session, |
| mock_cursor, |
| None, |
| False, |
| ), |
| ] |
| ) |
| |
| @mock.patch("superset.sql_lab.get_query") |
| @mock.patch("superset.sql_lab.execute_sql_statement") |
| def test_execute_sql_statements_ctas( |
| self, mock_execute_sql_statement, mock_get_query |
| ): |
| sql = """ |
| -- comment |
| SET @value = 42; |
| SELECT @value AS foo; |
| -- comment |
| """ |
| mock_session = mock.MagicMock() |
| mock_query = mock.MagicMock() |
| mock_query.database.allow_run_async = False |
| mock_cursor = mock.MagicMock() |
| mock_query.database.get_sqla_engine.return_value.raw_connection.return_value.cursor.return_value = ( |
| mock_cursor |
| ) |
| mock_query.database.db_engine_spec.run_multiple_statements_as_one = False |
| mock_get_query.return_value = mock_query |
| |
| # set the query to CTAS |
| mock_query.select_as_cta = True |
| mock_query.ctas_method = CtasMethod.TABLE |
| |
| execute_sql_statements( |
| query_id=1, |
| rendered_query=sql, |
| return_results=True, |
| store_results=False, |
| user_name="admin", |
| session=mock_session, |
| start_time=None, |
| expand_data=False, |
| log_params=None, |
| ) |
| mock_execute_sql_statement.assert_has_calls( |
| [ |
| mock.call( |
| "SET @value = 42", |
| mock_query, |
| "admin", |
| mock_session, |
| mock_cursor, |
| None, |
| False, |
| ), |
| mock.call( |
| "SELECT @value AS foo", |
| mock_query, |
| "admin", |
| mock_session, |
| mock_cursor, |
| None, |
| True, # apply_ctas |
| ), |
| ] |
| ) |
| |
| # try invalid CTAS |
| sql = "DROP TABLE my_table" |
| with pytest.raises(SqlLabException) as excinfo: |
| execute_sql_statements( |
| query_id=1, |
| rendered_query=sql, |
| return_results=True, |
| store_results=False, |
| user_name="admin", |
| session=mock_session, |
| start_time=None, |
| expand_data=False, |
| log_params=None, |
| ) |
| assert str(excinfo.value) == ( |
| "CTAS (create table as select) can only be run with " |
| "a query where the last statement is a SELECT. Please " |
| "make sure your query has a SELECT as its last " |
| "statement. Then, try running your query again." |
| ) |
| |
| # try invalid CVAS |
| mock_query.ctas_method = CtasMethod.VIEW |
| sql = """ |
| -- comment |
| SET @value = 42; |
| SELECT @value AS foo; |
| -- comment |
| """ |
| with pytest.raises(SqlLabException) as excinfo: |
| execute_sql_statements( |
| query_id=1, |
| rendered_query=sql, |
| return_results=True, |
| store_results=False, |
| user_name="admin", |
| session=mock_session, |
| start_time=None, |
| expand_data=False, |
| log_params=None, |
| ) |
| assert str(excinfo.value) == ( |
| "CVAS (create view as select) can only be run with a " |
| "query with a single SELECT statement. Please make " |
| "sure your query has only a SELECT statement. Then, " |
| "try running your query again." |
| ) |
