fix(extensions): add cache headers and strip Vary: Cookie for extension static assets (#40120)

Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
(cherry picked from commit a06e6ea19b)

superset-frontend/src/SqlLab/components/EstimateQueryCostButton/index.tsx

@@ -113,7 +113,7 @@ const EstimateQueryCostButton = ({
         modalBody={renderModalBody()}
         triggerNode={
           <Button
-            color="primary"
+            color="default"
             variant="text"
             style={{ height: 32, padding: '4px 15px' }}
             onClick={onClickHandler}

superset-frontend/src/extensions/ExtensionsLoader.test.ts

@@ -16,6 +16,7 @@
  * specific language governing permissions and limitations
  * under the License.
  */
+import { SupersetClient } from '@superset-ui/core';
 import { logging } from '@apache-superset/core/utils';
 import type { common as core } from '@apache-superset/core';
 import ExtensionsLoader from './ExtensionsLoader';
@@ -111,3 +112,33 @@ test('handles initialization errors gracefully', async () => {
   errorSpy.mockRestore();
   appendChildSpy.mockRestore();
 });
+
+test('logs success after initializeExtensions completes', async () => {
+  const loader = ExtensionsLoader.getInstance();
+  const infoSpy = jest.spyOn(logging, 'info').mockImplementation();
+  jest.spyOn(SupersetClient, 'get').mockResolvedValue({
+    json: { result: [] },
+  } as any);
+
+  await loader.initializeExtensions();
+
+  expect(infoSpy).toHaveBeenCalledWith('Extensions initialized successfully.');
+
+  infoSpy.mockRestore();
+});
+
+test('logs error when initializeExtensions fails', async () => {
+  const loader = ExtensionsLoader.getInstance();
+  const errorSpy = jest.spyOn(logging, 'error').mockImplementation();
+  const fetchError = new Error('Network error');
+  jest.spyOn(SupersetClient, 'get').mockRejectedValue(fetchError);
+
+  await loader.initializeExtensions();
+
+  expect(errorSpy).toHaveBeenCalledWith(
+    'Error setting up extensions:',
+    fetchError,
+  );
+
+  errorSpy.mockRestore();
+});

superset-frontend/src/extensions/ExtensionsLoader.ts

@@ -34,6 +34,8 @@ class ExtensionsLoader {
 
   private extensionIndex: Map<string, Extension> = new Map();
 
+  private initializationPromise: Promise<void> | null = null;
+
   // eslint-disable-next-line no-useless-constructor
   private constructor() {
     // Private constructor for singleton pattern
@@ -54,16 +56,27 @@ class ExtensionsLoader {
    * Initializes extensions by fetching the list from the API and loading each one.
    * @throws Error if initialization fails.
    */
-  public async initializeExtensions(): Promise<void> {
-    const response = await SupersetClient.get({
-      endpoint: '/api/v1/extensions/',
-    });
-    const extensions: Extension[] = response.json.result;
-    await Promise.all(
-      extensions.map(async extension => {
-        await this.initializeExtension(extension);
-      }),
-    );
+  public initializeExtensions(): Promise<void> {
+    if (this.initializationPromise) {
+      return this.initializationPromise;
+    }
+    this.initializationPromise = (async () => {
+      try {
+        const response = await SupersetClient.get({
+          endpoint: '/api/v1/extensions/',
+        });
+        const extensions: Extension[] = response.json.result;
+        await Promise.all(
+          extensions.map(async extension => {
+            await this.initializeExtension(extension);
+          }),
+        );
+        logging.info('Extensions initialized successfully.');
+      } catch (error) {
+        logging.error('Error setting up extensions:', error);
+      }
+    })();
+    return this.initializationPromise;
   }
 
   /**

superset-frontend/src/extensions/ExtensionsStartup.test.tsx

@@ -18,7 +18,6 @@
  */
 import { render, waitFor } from 'spec/helpers/testing-library';
 import { FeatureFlag, isFeatureEnabled } from '@superset-ui/core';
-import { logging } from '@apache-superset/core/utils';
 import fetchMock from 'fetch-mock';
 import ExtensionsStartup from './ExtensionsStartup';
 import ExtensionsLoader from './ExtensionsLoader';
@@ -192,14 +191,12 @@ test('only initializes once even with multiple renders', async () => {
   loader.initializeExtensions = originalInitialize;
 });
 
-test('initializes ExtensionsLoader and logs success when EnableExtensions feature flag is enabled', async () => {
+test('initializes ExtensionsLoader when EnableExtensions feature flag is enabled', async () => {
   // Ensure feature flag is enabled
   mockIsFeatureEnabled.mockImplementation(
     (flag: FeatureFlag) => flag === FeatureFlag.EnableExtensions,
   );
 
-  const infoSpy = jest.spyOn(logging, 'info').mockImplementation();
-
   // Mock the initializeExtensions method to succeed
   const originalInitialize = ExtensionsLoader.prototype.initializeExtensions;
   ExtensionsLoader.prototype.initializeExtensions = jest
@@ -220,15 +217,10 @@ test('initializes ExtensionsLoader and logs success when EnableExtensions featur
     expect(
       ExtensionsLoader.prototype.initializeExtensions,
     ).toHaveBeenCalledTimes(1);
-    // Verify success message was logged
-    expect(infoSpy).toHaveBeenCalledWith(
-      'Extensions initialized successfully.',
-    );
   });
 
   // Restore original method
   ExtensionsLoader.prototype.initializeExtensions = originalInitialize;
-  infoSpy.mockRestore();
 });
 
 test('does not initialize ExtensionsLoader when EnableExtensions feature flag is disabled', async () => {
@@ -259,38 +251,36 @@ test('does not initialize ExtensionsLoader when EnableExtensions feature flag is
   initializeSpy.mockRestore();
 });
 
-test('logs error when ExtensionsLoader initialization fails', async () => {
+test('continues rendering children even when ExtensionsLoader initialization fails', async () => {
   // Ensure feature flag is enabled
   mockIsFeatureEnabled.mockReturnValue(true);
 
-  const errorSpy = jest.spyOn(logging, 'error').mockImplementation();
-
-  // Mock the initializeExtensions method to throw an error
+  // Mock the initializeExtensions method to reject — ExtensionsLoader handles
+  // its own error logging internally
   const originalInitialize = ExtensionsLoader.prototype.initializeExtensions;
   ExtensionsLoader.prototype.initializeExtensions = jest
     .fn()
-    .mockImplementation(() => {
-      throw new Error('Test initialization error');
-    });
+    .mockImplementation(() => Promise.resolve());
 
-  render(<ExtensionsStartup />, {
-    useRedux: true,
-    initialState: mockInitialState,
-  });
+  const { container } = render(
+    <ExtensionsStartup>
+      <div data-testid="child" />
+    </ExtensionsStartup>,
+    {
+      useRedux: true,
+      initialState: mockInitialState,
+    },
+  );
 
   await waitFor(() => {
-    // Verify feature flag was checked
     expect(mockIsFeatureEnabled).toHaveBeenCalledWith(
       FeatureFlag.EnableExtensions,
     );
-    // Verify error was logged
-    expect(errorSpy).toHaveBeenCalledWith(
-      'Error setting up extensions:',
-      expect.any(Error),
-    );
+    expect(
+      container.querySelector('[data-testid="child"]'),
+    ).toBeInTheDocument();
   });
 
   // Restore original method
   ExtensionsLoader.prototype.initializeExtensions = originalInitialize;
-  errorSpy.mockRestore();
 });

superset-frontend/src/extensions/ExtensionsStartup.tsx

@@ -82,17 +82,7 @@ const ExtensionsStartup: React.FC<{ children?: React.ReactNode }> = ({
 
     const setup = async () => {
       if (isFeatureEnabled(FeatureFlag.EnableExtensions)) {
-        try {
-          await ExtensionsLoader.getInstance().initializeExtensions();
-          supersetCore.utils.logging.info(
-            'Extensions initialized successfully.',
-          );
-        } catch (error) {
-          supersetCore.utils.logging.error(
-            'Error setting up extensions:',
-            error,
-          );
-        }
+        await ExtensionsLoader.getInstance().initializeExtensions();
       }
       setInitialized(true);
     };

superset/app.py

@@ -36,6 +36,7 @@ else:
 from flask import Flask, Response
 from werkzeug.exceptions import NotFound
 
+from superset.extensions.cache_middleware import ExtensionCacheMiddleware
 from superset.extensions.local_extensions_watcher import (
     start_local_extensions_watcher_thread,
 )
@@ -44,7 +45,7 @@ from superset.initialization import SupersetAppInitializer
 logger = logging.getLogger(__name__)
 
 
-def create_app(
+def create_app(  # noqa: C901
     superset_config_module: Optional[str] = None,
     superset_app_root: Optional[str] = None,
 ) -> Flask:
@@ -66,7 +67,6 @@ def create_app(
             or app.config["APPLICATION_ROOT"],
         )
         if app_root != "/":
-            app.wsgi_app = AppRootMiddleware(app.wsgi_app, app_root)
             # If not set, manually configure options that depend on the
             # value of app_root so things work out of the box
             if not app.config["STATIC_ASSETS_PREFIX"]:
@@ -93,6 +93,13 @@ def create_app(
         app_initializer = app.config.get("APP_INITIALIZER", SupersetAppInitializer)(app)
         app_initializer.init_app()
 
+        # Must be applied before AppRootMiddleware so the path prefix
+        # is stripped before the extension asset path regex runs.
+        app.wsgi_app = ExtensionCacheMiddleware(app.wsgi_app)
+
+        if app_root != "/":
+            app.wsgi_app = AppRootMiddleware(app.wsgi_app, app_root)
+
         # Set up LOCAL_EXTENSIONS file watcher when in debug mode
         if app.debug:
             start_local_extensions_watcher_thread(app)

superset/extensions/api.py

@@ -225,4 +225,9 @@ class ExtensionsRestApi(BaseApi):
         if not mimetype:
             mimetype = "application/octet-stream"
 
-        return send_file(BytesIO(chunk), mimetype=mimetype)
+        response = send_file(BytesIO(chunk), mimetype=mimetype)
+        # Chunk filenames include a content hash, so they are immutable.
+        response.cache_control.max_age = 31536000
+        response.cache_control.public = True
+        response.cache_control.immutable = True
+        return response

superset/extensions/cache_middleware.py

@@ -0,0 +1,73 @@
+# 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.
+from __future__ import annotations
+
+import re
+from types import TracebackType
+from typing import Callable, Iterable, TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from _typeshed.wsgi import StartResponse, WSGIApplication, WSGIEnvironment
+
+# Matches only the static asset endpoint: /api/v1/extensions/<publisher>/<name>/<file>
+# Does not match the list (/), get (/<publisher>/<name>), or info (/_info) endpoints.
+_ASSET_PATH_RE = re.compile(r"^/api/v1/extensions/[^/]+/[^/]+/[^/]+$")
+
+
+class ExtensionCacheMiddleware:
+    """Strip 'Cookie' from the Vary header on extension asset responses.
+
+    Flask's session interface appends Vary: Cookie unconditionally after every
+    after_request hook runs, so it cannot be removed at the view layer. This
+    middleware intercepts the WSGI response at the lowest level, after all
+    Flask processing is complete.
+    """
+
+    def __init__(self, wsgi_app: WSGIApplication) -> None:
+        self.wsgi_app = wsgi_app
+
+    def __call__(
+        self, environ: WSGIEnvironment, start_response: StartResponse
+    ) -> Iterable[bytes]:
+        path = environ.get("PATH_INFO", "")
+        if not _ASSET_PATH_RE.match(path):
+            return self.wsgi_app(environ, start_response)
+
+        def patched_start_response(
+            status: str,
+            response_headers: list[tuple[str, str]],
+            exc_info: (
+                tuple[type[BaseException], BaseException, TracebackType]
+                | tuple[None, None, None]
+                | None
+            ) = None,
+        ) -> Callable[[bytes], object]:
+            new_headers = []
+            for name, value in response_headers:
+                if name.lower() == "vary":
+                    parts = [
+                        v.strip()
+                        for v in value.split(",")
+                        if v.strip().lower() != "cookie"
+                    ]
+                    if parts:
+                        new_headers.append((name, ", ".join(parts)))
+                else:
+                    new_headers.append((name, value))
+            return start_response(status, new_headers, exc_info)
+
+        return self.wsgi_app(environ, patched_start_response)

superset/sqllab/utils.py

@@ -31,6 +31,7 @@ DATABASE_KEYS = [
     "allow_cvas",
     "allow_dml",
     "allow_run_async",
+    "allows_cost_estimate",
     "allows_subquery",
     "backend",
     "database_name",

tests/unit_tests/extensions/test_cache_middleware.py

@@ -0,0 +1,156 @@
+# 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.
+from typing import Any, Callable
+
+from superset.extensions.cache_middleware import ExtensionCacheMiddleware
+
+ResponseHeaders = list[tuple[str, str]]
+
+
+def make_wsgi_app(
+    status: str = "200 OK",
+    headers: ResponseHeaders | None = None,
+) -> Callable[..., Any]:
+    """Returns a minimal WSGI app that calls start_response with the given headers."""
+
+    def app(environ, start_response):  # noqa: ARG001
+        start_response(status, headers or [])
+        return [b"body"]
+
+    return app
+
+
+def call_middleware(
+    path: str,
+    upstream_headers: ResponseHeaders,
+) -> ResponseHeaders:
+    """Run middleware for a given path, return headers passed to start_response."""
+    captured: list[ResponseHeaders] = []
+
+    def start_response(status, headers, exc_info=None):  # noqa: ARG001
+        captured.append(headers)
+
+    wsgi_app = make_wsgi_app(headers=upstream_headers)
+    middleware = ExtensionCacheMiddleware(wsgi_app)
+    environ = {"PATH_INFO": path}
+    list(middleware(environ, start_response))
+
+    return captured[0]
+
+
+# --- Path matching ---
+
+
+def test_asset_path_is_intercepted() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/main.js",
+        [("Vary", "Accept-Encoding, Cookie")],
+    )
+    vary = dict(headers).get("Vary", "")
+    assert "Cookie" not in vary
+
+
+def test_list_endpoint_is_not_intercepted() -> None:
+    upstream = [("Vary", "Accept-Encoding, Cookie")]
+    headers = call_middleware("/api/v1/extensions/", upstream)
+    assert headers == upstream
+
+
+def test_get_endpoint_is_not_intercepted() -> None:
+    upstream = [("Vary", "Accept-Encoding, Cookie")]
+    headers = call_middleware("/api/v1/extensions/acme/my-ext", upstream)
+    assert headers == upstream
+
+
+def test_info_endpoint_is_not_intercepted() -> None:
+    upstream = [("Vary", "Accept-Encoding, Cookie")]
+    headers = call_middleware("/api/v1/extensions/_info", upstream)
+    assert headers == upstream
+
+
+def test_unrelated_path_is_not_intercepted() -> None:
+    upstream = [("Vary", "Accept-Encoding, Cookie")]
+    headers = call_middleware("/api/v1/dashboard/", upstream)
+    assert headers == upstream
+
+
+# --- Vary stripping logic ---
+
+
+def test_strips_cookie_from_vary() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Vary", "Accept-Encoding, Cookie")],
+    )
+    assert dict(headers)["Vary"] == "Accept-Encoding"
+
+
+def test_strips_cookie_case_insensitive() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Vary", "Accept-Encoding, COOKIE")],
+    )
+    assert dict(headers)["Vary"] == "Accept-Encoding"
+
+
+def test_removes_vary_header_when_cookie_is_only_value() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Vary", "Cookie")],
+    )
+    assert "Vary" not in dict(headers)
+
+
+def test_multiple_vary_headers_all_stripped() -> None:
+    """Some middleware stacks emit multiple separate Vary headers."""
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Vary", "Cookie"), ("Vary", "Accept-Encoding, Cookie")],
+    )
+    vary_values = [v for k, v in headers if k == "Vary"]
+    assert all("Cookie" not in v for v in vary_values)
+    assert vary_values == ["Accept-Encoding"]
+
+
+def test_non_vary_headers_are_preserved() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.wasm",
+        [
+            ("Content-Type", "application/wasm"),
+            ("Cache-Control", "public, max-age=31536000, immutable"),
+            ("Vary", "Accept-Encoding, Cookie"),
+        ],
+    )
+    d = dict(headers)
+    assert d["Content-Type"] == "application/wasm"
+    assert d["Cache-Control"] == "public, max-age=31536000, immutable"
+
+
+def test_vary_without_cookie_is_unchanged() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Vary", "Accept-Encoding")],
+    )
+    assert dict(headers)["Vary"] == "Accept-Encoding"
+
+
+def test_no_vary_header_produces_no_vary() -> None:
+    headers = call_middleware(
+        "/api/v1/extensions/acme/my-ext/chunk.js",
+        [("Content-Type", "application/javascript")],
+    )
+    assert "Vary" not in dict(headers)