feat: add new SQLLAB_FORCE_RUN_ASYNC feature flag (#29231)

Co-authored-by: Geido <60598000+geido@users.noreply.github.com>

superset-frontend/packages/superset-ui-core/src/utils/featureFlags.ts

@@ -57,6 +57,7 @@ export enum FeatureFlag {
   TaggingSystem = 'TAGGING_SYSTEM',
   Thumbnails = 'THUMBNAILS',
   UseAnalagousColors = 'USE_ANALAGOUS_COLORS',
+  ForceSqlLabRunAsync = 'SQLLAB_FORCE_RUN_ASYNC',
 }
 
 export type ScheduleQueriesProps = {

superset-frontend/src/components/InfoTooltip/index.tsx

@@ -17,12 +17,12 @@
  * under the License.
  */
 
-import { styled } from '@superset-ui/core';
+import { styled, useTheme } from '@superset-ui/core';
 import { Tooltip } from 'src/components/Tooltip';
 import Icons from 'src/components/Icons';
 
 export interface InfoTooltipProps {
-  className?: string;
+  iconStyle?: React.CSSProperties;
   tooltip: string;
   placement?:
     | 'bottom'
@@ -46,9 +46,6 @@ export interface InfoTooltipProps {
 
 const StyledTooltip = styled(Tooltip)`
   cursor: pointer;
-  path:first-of-type {
-    fill: ${({ theme }) => theme.colors.grayscale.base};
-  }
 `;
 
 const StyledTooltipTitle = styled.span`
@@ -68,12 +65,18 @@ const defaultColor = 'rgba(0,0,0,0.9)';
 
 export default function InfoTooltip({
   tooltip,
+  iconStyle = {},
   placement = 'right',
   trigger = 'hover',
   overlayStyle = defaultOverlayStyle,
   bgColor = defaultColor,
   viewBox = '0 -1 24 24',
 }: InfoTooltipProps) {
+  const theme = useTheme();
+  const alteredIconStyle = {
+    ...iconStyle,
+    color: iconStyle.color || theme.colors.grayscale.base,
+  };
   return (
     <StyledTooltip
       title={<StyledTooltipTitle>{tooltip}</StyledTooltipTitle>}
@@ -82,7 +85,7 @@ export default function InfoTooltip({
       overlayStyle={overlayStyle}
       color={bgColor}
     >
-      <Icons.InfoSolidSmall className="info-solid-small" viewBox={viewBox} />
+      <Icons.InfoSolidSmall style={alteredIconStyle} viewBox={viewBox} />
     </StyledTooltip>
   );
 }

superset-frontend/src/features/databases/DatabaseModal/ExtraOptions.tsx

@@ -20,8 +20,11 @@ import { ChangeEvent, EventHandler } from 'react';
 import cx from 'classnames';
 import {
   t,
-  SupersetTheme,
   DatabaseConnectionExtension,
+  isFeatureEnabled,
+  SupersetTheme,
+  useTheme,
+  FeatureFlag,
 } from '@superset-ui/core';
 import InfoTooltip from 'src/components/InfoTooltip';
 import IndeterminateCheckbox from 'src/components/IndeterminateCheckbox';
@@ -68,10 +71,16 @@ const ExtraOptions = ({
     }
     return value;
   });
-
+  const theme = useTheme();
   const ExtraExtensionComponent = extraExtension?.component;
   const ExtraExtensionLogo = extraExtension?.logo;
   const ExtensionDescription = extraExtension?.description;
+  const allowRunAsync = isFeatureEnabled(FeatureFlag.ForceSqlLabRunAsync)
+    ? true
+    : !!db?.allow_run_async;
+  const isAllowRunAsyncDisabled = isFeatureEnabled(
+    FeatureFlag.ForceSqlLabRunAsync,
+  );
 
   return (
     <Collapse
@@ -319,7 +328,7 @@ const ExtraOptions = ({
             <IndeterminateCheckbox
               id="allow_run_async"
               indeterminate={false}
-              checked={!!db?.allow_run_async}
+              checked={allowRunAsync}
               onChange={onInputChange}
               labelText={t('Asynchronous query execution')}
             />
@@ -331,6 +340,14 @@ const ExtraOptions = ({
                   'backend. Refer to the installation docs for more information.',
               )}
             />
+            {isAllowRunAsyncDisabled && (
+              <InfoTooltip
+                iconStyle={{ color: theme.colors.error.base }}
+                tooltip={t(
+                  'This option has been disabled by the administrator.',
+                )}
+              />
+            )}
           </div>
         </StyledInputContainer>
         <StyledInputContainer css={{ no_margin_bottom }}>

superset/config.py

@@ -537,6 +537,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,
 }
 
 # ------------------------------

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")):

tests/unit_tests/sql_lab_execution_context.py

@@ -0,0 +1,96 @@
+# 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.
+# pylint: disable=import-outside-toplevel, invalid-name, unused-argument, too-many-locals
+
+import pytest
+
+from superset.sqllab.sqllab_execution_context import (
+    CreateTableAsSelect,
+    SqlJsonExecutionContext,
+)
+from tests.unit_tests.conftest import with_feature_flags
+
+
+@pytest.fixture
+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)
+@pytest.mark.parametrize("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)
+@pytest.mark.parametrize(
+    "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"