Use explicit context API

superset/extensions/context.py

@@ -16,75 +16,118 @@
 # under the License.
 
 """
-Extension Context Management - provides ambient context during extension loading.
+Extension Context Management - provides ambient context for extensions.
 
-This module provides a thread-local context system that allows decorators to
-automatically detect whether they are being called in host or extension code
-during extension loading.
+This module provides a context system using Python's contextvars that allows
+extensions to access their context (metadata and scoped resources) via get_context().
+
+The context is set during extension loading and when extension callbacks are invoked.
+Uses ContextVar for thread-safe and async-safe context management with automatic
+save/restore for nested contexts.
 """
 
 from __future__ import annotations
 
-import contextlib
-from threading import local
-from typing import Any, Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import Any, Iterator
 
 from superset_core.extensions.types import Manifest
 
-# Thread-local storage for extension context
-_extension_context: local = local()
+
+class ExtensionStorage:
+    """Extension storage with all available tiers."""
+
+    @property
+    def ephemeral(self) -> Any:
+        from superset.extensions.storage.ephemeral_state import EphemeralStateImpl
+
+        return EphemeralStateImpl
 
 
-class ExtensionContext:
-    """Manages ambient extension context during loading."""
-
-    def __init__(self, manifest: Manifest):
-        self.manifest = manifest
-
-    def __enter__(self) -> "ExtensionContext":
-        if getattr(_extension_context, "current", None) is not None:
-            current_extension = _extension_context.current.manifest.id
-            raise RuntimeError(
-                f"Cannot initialize extension {self.manifest.id} while extension "
-                f"{current_extension} is already being initialized. "
-                f"Nested extension initialization is not supported."
-            )
-
-        _extension_context.current = self
-        return self
-
-    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
-        # Clear the current context
-        _extension_context.current = None
-
-
-class ExtensionContextWrapper:
-    """Wrapper for extension context with extensible properties."""
+class ConcreteExtensionContext:
+    """Concrete implementation of ExtensionContext for the host."""
 
     def __init__(self, manifest: Manifest):
         self._manifest = manifest
+        self._storage = ExtensionStorage()
+
+    @property
+    def extension(self) -> Manifest:
+        """Extension metadata (new API)."""
+        return self._manifest
 
     @property
     def manifest(self) -> Manifest:
-        """Get the extension manifest."""
+        """Extension manifest (for backward compatibility)."""
         return self._manifest
 
-    # Future: Add other context properties here
-    # @property
-    # def security_context(self) -> SecurityContext: ...
-    # @property
-    # def build_info(self) -> BuildInfo: ...
+    @property
+    def storage(self) -> ExtensionStorage:
+        return self._storage
 
 
-def get_current_extension_context() -> ExtensionContextWrapper | None:
-    """Get the currently active extension context wrapper, or None if in host code."""
-    if context := getattr(_extension_context, "current", None):
-        return ExtensionContextWrapper(context.manifest)
-    return None
+# Context variable for ambient extension context pattern.
+# Thread-safe and async-safe via Python's contextvars.
+_current_context: ContextVar[ConcreteExtensionContext | None] = ContextVar(
+    "extension_context", default=None
+)
 
 
-@contextlib.contextmanager
-def extension_context(manifest: Manifest) -> Generator[None, None, None]:
-    """Context manager for setting extension context during loading."""
-    with ExtensionContext(manifest):
+def get_context() -> ConcreteExtensionContext:
+    """
+    Get the current extension's context.
+
+    This is the host implementation that replaces the stub in superset_core.
+
+    :returns: The current extension's context.
+    :raises RuntimeError: If called outside of an extension context.
+    """
+    context = _current_context.get()
+    if context is None:
+        raise RuntimeError(
+            "get_context() must be called within an extension context. "
+            "Ensure this code is being executed during extension loading or "
+            "within an extension callback."
+        )
+    return context
+
+
+def get_current_extension_context() -> ConcreteExtensionContext | None:
+    """Get the currently active extension context, or None if in host code."""
+    return _current_context.get()
+
+
+@contextmanager
+def use_context(ctx: ConcreteExtensionContext) -> Iterator[None]:
+    """
+    Context manager to set ambient context for extension execution.
+
+    Used to establish the ambient context before executing extension code.
+    The context is automatically restored after execution, supporting nested
+    context switches.
+
+    :param ctx: ExtensionContext to set as the current context
+    :yields: None
+    """
+    token = _current_context.set(ctx)
+    try:
         yield
+    finally:
+        _current_context.reset(token)
+
+
+@contextmanager
+def extension_context(manifest: Manifest) -> Iterator[ConcreteExtensionContext]:
+    """
+    Context manager for setting extension context during loading.
+
+    Creates a new ExtensionContext for the given manifest and sets it as
+    the current context. Supports nested contexts via ContextVar tokens.
+
+    :param manifest: The extension manifest
+    :yields: The created ExtensionContext
+    """
+    ctx = ConcreteExtensionContext(manifest)
+    with use_context(ctx):
+        yield ctx