This is an automated email from the ASF dual-hosted git repository. maximebeauchemin pushed a commit to branch ff_force_async in repository https://gitbox.apache.org/repos/asf/superset.git
commit ab2282f9dbfd0c3c8a7295b82aa89232dcb12494 Author: Maxime Beauchemin <[email protected]> AuthorDate: Wed Jun 12 12:55:12 2024 -0700 feat: add new SQLLAB_FORCE_RUN_ASYNC feature flag When setting up a celery backend to run asynchronous workloads, the administrator can set any given database connection to use (or not use) that backend. This PR introduced a new feature flag SQLLAB_FORCE_RUN_ASYNC (False by default) that disregards the database configuration, and forces SQL Lab to use the async backend. NOTE: superset/sqllab/sqllab_execution_context.py didn't have unit tests so I wrote started a new module to run basic tests around it. --- superset/config.py | 2 + superset/sqllab/sqllab_execution_context.py | 4 +- tests/unit_tests/sql_lab_execution_context.py | 78 +++++++++++++++++++++++++++ 3 files changed, 83 insertions(+), 1 deletion(-) diff --git a/superset/config.py b/superset/config.py index fa4ad4a29f..2b9dafb394 100644 --- a/superset/config.py +++ b/superset/config.py @@ -530,6 +530,8 @@ DEFAULT_FEATURE_FLAGS: dict[str, bool] = { "PLAYWRIGHT_REPORTS_AND_THUMBNAILS": False, # Set to True to enable experimental chart plugins "CHART_PLUGINS_EXPERIMENTAL": False, + # Regardless of database configuration settings, force SQLLAB to run async using Celery + "SQLLAB_FORCE_RUN_ASYNC": False, } # ------------------------------ diff --git a/superset/sqllab/sqllab_execution_context.py b/superset/sqllab/sqllab_execution_context.py index f67e2bf5a9..5ca180d101 100644 --- a/superset/sqllab/sqllab_execution_context.py +++ b/superset/sqllab/sqllab_execution_context.py @@ -77,7 +77,9 @@ class SqlJsonExecutionContext: # pylint: disable=too-many-instance-attributes self.schema = cast(str, query_params.get("schema")) self.sql = cast(str, query_params.get("sql")) self.template_params = self._get_template_params(query_params) - self.async_flag = cast(bool, query_params.get("runAsync")) + self.async_flag = is_feature_enabled("SQLLAB_FORCE_RUN_ASYNC") or cast( + bool, query_params.get("runAsync") + ) self.limit = self._get_limit_param(query_params) self.status = cast(str, query_params.get("status")) if cast(bool, query_params.get("select_as_cta")): diff --git a/tests/unit_tests/sql_lab_execution_context.py b/tests/unit_tests/sql_lab_execution_context.py new file mode 100644 index 0000000000..d7adf572c5 --- /dev/null +++ b/tests/unit_tests/sql_lab_execution_context.py @@ -0,0 +1,78 @@ +import pytest + +from superset.sqllab.sqllab_execution_context import ( + CreateTableAsSelect, + SqlJsonExecutionContext, +) +from tests.unit_tests.conftest import with_feature_flags + + [email protected] +def query_params(): + return { + "database_id": 1, + "catalog": "default", + "schema": "public", + "sql": "SELECT * FROM table", + "templateParams": "{}", + "runAsync": False, + "queryLimit": 1000, + "status": "success", + "select_as_cta": False, + "client_id": "client123", + "sql_editor_id": "editor123", + "tab": "tab123", + "expand_data": False, + } + + +def test_sql_json_execution_context_init(query_params): + context = SqlJsonExecutionContext(query_params) + assert context.database_id == 1 + assert context.catalog == "default" + assert context.schema == "public" + assert context.sql == "SELECT * FROM table" + assert context.template_params == {} + assert context.async_flag is False + assert context.limit == 1000 + assert context.status == "success" + assert context.client_id == "client123" + assert context.sql_editor_id == "editor123" + assert context.tab_name == "tab123" + assert context.expand_data is False + + +@with_feature_flags(SQLLAB_FORCE_RUN_ASYNC=True) [email protected]("runAsync, expected_async_flag", [(True, True), (False, True)]) +def test_sql_json_execution_context_feature_flag_false( + mocker, query_params, runAsync, expected_async_flag +): + query_params["runAsync"] = runAsync + context = SqlJsonExecutionContext(query_params) + assert context.async_flag == expected_async_flag + assert context.is_run_asynchronous() == expected_async_flag + + +@with_feature_flags(SQLLAB_FORCE_RUN_ASYNC=False) [email protected]( + "runAsync, expected_async_flag", [(True, True), (False, False)] +) +def test_sql_json_execution_context_feature_flag_true( + mocker, query_params, runAsync, expected_async_flag +): + query_params["runAsync"] = runAsync + context = SqlJsonExecutionContext(query_params) + assert context.async_flag == expected_async_flag + assert context.is_run_asynchronous() == expected_async_flag + + +def test_create_table_as_select(): + query_params = { + "ctas_method": "TABLE", + "schema": "public", + "tmp_table_name": "temp_table", + } + ctas = CreateTableAsSelect.create_from(query_params) + assert ctas.ctas_method == "TABLE" + assert ctas.target_schema_name == "public" + assert ctas.target_table_name == "temp_table"
