diff --git a/docs/developer_docs/extensions/storage.md b/docs/developer_docs/extensions/storage.md index 7a0e84235db..596d378435d 100644 --- a/docs/developer_docs/extensions/storage.md +++ b/docs/developer_docs/extensions/storage.md @@ -34,6 +34,7 @@ Each extension receives its own isolated storage namespace. When Superset loads | ---- | ----------------- | ------------------------------------------ | -------------------------------------- | | 1 | Browser storage | `ctx.storage.local`, `ctx.storage.session` | UI state, wizard progress, draft forms | | 2 | Server-side cache | `ctx.storage.ephemeral` | Job progress, temporary results | +| 3 | Database | `ctx.storage.persistent` | User preferences, durable config | ## Tier 1: Local State @@ -193,6 +194,79 @@ result = ctx.storage.ephemeral.shared.get('shared_result') - Subject to cache eviction under memory pressure - TTL-based expiration (data disappears after timeout) +## Tier 3: Persistent State + +Database-backed storage that survives server restarts, cache evictions, and browser clears. Use for any data that must not be lost. + +### Frontend Usage + +```typescript +import { getContext } from '@apache-superset/core/extensions'; + +const ctx = getContext(); + +// Store user preferences +await ctx.storage.persistent.set('preferences', { theme: 'dark', locale: 'en' }); + +// Retrieve +const prefs = await ctx.storage.persistent.get('preferences'); + +// Remove +await ctx.storage.persistent.remove('preferences'); +``` + +### Backend Usage + +```python +from superset_core.extensions.context import get_context + +ctx = get_context() + +# Store user preferences +ctx.storage.persistent.set('preferences', {'theme': 'dark', 'locale': 'en'}) + +# Retrieve +prefs = ctx.storage.persistent.get('preferences') + +# Remove +ctx.storage.persistent.remove('preferences') +``` + +### Shared State + +For data that should be visible to all users of the extension: + +```typescript +import { getContext } from '@apache-superset/core/extensions'; + +const ctx = getContext(); + +await ctx.storage.persistent.shared.set('global_config', { version: 2 }); +const config = await ctx.storage.persistent.shared.get('global_config'); +``` + +```python +from superset_core.extensions.context import get_context + +ctx = get_context() + +ctx.storage.persistent.shared.set('global_config', {'version': 2}) +config = ctx.storage.persistent.shared.get('global_config') +``` + +### When to Use Tier 3 + +- User preferences and settings +- Extension configuration that must survive restarts +- Saved state that needs to roam across devices and browsers +- Any data where loss is unacceptable + +### Limitations + +- Higher latency than Tiers 1–2 (database round-trip per operation) +- Subject to the 16 MB value size limit +- Requires a database migration when first deployed + ## Key Patterns All storage keys are automatically namespaced: @@ -210,7 +284,9 @@ This ensures: ## Configuration -Administrators can configure Tier 2 storage in `superset_config.py`: +### Tier 2: Ephemeral Storage + +Administrators can configure the server-side cache backend in `superset_config.py`: ```python EXTENSIONS_STORAGE = { @@ -224,3 +300,17 @@ EXTENSIONS_STORAGE = { ``` For development, the default `SupersetMetastoreCache` stores data in the metadata database. + +### Tier 3: Persistent Storage + +Tier 3 values are stored in the `extension_storage` database table. The encryption infrastructure is in place (Fernet-based, keyed from `EXTENSION_STORAGE_ENCRYPTION_KEYS`), but values written through the standard storage API are stored unencrypted by default. Encryption is available at the DAO layer for backend extensions that call `ExtensionStorageDAO.set(..., is_encrypted=True)` directly. + +```python +# Optional: override the encryption key(s) used for Tier 3 persistent storage. +# Falls back to SECRET_KEY when not set. +# Rotate keys by prepending the new key — all keys are tried on decryption. +EXTENSION_STORAGE_ENCRYPTION_KEYS = [ + "my-new-key-base64url-encoded", # used for new writes + "my-old-key-base64url-encoded", # kept for reading old values +] +``` diff --git a/superset-core/src/superset_core/extensions/storage/__init__.py b/superset-core/src/superset_core/extensions/storage/__init__.py index c45b9b2ee09..2e8f23eca72 100644 --- a/superset-core/src/superset_core/extensions/storage/__init__.py +++ b/superset-core/src/superset_core/extensions/storage/__init__.py @@ -50,10 +50,13 @@ Usage: ephemeral_state.shared.get('job_progress') ephemeral_state.shared.set('job_progress', {'pct': 42}, ttl=3600) - # Future: Persistent state - # from superset_core.extensions.storage import persistent_state - # persistent_state.get('config') - # persistent_state.for_resource('dashboard', uuid).get('settings') + # Tier 3: Persistent state + from superset_core.extensions.storage import persistent_state + persistent_state.get('config') + persistent_state.set('config', {'version': 2}) """ -from superset_core.extensions.storage import ephemeral_state # noqa: F401 +from superset_core.extensions.storage import ( + ephemeral_state, # noqa: F401 + persistent_state, # noqa: F401 +) diff --git a/superset-core/src/superset_core/extensions/storage/persistent_state.py b/superset-core/src/superset_core/extensions/storage/persistent_state.py new file mode 100644 index 00000000000..7a8f89ea352 --- /dev/null +++ b/superset-core/src/superset_core/extensions/storage/persistent_state.py @@ -0,0 +1,129 @@ +# 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. + +""" +Persistent State API for superset-core extensions (Tier 3 Storage). + +Provides durable KV storage backed by a dedicated database table. +Data survives server restarts, cache evictions, and browser clears. +Suitable for user preferences, saved state, and any data that must not be lost. + +Host implementations will replace these functions during initialization +with concrete implementations providing actual functionality. + +Database keys are namespaced automatically: +- User-scoped (default): (extension_id, user_id, key) +- Shared (global): (extension_id, null, key) + +Usage: + from superset_core.extensions.storage import persistent_state + + # User-scoped state (default - private to current user) + persistent_state.get('preferences') + persistent_state.set('preferences', {'theme': 'dark'}) + persistent_state.remove('preferences') + + # Shared state (explicit opt-in - visible to all users) + persistent_state.shared.get('global_config') + persistent_state.shared.set('global_config', {'version': 2}) + persistent_state.shared.remove('global_config') +""" + +from typing import Any, Protocol + + +class PersistentStateAccessor(Protocol): + """Protocol for scoped persistent state access.""" + + def get(self, key: str) -> Any: + """Get a value from persistent state.""" + ... + + def set(self, key: str, value: Any) -> None: + """Set a value in persistent state.""" + ... + + def remove(self, key: str) -> None: + """Remove a value from persistent state.""" + ... + + +def get(key: str) -> Any: + """ + Get a value from user-scoped persistent state. + + Data is automatically scoped to the current authenticated user. + Other users cannot see or modify this data. + + Host implementations will replace this function during initialization + with a concrete implementation providing actual functionality. + + :param key: The key to retrieve. + :returns: The stored value, or None if not found. + """ + raise NotImplementedError("Function will be replaced during initialization") + + +def set(key: str, value: Any) -> None: + """ + Set a value in user-scoped persistent state. + + Data is automatically scoped to the current authenticated user. + Other users cannot see or modify this data. + Data persists indefinitely until explicitly removed. + + Host implementations will replace this function during initialization + with a concrete implementation providing actual functionality. + + :param key: The key to store. + :param value: The value to store (must be JSON-serializable). + """ + raise NotImplementedError("Function will be replaced during initialization") + + +def remove(key: str) -> None: + """ + Remove a value from user-scoped persistent state. + + Data is automatically scoped to the current authenticated user. + + Host implementations will replace this function during initialization + with a concrete implementation providing actual functionality. + + :param key: The key to remove. + """ + raise NotImplementedError("Function will be replaced during initialization") + + +class _SharedStub: + """Stub for shared accessor that raises NotImplementedError on any operation.""" + + def get(self, key: str) -> Any: + raise NotImplementedError("Accessor will be replaced during initialization") + + def set(self, key: str, value: Any) -> None: + raise NotImplementedError("Accessor will be replaced during initialization") + + def remove(self, key: str) -> None: + raise NotImplementedError("Accessor will be replaced during initialization") + + +#: Shared (global) persistent state accessor. +#: Data stored via this accessor is visible to all users of the extension. +#: WARNING: Do not store user-specific or sensitive data here. +#: Host implementations will replace this during initialization. +shared: PersistentStateAccessor = _SharedStub() diff --git a/superset-frontend/packages/superset-core/src/extensions/index.ts b/superset-frontend/packages/superset-core/src/extensions/index.ts index 0bcd477f400..4f73fa5c37a 100644 --- a/superset-frontend/packages/superset-core/src/extensions/index.ts +++ b/superset-frontend/packages/superset-core/src/extensions/index.ts @@ -73,6 +73,13 @@ export interface ExtensionStorage { * Use `.shared` for data visible to all users. */ ephemeral: StorageTier; + + /** + * Durable database-backed storage (Tier 3). + * Data survives server restarts and cache evictions. + * Use `.shared` for data visible to all users. + */ + persistent: StorageTier; } /** diff --git a/superset-frontend/packages/superset-core/src/storage/index.ts b/superset-frontend/packages/superset-core/src/storage/index.ts index 8f8a0ee671c..611a74f9a5d 100644 --- a/superset-frontend/packages/superset-core/src/storage/index.ts +++ b/superset-frontend/packages/superset-core/src/storage/index.ts @@ -56,4 +56,5 @@ export * as localState from './localState'; export * as sessionState from './sessionState'; export * as ephemeralState from './ephemeralState'; +export * as persistentState from './persistentState'; export * from './types'; diff --git a/superset-frontend/packages/superset-core/src/storage/persistentState.ts b/superset-frontend/packages/superset-core/src/storage/persistentState.ts new file mode 100644 index 00000000000..2a7dbd86204 --- /dev/null +++ b/superset-frontend/packages/superset-core/src/storage/persistentState.ts @@ -0,0 +1,123 @@ +/** + * 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. + */ + +import type { JsonValue, StorageAccessor } from './types'; + +/** + * @fileoverview Persistent State API for Superset extensions (Tier 3 Storage). + * + * Provides durable KV storage backed by a dedicated database table. + * Data survives server restarts, cache evictions, and browser clears. + * Suitable for user preferences, saved state, and any data that must + * not be lost. + * + * By default, all operations are user-scoped (private to the current user). + * Use `shared` to access state that is visible to all users of the extension. + * + * Database keys are namespaced automatically: + * - User-scoped (default): (extension_id, user_id, key) + * - Shared (global): (extension_id, null, key) + * + * @example + * ```typescript + * import { persistentState } from '@apache-superset/core/storage'; + * + * // User-scoped state (default - private to current user) + * const prefs = await persistentState.get('preferences'); + * await persistentState.set('preferences', { theme: 'dark', locale: 'en' }); + * await persistentState.remove('preferences'); + * + * // Shared state (explicit opt-in - visible to all users) + * const config = await persistentState.shared.get('global_config'); + * await persistentState.shared.set('global_config', { version: 2 }); + * await persistentState.shared.remove('global_config'); + * ``` + */ + +/** + * Get a value from user-scoped persistent state. + * + * Data is automatically scoped to the current authenticated user. + * Other users cannot see or modify this data. + * + * @param key The key to retrieve. + * @returns The stored value, or null if not found. + * + * @example + * ```typescript + * const prefs = await persistentState.get('preferences'); + * if (prefs !== null) { + * applyPreferences(prefs); + * } + * ``` + */ +export declare function get(key: string): Promise; + +/** + * Set a value in user-scoped persistent state. + * + * Data is automatically scoped to the current authenticated user. + * Other users cannot see or modify this data. + * Data persists indefinitely until explicitly removed. + * + * @param key The key to store. + * @param value The value to store (must be JSON-serializable). + * + * @example + * ```typescript + * await persistentState.set('preferences', { theme: 'dark', locale: 'en' }); + * ``` + */ +export declare function set(key: string, value: JsonValue): Promise; + +/** + * Remove a value from user-scoped persistent state. + * + * @param key The key to remove. + * + * @example + * ```typescript + * await persistentState.remove('preferences'); + * ``` + */ +export declare function remove(key: string): Promise; + +/** + * Shared (global) persistent state accessor. + * + * Accessor for state that is shared across all users of the extension. + * Use this for extension-wide configuration, shared datasets, or any + * data that should be accessible to all users regardless of identity. + * + * WARNING: Data stored via shared is visible to all users of the extension. + * Do not store user-specific or sensitive data here. + * + * @example + * ```typescript + * // Read shared extension config + * const config = await persistentState.shared.get('global_config'); + * + * // Update shared config (typically admin-only) + * await persistentState.shared.set('global_config', { version: 2 }); + * + * // Remove shared config entry + * await persistentState.shared.remove('global_config'); + * ``` + */ +export declare const shared: StorageAccessor; diff --git a/superset-frontend/src/core/storage/index.ts b/superset-frontend/src/core/storage/index.ts index 8fc53840bed..b26f9a411f8 100644 --- a/superset-frontend/src/core/storage/index.ts +++ b/superset-frontend/src/core/storage/index.ts @@ -24,3 +24,4 @@ export { createBrowserStorage } from './localState'; export { createEphemeralState } from './ephemeralState'; +export { createPersistentState } from './persistentState'; diff --git a/superset-frontend/src/core/storage/persistentState.ts b/superset-frontend/src/core/storage/persistentState.ts new file mode 100644 index 00000000000..99354ff351b --- /dev/null +++ b/superset-frontend/src/core/storage/persistentState.ts @@ -0,0 +1,80 @@ +/** + * 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. + */ + +import { type storage as StorageTypes } from '@apache-superset/core'; +import { SupersetClient } from '@superset-ui/core'; + +/** + * Create persistent state (database-backed) bound to an extension ID. + */ +export function createPersistentState( + extensionId: string, +): typeof StorageTypes.persistentState { + const MAX_KEY_LENGTH = 255; + + const buildUrl = (key: string, shared?: boolean): string => { + if (key.length > MAX_KEY_LENGTH) { + throw new Error( + `Persistent storage key must be ${MAX_KEY_LENGTH} characters or less.`, + ); + } + const basePath = '/api/v1/extensions/storage/persistent'; + const encodedId = encodeURIComponent(extensionId); + const encodedKey = encodeURIComponent(key); + const url = `${basePath}/${encodedId}/${encodedKey}`; + return shared ? `${url}?shared=true` : url; + }; + + const shared: StorageTypes.StorageAccessor = { + get: async (key: string) => { + const response = await SupersetClient.get({ + endpoint: buildUrl(key, true), + }); + return response.json?.result ?? null; + }, + set: async (key: string, value: StorageTypes.JsonValue) => { + await SupersetClient.put({ + endpoint: buildUrl(key, true), + body: JSON.stringify({ value }), + headers: { 'Content-Type': 'application/json' }, + }); + }, + remove: async (key: string) => { + await SupersetClient.delete({ endpoint: buildUrl(key, true) }); + }, + }; + + return { + get: async (key: string) => { + const response = await SupersetClient.get({ endpoint: buildUrl(key) }); + return response.json?.result ?? null; + }, + set: async (key: string, value: StorageTypes.JsonValue) => { + await SupersetClient.put({ + endpoint: buildUrl(key), + body: JSON.stringify({ value }), + headers: { 'Content-Type': 'application/json' }, + }); + }, + remove: async (key: string) => { + await SupersetClient.delete({ endpoint: buildUrl(key) }); + }, + shared, + }; +} diff --git a/superset-frontend/src/extensions/ExtensionContext.ts b/superset-frontend/src/extensions/ExtensionContext.ts index 1feea23b373..e3d7a4726f4 100644 --- a/superset-frontend/src/extensions/ExtensionContext.ts +++ b/superset-frontend/src/extensions/ExtensionContext.ts @@ -20,7 +20,11 @@ import type { extensions as extensionsApi, common, } from '@apache-superset/core'; -import { createBrowserStorage, createEphemeralState } from 'src/core/storage'; +import { + createBrowserStorage, + createEphemeralState, + createPersistentState, +} from 'src/core/storage'; type ExtensionContextType = extensionsApi.ExtensionContext; type Extension = common.Extension; @@ -44,6 +48,7 @@ class ExtensionContext implements ExtensionContextType { local: createBrowserStorage(localStorage, id), session: createBrowserStorage(sessionStorage, id), ephemeral: createEphemeralState(id), + persistent: createPersistentState(id), }; } return this._storage; diff --git a/superset/core/api/core_api_injection.py b/superset/core/api/core_api_injection.py index fdbce562ea8..93def181b06 100644 --- a/superset/core/api/core_api_injection.py +++ b/superset/core/api/core_api_injection.py @@ -232,8 +232,10 @@ def inject_storage_implementations() -> None: implementations from Superset. """ import superset_core.extensions.storage.ephemeral_state as core_ephemeral_state + import superset_core.extensions.storage.persistent_state as core_persistent_state from superset.extensions.storage.ephemeral_state import EphemeralStateImpl + from superset.extensions.storage.persistent_state_impl import PersistentStateImpl # Replace abstract functions with concrete implementations core_ephemeral_state.get = EphemeralStateImpl.get @@ -241,6 +243,11 @@ def inject_storage_implementations() -> None: core_ephemeral_state.remove = EphemeralStateImpl.remove core_ephemeral_state.shared = EphemeralStateImpl.shared + core_persistent_state.get = PersistentStateImpl.get + core_persistent_state.set = PersistentStateImpl.set + core_persistent_state.remove = PersistentStateImpl.remove + core_persistent_state.shared = PersistentStateImpl.shared + def inject_extension_context() -> None: """ diff --git a/superset/extensions/context.py b/superset/extensions/context.py index b299d949bac..8748ddaf808 100644 --- a/superset/extensions/context.py +++ b/superset/extensions/context.py @@ -44,6 +44,14 @@ class ExtensionStorage: return EphemeralStateImpl + @property + def persistent(self) -> Any: + from superset.extensions.storage.persistent_state_impl import ( + PersistentStateImpl, + ) + + return PersistentStateImpl + class ConcreteExtensionContext: """Concrete implementation of ExtensionContext for the host.""" diff --git a/superset/extensions/storage/api.py b/superset/extensions/storage/api.py index ea67105c945..15fd9cdb5ce 100644 --- a/superset/extensions/storage/api.py +++ b/superset/extensions/storage/api.py @@ -34,8 +34,11 @@ from flask.wrappers import Response from flask_appbuilder.api import BaseApi, expose, protect, safe from superset.extensions import cache_manager +from superset.extensions.storage.persistent_state_dao import ExtensionStorageDAO from superset.extensions.types import LoadedExtension from superset.extensions.utils import get_extensions +from superset.utils import json +from superset.utils.decorators import transaction # Key separator SEPARATOR = ":" @@ -264,3 +267,158 @@ class ExtensionStorageRestApi(BaseApi): cache_manager.extension_ephemeral_state_cache.delete(cache_key) return self.response(200, message="Value deleted successfully") + + @protect() + @safe + @expose("/persistent//", methods=("GET",)) + def get_persistent(self, extension_id: str, key: str, **kwargs: Any) -> Response: + """Get a value from persistent state. + --- + get: + summary: Get a value from persistent state + parameters: + - in: path + name: extension_id + schema: + type: string + required: true + description: Extension ID (publisher.name) + - in: path + name: key + schema: + type: string + required: true + description: Storage key + - in: query + name: shared + schema: + type: boolean + required: false + description: If true, read from shared state visible to all users + responses: + 200: + description: Value retrieved successfully + content: + application/json: + schema: + type: object + properties: + result: + description: The stored value + 404: + description: Extension not found + """ + extension = _get_extension_or_404(extension_id) + if not extension: + return self.response_404("Extension not found") + + shared = request.args.get("shared", "false").lower() == "true" + user_fk = None if shared else g.user.id + raw = ExtensionStorageDAO.get_value(extension_id, key, user_fk=user_fk) + value = json.loads(raw) if raw is not None else None + + return self.response(200, result=value) + + @protect() + @safe + @expose("/persistent//", methods=("PUT",)) + @transaction() + def set_persistent(self, extension_id: str, key: str, **kwargs: Any) -> Response: + """Set a value in persistent state. + --- + put: + summary: Set a value in persistent state + parameters: + - in: path + name: extension_id + schema: + type: string + required: true + description: Extension ID (publisher.name) + - in: path + name: key + schema: + type: string + required: true + description: Storage key + - in: query + name: shared + schema: + type: boolean + required: false + description: If true, store as shared state visible to all users + requestBody: + required: true + content: + application/json: + schema: + type: object + properties: + value: + description: The value to store (must be JSON-serializable) + responses: + 200: + description: Value stored successfully + 400: + description: Invalid request body + 404: + description: Extension not found + """ + extension = _get_extension_or_404(extension_id) + if not extension: + return self.response_404("Extension not found") + + body = request.get_json(silent=True) or {} + if "value" not in body: + return self.response_400("Request body must contain 'value' field") + + shared = request.args.get("shared", "false").lower() == "true" + user_fk = None if shared else g.user.id + value_bytes = json.dumps(body["value"]).encode() + ExtensionStorageDAO.set(extension_id, key, value_bytes, user_fk=user_fk) + + return self.response(200, message="Value stored successfully") + + @protect() + @safe + @expose("/persistent//", methods=("DELETE",)) + @transaction() + def delete_persistent(self, extension_id: str, key: str, **kwargs: Any) -> Response: + """Delete a value from persistent state. + --- + delete: + summary: Delete a value from persistent state + parameters: + - in: path + name: extension_id + schema: + type: string + required: true + description: Extension ID (publisher.name) + - in: path + name: key + schema: + type: string + required: true + description: Storage key + - in: query + name: shared + schema: + type: boolean + required: false + description: If true, delete from shared state + responses: + 200: + description: Value deleted successfully + 404: + description: Extension not found + """ + extension = _get_extension_or_404(extension_id) + if not extension: + return self.response_404("Extension not found") + + shared = request.args.get("shared", "false").lower() == "true" + user_fk = None if shared else g.user.id + ExtensionStorageDAO.delete(extension_id, key, user_fk=user_fk) + + return self.response(200, message="Value deleted successfully") diff --git a/superset/extensions/storage/persistent_state_dao.py b/superset/extensions/storage/persistent_state_dao.py new file mode 100644 index 00000000000..1c01b807973 --- /dev/null +++ b/superset/extensions/storage/persistent_state_dao.py @@ -0,0 +1,277 @@ +# 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 base64 +import hashlib +import logging + +from cryptography.fernet import Fernet, InvalidToken, MultiFernet +from flask import current_app +from sqlalchemy import and_ + +from superset import db +from superset.extensions.storage.persistent_state_model import ExtensionStorage + +logger = logging.getLogger(__name__) + + +def _key_to_fernet(raw_key: str | bytes) -> Fernet: + """Derive a Fernet instance from an arbitrary-length secret string. + + SHA-256 compresses the key to exactly 32 bytes, which are then + base64url-encoded to satisfy Fernet's key format requirement. + """ + if isinstance(raw_key, str): + raw_key = raw_key.encode() + return Fernet(base64.urlsafe_b64encode(hashlib.sha256(raw_key).digest())) + + +def _fernet() -> MultiFernet: + """Return a MultiFernet built from EXTENSION_STORAGE_ENCRYPTION_KEYS. + + Falls back to SECRET_KEY when the config list is absent. The first key in + the list is used for new encryptions; all keys are tried on decryption, + enabling zero-downtime rotation: add the new key at the front of + EXTENSION_STORAGE_ENCRYPTION_KEYS, then run ``superset rotate-extension- + storage-keys`` to re-encrypt every row with the new key. + """ + raw_keys: list[str | bytes] = current_app.config.get( + "EXTENSION_STORAGE_ENCRYPTION_KEYS" + ) or [current_app.config["SECRET_KEY"]] + return MultiFernet([_key_to_fernet(k) for k in raw_keys]) + + +def _scope_filter( + extension_id: str, + key: str, + user_fk: int | None = None, + resource_type: str | None = None, + resource_uuid: str | None = None, +) -> list[object]: + """Build the SQLAlchemy filter list for a scoped lookup.""" + filters: list[object] = [ + ExtensionStorage.extension_id == extension_id, + ExtensionStorage.key == key, + ] + if user_fk is None: + filters.append(ExtensionStorage.user_fk.is_(None)) + else: + filters.append(ExtensionStorage.user_fk == user_fk) + if resource_type is None: + filters.append(ExtensionStorage.resource_type.is_(None)) + else: + filters.append(ExtensionStorage.resource_type == resource_type) + if resource_uuid is None: + filters.append(ExtensionStorage.resource_uuid.is_(None)) + else: + filters.append(ExtensionStorage.resource_uuid == resource_uuid) + return filters + + +class ExtensionStorageDAO: + """Persistent key-value store for extensions. + + Provides scoped get/set/delete and list operations covering the three + storage scopes defined by the Tier 3 proposal: + + * Global scope — user_fk=None, resource_type=None + * User scope — user_fk=, resource_type=None + * Resource scope — resource_type + resource_uuid set + """ + + # ── Read ───────────────────────────────────────────────────────────────── + + @staticmethod + def get( + extension_id: str, + key: str, + user_fk: int | None = None, + resource_type: str | None = None, + resource_uuid: str | None = None, + ) -> ExtensionStorage | None: + """Return the raw storage entry. The value field may be encrypted.""" + entry = ( + db.session.query(ExtensionStorage) + .filter( + and_( + *_scope_filter( + extension_id, key, user_fk, resource_type, resource_uuid + ) + ) + ) + .first() + ) + return entry + + @staticmethod + def get_value( + extension_id: str, + key: str, + user_fk: int | None = None, + resource_type: str | None = None, + resource_uuid: str | None = None, + ) -> bytes | None: + """Return the raw (decrypted) value bytes, or None if not found.""" + entry = ExtensionStorageDAO.get( + extension_id, key, user_fk, resource_type, resource_uuid + ) + if entry is None: + return None + if entry.is_encrypted: + try: + return _fernet().decrypt(entry.value) + except InvalidToken: + logger.error( + "Failed to decrypt extension storage value for " + "extension_id=%s key=%s — possible key rotation issue", + extension_id, + key, + ) + return None + return entry.value + + # ── Write (upsert) ──────────────────────────────────────────────────────── + + @staticmethod + def set( + extension_id: str, + key: str, + value: bytes, + value_type: str = "application/json", + user_fk: int | None = None, + resource_type: str | None = None, + resource_uuid: str | None = None, + category: str | None = None, + description: str | None = None, + is_encrypted: bool = False, + ) -> ExtensionStorage: + """Upsert a storage entry. Encrypts value when is_encrypted=True.""" + stored_value = _fernet().encrypt(value) if is_encrypted else value + + entry = ( + db.session.query(ExtensionStorage) + .filter( + and_( + *_scope_filter( + extension_id, key, user_fk, resource_type, resource_uuid + ) + ) + ) + .first() + ) + if entry is not None: + entry.value = stored_value + entry.value_type = value_type + entry.category = category + entry.description = description + entry.is_encrypted = is_encrypted + else: + entry = ExtensionStorage( + extension_id=extension_id, + key=key, + value=stored_value, + value_type=value_type, + user_fk=user_fk, + resource_type=resource_type, + resource_uuid=resource_uuid, + category=category, + description=description, + is_encrypted=is_encrypted, + ) + db.session.add(entry) + db.session.flush() + return entry + + # ── Delete ──────────────────────────────────────────────────────────────── + + @staticmethod + def delete( + extension_id: str, + key: str, + user_fk: int | None = None, + resource_type: str | None = None, + resource_uuid: str | None = None, + ) -> bool: + """Delete an entry. Returns True if a row was removed.""" + entry = ( + db.session.query(ExtensionStorage) + .filter( + and_( + *_scope_filter( + extension_id, key, user_fk, resource_type, resource_uuid + ) + ) + ) + .first() + ) + if entry is None: + return False + db.session.delete(entry) + db.session.flush() + return True + + # ── List ────────────────────────────────────────────────────────────────── + + @staticmethod + def list_global( + extension_id: str, + category: str | None = None, + ) -> list[ExtensionStorage]: + """List all global (user_fk=NULL, resource_type=NULL) entries.""" + q = db.session.query(ExtensionStorage).filter( + ExtensionStorage.extension_id == extension_id, + ExtensionStorage.user_fk.is_(None), + ExtensionStorage.resource_type.is_(None), + ) + if category is not None: + q = q.filter(ExtensionStorage.category == category) + return q.order_by(ExtensionStorage.key).all() + + @staticmethod + def list_user( + extension_id: str, + user_fk: int, + category: str | None = None, + ) -> list[ExtensionStorage]: + """List all user-scoped entries (resource_type=NULL).""" + q = db.session.query(ExtensionStorage).filter( + ExtensionStorage.extension_id == extension_id, + ExtensionStorage.user_fk == user_fk, + ExtensionStorage.resource_type.is_(None), + ) + if category is not None: + q = q.filter(ExtensionStorage.category == category) + return q.order_by(ExtensionStorage.key).all() + + @staticmethod + def list_resource( + extension_id: str, + resource_type: str, + resource_uuid: str, + category: str | None = None, + ) -> list[ExtensionStorage]: + """List all entries linked to a specific resource.""" + q = db.session.query(ExtensionStorage).filter( + ExtensionStorage.extension_id == extension_id, + ExtensionStorage.resource_type == resource_type, + ExtensionStorage.resource_uuid == resource_uuid, + ) + if category is not None: + q = q.filter(ExtensionStorage.category == category) + return q.order_by(ExtensionStorage.key).all() diff --git a/superset/extensions/storage/persistent_state_impl.py b/superset/extensions/storage/persistent_state_impl.py new file mode 100644 index 00000000000..205692019d0 --- /dev/null +++ b/superset/extensions/storage/persistent_state_impl.py @@ -0,0 +1,165 @@ +# 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. + +""" +Host implementation for Persistent State (Tier 3 Storage). + +Provides the concrete database-backed implementation that is injected into +superset_core.extensions.storage.persistent_state at startup. +""" + +from __future__ import annotations + +from typing import Any + +from flask import g + +from superset.extensions.context import get_current_extension_context +from superset.extensions.storage.persistent_state_dao import ExtensionStorageDAO +from superset.utils import json +from superset.utils.decorators import transaction + + +def _get_extension_id() -> str: + """Get the current extension ID from context.""" + context = get_current_extension_context() + if context is None: + raise RuntimeError( + "persistent_state can only be used within an extension context. " + "Ensure this code is being executed during extension loading or " + "within an extension API request handler." + ) + return context.manifest.id + + +def _get_current_user_id() -> int: + """Get the current authenticated user's ID.""" + user = getattr(g, "user", None) + if user is None or not hasattr(user, "id"): + raise RuntimeError( + "persistent_state requires an authenticated user. " + "Ensure the request has been authenticated." + ) + return user.id + + +def _decode(raw: bytes | None) -> Any: + """Decode stored bytes back to a Python value.""" + if raw is None: + return None + return json.loads(raw) + + +def _encode(value: Any) -> bytes: + """Encode a Python value for database storage.""" + return json.dumps(value).encode() + + +class SharedPersistentStateAccessor: + """ + Accessor for shared (global) persistent state. + + Data stored via this accessor is visible to all users of the extension. + Extension ID is resolved lazily on each operation from the current context. + """ + + def get(self, key: str) -> Any: + """ + Get a value from shared persistent state. + + :param key: The key to retrieve. + :returns: The stored value, or None if not found. + """ + extension_id = _get_extension_id() + raw = ExtensionStorageDAO.get_value(extension_id, key, user_fk=None) + return _decode(raw) + + @transaction() + def set(self, key: str, value: Any) -> None: + """ + Set a value in shared persistent state. + + :param key: The key to store. + :param value: The value to store (must be JSON-serializable). + """ + extension_id = _get_extension_id() + ExtensionStorageDAO.set(extension_id, key, _encode(value), user_fk=None) + + @transaction() + def remove(self, key: str) -> None: + """ + Remove a value from shared persistent state. + + :param key: The key to remove. + """ + extension_id = _get_extension_id() + ExtensionStorageDAO.delete(extension_id, key, user_fk=None) + + +class PersistentStateImpl: + """ + Host implementation for persistent state operations. + + This class provides the concrete implementation that is injected into + superset_core.extensions.storage.persistent_state. + + By default, all operations are user-scoped (private to the current user). + Use `shared` to access state that is visible to all users. + """ + + @staticmethod + def get(key: str) -> Any: + """ + Get a value from user-scoped persistent state. + + :param key: The key to retrieve. + :returns: The stored value, or None if not found. + """ + extension_id = _get_extension_id() + user_id = _get_current_user_id() + raw = ExtensionStorageDAO.get_value(extension_id, key, user_fk=user_id) + return _decode(raw) + + @staticmethod + @transaction() + def set(key: str, value: Any) -> None: + """ + Set a value in user-scoped persistent state. + + :param key: The key to store. + :param value: The value to store (must be JSON-serializable). + """ + extension_id = _get_extension_id() + user_id = _get_current_user_id() + ExtensionStorageDAO.set(extension_id, key, _encode(value), user_fk=user_id) + + @staticmethod + @transaction() + def remove(key: str) -> None: + """ + Remove a value from user-scoped persistent state. + + :param key: The key to remove. + """ + extension_id = _get_extension_id() + user_id = _get_current_user_id() + ExtensionStorageDAO.delete(extension_id, key, user_fk=user_id) + + #: Shared (global) persistent state accessor. + #: Data stored via this accessor is visible to all users of the extension. + #: WARNING: Do not store user-specific or sensitive data here. + shared: SharedPersistentStateAccessor = SharedPersistentStateAccessor() diff --git a/superset/extensions/storage/persistent_state_model.py b/superset/extensions/storage/persistent_state_model.py new file mode 100644 index 00000000000..9790956b618 --- /dev/null +++ b/superset/extensions/storage/persistent_state_model.py @@ -0,0 +1,126 @@ +# 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. + +import uuid as uuid_module + +from flask_appbuilder import Model +from sqlalchemy import ( + Boolean, + Column, + ForeignKey, + Index, + Integer, + LargeBinary, + String, + Text, + UniqueConstraint, +) +from sqlalchemy.orm import backref, relationship +from sqlalchemy_utils import UUIDType + +from superset.models.helpers import AuditMixinNullable + +# 16 MB — matches the KeyValue store limit. +EXTENSION_STORAGE_MAX_SIZE = 2**24 - 1 + + +class ExtensionStorage(AuditMixinNullable, Model): + """Generic persistent key-value storage for extensions (Tier 3). + + Each row is identified by (extension_id, user_fk, resource_type, + resource_uuid, key): + + * Global scope — user_fk IS NULL, resource_type IS NULL + * User scope — user_fk set, resource_type IS NULL + * Resource scope — resource_type + resource_uuid set (user_fk optional) + + The payload is stored as raw bytes (value) with a MIME-type hint + (value_type). When is_encrypted is True the value has been encrypted + at the DAO layer using Fernet and must be decrypted before use. + """ + + __tablename__ = "extension_storage" + + id = Column(Integer, primary_key=True, autoincrement=True) + uuid = Column( + UUIDType(binary=True), + default=uuid_module.uuid4, + unique=True, + nullable=False, + ) + + # Extension identity + extension_id = Column(String(255), nullable=False) + + # Scope discriminators — all nullable; NULLs define the scope (see docstring) + user_fk = Column( + Integer, + ForeignKey( + "ab_user.id", + ondelete="SET NULL", + name="fk_extension_storage_user_fk_ab_user", + ), + nullable=True, + ) + resource_type = Column(String(64), nullable=True) + resource_uuid = Column(String(36), nullable=True) + + # Storage key within the scope + key = Column(String(255), nullable=False) + + # Optional metadata + category = Column(String(64), nullable=True) + description = Column(Text, nullable=True) + + # Payload + value = Column(LargeBinary(EXTENSION_STORAGE_MAX_SIZE), nullable=False) + value_type = Column(String(255), nullable=False, default="application/json") + is_encrypted = Column(Boolean, nullable=False, default=False) + + user = relationship( + "User", + backref=backref("extension_storage_entries", cascade="all, delete-orphan"), + foreign_keys=[user_fk], + ) + + __table_args__ = ( + # Unique constraint prevents duplicate rows from concurrent writes + UniqueConstraint( + "extension_id", + "user_fk", + "resource_type", + "resource_uuid", + "key", + name="uq_extension_storage_scoped_key", + ), + # Composite index covering all lookup dimensions + Index( + "ix_ext_storage_lookup", + "extension_id", + "user_fk", + "resource_type", + "resource_uuid", + "key", + ), + Index("ix_ext_storage_extension_id", "extension_id"), + ) + + def __repr__(self) -> str: + return ( + f"" + ) diff --git a/superset/migrations/versions/2026-04-07_12-00_e5f6a7b8c9d0_add_extension_storage_table.py b/superset/migrations/versions/2026-04-07_12-00_e5f6a7b8c9d0_add_extension_storage_table.py new file mode 100644 index 00000000000..5f0088d145b --- /dev/null +++ b/superset/migrations/versions/2026-04-07_12-00_e5f6a7b8c9d0_add_extension_storage_table.py @@ -0,0 +1,107 @@ +# 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. +"""add_extension_storage_table + +Revision ID: e5f6a7b8c9d0 +Revises: a1b2c3d4e5f6 +Create Date: 2026-04-07 12:00:00.000000 + +""" + +# revision identifiers, used by Alembic. +revision = "e5f6a7b8c9d0" +down_revision = "a1b2c3d4e5f6" + +import sqlalchemy as sa # noqa: E402 +from alembic import op # noqa: E402 +from sqlalchemy_utils import UUIDType # noqa: E402 + + +def upgrade() -> None: + op.create_table( + "extension_storage", + sa.Column("id", sa.Integer(), nullable=False, autoincrement=True), + sa.Column("uuid", UUIDType(binary=True), nullable=False), + sa.Column("extension_id", sa.String(255), nullable=False), + sa.Column("user_fk", sa.Integer(), nullable=True), + sa.Column("resource_type", sa.String(64), nullable=True), + sa.Column("resource_uuid", sa.String(36), nullable=True), + sa.Column("key", sa.String(255), nullable=False), + sa.Column("category", sa.String(64), nullable=True), + sa.Column("description", sa.Text(), nullable=True), + sa.Column("value", sa.LargeBinary(2**24 - 1), nullable=False), + sa.Column( + "value_type", + sa.String(255), + nullable=False, + server_default="application/json", + ), + sa.Column( + "is_encrypted", + sa.Boolean(), + nullable=False, + server_default=sa.false(), + ), + sa.Column("created_on", sa.DateTime(), nullable=True), + sa.Column("changed_on", sa.DateTime(), nullable=True), + sa.Column("created_by_fk", sa.Integer(), nullable=True), + sa.Column("changed_by_fk", sa.Integer(), nullable=True), + sa.ForeignKeyConstraint( + ["user_fk"], + ["ab_user.id"], + name="fk_extension_storage_user_fk_ab_user", + ondelete="SET NULL", + ), + sa.ForeignKeyConstraint( + ["created_by_fk"], + ["ab_user.id"], + name="fk_extension_storage_created_by_fk_ab_user", + ondelete="SET NULL", + ), + sa.ForeignKeyConstraint( + ["changed_by_fk"], + ["ab_user.id"], + name="fk_extension_storage_changed_by_fk_ab_user", + ondelete="SET NULL", + ), + sa.PrimaryKeyConstraint("id"), + sa.UniqueConstraint("uuid"), + sa.UniqueConstraint( + "extension_id", + "user_fk", + "resource_type", + "resource_uuid", + "key", + name="uq_extension_storage_scoped_key", + ), + ) + op.create_index( + "ix_ext_storage_extension_id", + "extension_storage", + ["extension_id"], + ) + op.create_index( + "ix_ext_storage_lookup", + "extension_storage", + ["extension_id", "user_fk", "resource_type", "resource_uuid", "key"], + ) + + +def downgrade() -> None: + op.drop_index("ix_ext_storage_lookup", "extension_storage") + op.drop_index("ix_ext_storage_extension_id", "extension_storage") + op.drop_table("extension_storage") diff --git a/tests/unit_tests/extensions/storage/test_persistent_state.py b/tests/unit_tests/extensions/storage/test_persistent_state.py new file mode 100644 index 00000000000..f84629afc49 --- /dev/null +++ b/tests/unit_tests/extensions/storage/test_persistent_state.py @@ -0,0 +1,169 @@ +# 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. + +"""Tests for persistent state storage implementation.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from flask import Flask, g +from superset_core.extensions.types import Manifest + +from superset.extensions.context import ConcreteExtensionContext, use_context +from superset.extensions.storage.persistent_state_impl import ( + PersistentStateImpl, + SharedPersistentStateAccessor, +) +from superset.utils import json + + +@pytest.fixture +def app() -> Flask: + """Create a minimal Flask app for testing.""" + app = Flask(__name__) + app.config["TESTING"] = True + return app + + +def _create_context( + publisher: str = "test-org", name: str = "test-ext" +) -> ConcreteExtensionContext: + """Create test context with given extension identifiers.""" + manifest = Manifest.model_validate( + { + "id": f"{publisher}.{name}", + "publisher": publisher, + "name": name, + "displayName": f"Test {name}", + } + ) + return ConcreteExtensionContext(manifest) + + +def _set_user(user_id: int) -> None: + """Set a mock user on Flask's g object.""" + g.user = MagicMock(id=user_id) + + +@patch("superset.db") +def test_persistent_state_raises_without_context( + mock_db: MagicMock, app: Flask +) -> None: + """PersistentStateImpl operations raise RuntimeError without extension context.""" + with app.app_context(): + _set_user(1) + + with pytest.raises(RuntimeError, match="within an extension context"): + PersistentStateImpl.get("key") + + with pytest.raises(RuntimeError, match="within an extension context"): + PersistentStateImpl.set("key", "value") + + with pytest.raises(RuntimeError, match="within an extension context"): + PersistentStateImpl.remove("key") + + +def test_persistent_state_raises_without_user(app: Flask) -> None: + """PersistentStateImpl operations raise RuntimeError without authenticated user.""" + ctx = _create_context() + + with app.app_context(), use_context(ctx): + with pytest.raises(RuntimeError, match="requires an authenticated user"): + PersistentStateImpl.get("key") + + +@patch("superset.extensions.storage.persistent_state_impl.ExtensionStorageDAO") +def test_persistent_state_get_returns_value(mock_dao: MagicMock, app: Flask) -> None: + """PersistentStateImpl.get returns decoded value from DAO.""" + ctx = _create_context() + stored = json.dumps({"theme": "dark"}).encode() + mock_dao.get_value.return_value = stored + + with app.app_context(), use_context(ctx): + _set_user(42) + result = PersistentStateImpl.get("prefs") + + mock_dao.get_value.assert_called_once_with("test-org.test-ext", "prefs", user_fk=42) + assert result == {"theme": "dark"} + + +@patch("superset.extensions.storage.persistent_state_impl.ExtensionStorageDAO") +def test_persistent_state_get_returns_none_when_missing( + mock_dao: MagicMock, app: Flask +) -> None: + """PersistentStateImpl.get returns None when key does not exist.""" + ctx = _create_context() + mock_dao.get_value.return_value = None + + with app.app_context(), use_context(ctx): + _set_user(42) + result = PersistentStateImpl.get("missing") + + assert result is None + + +@patch("superset.db") +@patch("superset.extensions.storage.persistent_state_impl.ExtensionStorageDAO") +def test_persistent_state_set_encodes_value( + mock_dao: MagicMock, mock_db: MagicMock, app: Flask +) -> None: + """PersistentStateImpl.set encodes value as JSON bytes.""" + ctx = _create_context() + + with app.app_context(), use_context(ctx): + _set_user(42) + PersistentStateImpl.set("prefs", {"theme": "dark"}) + + expected_bytes = json.dumps({"theme": "dark"}).encode() + mock_dao.set.assert_called_once_with( + "test-org.test-ext", "prefs", expected_bytes, user_fk=42 + ) + + +@patch("superset.db") +@patch("superset.extensions.storage.persistent_state_impl.ExtensionStorageDAO") +def test_persistent_state_remove_deletes_entry( + mock_dao: MagicMock, mock_db: MagicMock, app: Flask +) -> None: + """PersistentStateImpl.remove calls DAO delete.""" + ctx = _create_context() + + with app.app_context(), use_context(ctx): + _set_user(42) + PersistentStateImpl.remove("prefs") + + mock_dao.delete.assert_called_once_with("test-org.test-ext", "prefs", user_fk=42) + + +@patch("superset.extensions.storage.persistent_state_impl.ExtensionStorageDAO") +def test_shared_accessor_uses_null_user_fk(mock_dao: MagicMock, app: Flask) -> None: + """SharedPersistentStateAccessor uses user_fk=None for global scope.""" + ctx = _create_context() + mock_dao.get_value.return_value = json.dumps("shared_value").encode() + + accessor = SharedPersistentStateAccessor() + + with app.app_context(), use_context(ctx): + _set_user(42) + result = accessor.get("config") + + mock_dao.get_value.assert_called_once_with( + "test-org.test-ext", "config", user_fk=None + ) + assert result == "shared_value" diff --git a/tests/unit_tests/extensions/storage/test_persistent_state_dao.py b/tests/unit_tests/extensions/storage/test_persistent_state_dao.py new file mode 100644 index 00000000000..a0175c4a8e1 --- /dev/null +++ b/tests/unit_tests/extensions/storage/test_persistent_state_dao.py @@ -0,0 +1,217 @@ +# 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. + +"""Tests for ExtensionStorageDAO — encryption, scoping, and CRUD behavior.""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest +from cryptography.fernet import InvalidToken +from flask import Flask + +from superset.extensions.storage.persistent_state_dao import ExtensionStorageDAO + + +@pytest.fixture +def app() -> Flask: + """Create a minimal Flask app for testing.""" + flask_app = Flask(__name__) + flask_app.config["TESTING"] = True + return flask_app + + +# ── get / get_value ─────────────────────────────────────────────────────────── + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_get_returns_none_when_not_found(mock_db: MagicMock, app: Flask) -> None: + """get() returns None when no entry exists for the given scope.""" + mock_db.session.query.return_value.filter.return_value.first.return_value = None + + with app.app_context(): + result = ExtensionStorageDAO.get("my-ext", "key", user_fk=1) + + assert result is None + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_get_value_returns_raw_bytes_for_unencrypted( + mock_db: MagicMock, app: Flask +) -> None: + """get_value() returns raw bytes unchanged when the entry is not encrypted.""" + entry = MagicMock() + entry.is_encrypted = False + entry.value = b'{"foo": "bar"}' + mock_db.session.query.return_value.filter.return_value.first.return_value = entry + + with app.app_context(): + result = ExtensionStorageDAO.get_value("my-ext", "key", user_fk=1) + + assert result == b'{"foo": "bar"}' + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_get_value_returns_none_when_not_found( + mock_db: MagicMock, app: Flask +) -> None: + """get_value() returns None when the key does not exist.""" + mock_db.session.query.return_value.filter.return_value.first.return_value = None + + with app.app_context(): + result = ExtensionStorageDAO.get_value("my-ext", "missing", user_fk=1) + + assert result is None + + +@patch("superset.extensions.storage.persistent_state_dao._fernet") +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_get_value_decrypts_encrypted_entry( + mock_db: MagicMock, mock_fernet_fn: MagicMock, app: Flask +) -> None: + """get_value() decrypts the stored value for encrypted entries.""" + entry = MagicMock() + entry.is_encrypted = True + entry.value = b"encrypted-bytes" + mock_db.session.query.return_value.filter.return_value.first.return_value = entry + mock_fernet_fn.return_value.decrypt.return_value = b'{"decrypted": true}' + + with app.app_context(): + result = ExtensionStorageDAO.get_value("my-ext", "key", user_fk=1) + + assert result == b'{"decrypted": true}' + mock_fernet_fn.return_value.decrypt.assert_called_once_with(b"encrypted-bytes") + + +@patch("superset.extensions.storage.persistent_state_dao._fernet") +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_get_value_returns_none_on_invalid_token( + mock_db: MagicMock, mock_fernet_fn: MagicMock, app: Flask +) -> None: + """get_value() returns None and logs an error when decryption fails.""" + entry = MagicMock() + entry.is_encrypted = True + entry.value = b"corrupted" + mock_db.session.query.return_value.filter.return_value.first.return_value = entry + mock_fernet_fn.return_value.decrypt.side_effect = InvalidToken() + + with app.app_context(): + result = ExtensionStorageDAO.get_value("my-ext", "key", user_fk=1) + + assert result is None + + +# ── set (upsert) ────────────────────────────────────────────────────────────── + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_set_creates_new_entry_when_absent(mock_db: MagicMock, app: Flask) -> None: + """set() adds a new entry when no existing entry is found.""" + mock_db.session.query.return_value.filter.return_value.first.return_value = None + + with app.app_context(): + ExtensionStorageDAO.set("my-ext", "key", b'{"value": 1}', user_fk=1) + + mock_db.session.add.assert_called_once() + mock_db.session.flush.assert_called_once() + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_set_updates_existing_entry(mock_db: MagicMock, app: Flask) -> None: + """set() updates in-place when an entry already exists (no duplicate row).""" + existing = MagicMock() + mock_db.session.query.return_value.filter.return_value.first.return_value = existing + + with app.app_context(): + ExtensionStorageDAO.set("my-ext", "key", b'{"new": true}', user_fk=1) + + assert existing.value == b'{"new": true}' + mock_db.session.add.assert_not_called() + mock_db.session.flush.assert_called_once() + + +@patch("superset.extensions.storage.persistent_state_dao._fernet") +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_set_encrypts_value_when_requested( + mock_db: MagicMock, mock_fernet_fn: MagicMock, app: Flask +) -> None: + """set() encrypts value bytes and sets is_encrypted=True when requested.""" + mock_db.session.query.return_value.filter.return_value.first.return_value = None + mock_fernet_fn.return_value.encrypt.return_value = b"ciphertext" + + with app.app_context(): + ExtensionStorageDAO.set( + "my-ext", "key", b"plaintext", user_fk=1, is_encrypted=True + ) + + mock_fernet_fn.return_value.encrypt.assert_called_once_with(b"plaintext") + added_entry = mock_db.session.add.call_args[0][0] + assert added_entry.value == b"ciphertext" + assert added_entry.is_encrypted is True + + +# ── delete ──────────────────────────────────────────────────────────────────── + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_delete_returns_true_when_entry_exists( + mock_db: MagicMock, app: Flask +) -> None: + """delete() returns True and removes the row when the entry is found.""" + entry = MagicMock() + mock_db.session.query.return_value.filter.return_value.first.return_value = entry + + with app.app_context(): + result = ExtensionStorageDAO.delete("my-ext", "key", user_fk=1) + + assert result is True + mock_db.session.delete.assert_called_once_with(entry) + mock_db.session.flush.assert_called_once() + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_delete_returns_false_when_not_found( + mock_db: MagicMock, app: Flask +) -> None: + """delete() returns False without touching the session when entry is absent.""" + mock_db.session.query.return_value.filter.return_value.first.return_value = None + + with app.app_context(): + result = ExtensionStorageDAO.delete("my-ext", "key", user_fk=1) + + assert result is False + mock_db.session.delete.assert_not_called() + + +# ── scoping ─────────────────────────────────────────────────────────────────── + + +@patch("superset.extensions.storage.persistent_state_dao.db") +def test_dao_user_and_shared_scopes_issue_independent_queries( + mock_db: MagicMock, app: Flask +) -> None: + """User-scoped (user_fk=N) and shared-scoped (user_fk=None) lookups are separate.""" + first_call = MagicMock(return_value=None) + mock_db.session.query.return_value.filter.return_value.first = first_call + + with app.app_context(): + ExtensionStorageDAO.get("my-ext", "key", user_fk=42) + ExtensionStorageDAO.get("my-ext", "key", user_fk=None) + + # Each scope issues its own independent DB query + assert first_call.call_count == 2