# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2025) # # Licensed 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. """Declares the CacheStorageContext dataclass, which contains parameter information for each function decorated by `@st.cache_data` (for example: ttl, max_entries etc.) Declares the CacheStorageManager protocol, which implementations are used to create CacheStorage instances and to optionally clear all cache storages, that were created by this manager, and to check if the context is valid for the storage. Declares the CacheStorage protocol, which implementations are used to store cached values for a single `@st.cache_data` decorated function serialized as bytes. How these classes work together ------------------------------- - CacheStorageContext : this is a dataclass that contains the parameters from `@st.cache_data` that are passed to the CacheStorageManager.create() method. - CacheStorageManager : each instance of this is able to create CacheStorage instances, and optionally to clear data of all cache storages. - CacheStorage : each instance of this is able to get, set, delete, and clear entries for a single `@st.cache_data` decorated function. ┌───────────────────────────────┐ │ │ │ CacheStorageManager │ │ │ │ - clear_all(optional) │ │ - check_context │ │ │ └──┬────────────────────────────┘ │ │ ┌──────────────────────┐ │ │ CacheStorage │ │ create(context)│ │ └────────────────► - get │ │ - set │ │ - delete │ │ - close (optional)│ │ - clear │ └──────────────────────┘ """ from __future__ import annotations from abc import abstractmethod from dataclasses import dataclass from typing import Literal, Protocol class CacheStorageError(Exception): """Base exception raised by the cache storage""" class CacheStorageKeyNotFoundError(CacheStorageError): """Raised when the key is not found in the cache storage""" class InvalidCacheStorageContext(CacheStorageError): """Raised if the cache storage manager is not able to work with provided CacheStorageContext. """ @dataclass(frozen=True) class CacheStorageContext: """Context passed to the cache storage during initialization This is the normalized parameters that are passed to CacheStorageManager.create() method. Parameters ---------- function_key: str A hash computed based on function name and source code decorated by `@st.cache_data` function_display_name: str The display name of the function that is decorated by `@st.cache_data` ttl_seconds : float or None The time-to-live for the keys in storage, in seconds. If None, the entry will never expire. max_entries : int or None The maximum number of entries to store in the cache storage. If None, the cache storage will not limit the number of entries. persist : Literal["disk"] or None The persistence mode for the cache storage. Legacy parameter, that used in Streamlit current cache storage implementation. Could be ignored by cache storage implementation, if storage does not support persistence or it persistent by default. """ function_key: str function_display_name: str ttl_seconds: float | None = None max_entries: int | None = None persist: Literal["disk"] | None = None class CacheStorage(Protocol): """Cache storage protocol, that should be implemented by the concrete cache storages. Used to store cached values for a single `@st.cache_data` decorated function serialized as bytes. CacheStorage instances should be created by `CacheStorageManager.create()` method. Notes ----- Threading: The methods of this protocol could be called from multiple threads. This is a responsibility of the concrete implementation to ensure thread safety guarantees. """ @abstractmethod def get(self, key: str) -> bytes: """Returns the stored value for the key. Raises ------ CacheStorageKeyNotFoundError Raised if the key is not in the storage. """ raise NotImplementedError @abstractmethod def set(self, key: str, value: bytes) -> None: """Sets the value for a given key""" raise NotImplementedError @abstractmethod def delete(self, key: str) -> None: """Delete a given key""" raise NotImplementedError @abstractmethod def clear(self) -> None: """Remove all keys for the storage""" raise NotImplementedError def close(self) -> None: """Closes the cache storage, it is optional to implement, and should be used to close open resources, before we delete the storage instance. e.g. close the database connection etc. """ pass class CacheStorageManager(Protocol): """Cache storage manager protocol, that should be implemented by the concrete cache storage managers. It is responsible for: - Creating cache storage instances for the specific decorated functions, - Validating the context for the cache storages. - Optionally clearing all cache storages in optimal way. It should be created during Runtime initialization. """ @abstractmethod def create(self, context: CacheStorageContext) -> CacheStorage: """Creates a new cache storage instance Please note that the ttl, max_entries and other context fields are specific for whole storage, not for individual key. Notes ----- Threading: Should be safe to call from any thread. """ raise NotImplementedError def clear_all(self) -> None: """Remove everything what possible from the cache storages in optimal way. meaningful default behaviour is to raise NotImplementedError, so this is not abstractmethod. The method is optional to implement: cache data API will fall back to remove all available storages one by one via storage.clear() method if clear_all raises NotImplementedError. Raises ------ NotImplementedError Raised if the storage manager does not provide an ability to clear all storages at once in optimal way. Notes ----- Threading: This method could be called from multiple threads. This is a responsibility of the concrete implementation to ensure thread safety guarantees. """ raise NotImplementedError def check_context(self, context: CacheStorageContext) -> None: """Checks if the context is valid for the storage manager. This method should not return anything, but log message or raise an exception if the context is invalid. In case of raising an exception, we not handle it and let the exception to be propagated. check_context is called only once at the moment of creating `@st.cache_data` decorator for specific function, so it is not called for every cache hit. Parameters ---------- context: CacheStorageContext The context to check for the storage manager, dummy function_key in context will be used, since it is not computed at the point of calling this method. Raises ------ InvalidCacheStorageContext Raised if the cache storage manager is not able to work with provided CacheStorageContext. When possible we should log message instead, since this exception will be propagated to the user. Notes ----- Threading: Should be safe to call from any thread. """ pass