streamlit-nightly 1.43.2.dev20250307__py3-none-any.whl

streamlit/runtime/caching/storage/local_disk_cache_storage.py

@@ -0,0 +1,223 @@
+# 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 LocalDiskCacheStorageManager class, which is used
+to create LocalDiskCacheStorage instances wrapped by InMemoryCacheStorageWrapper,
+InMemoryCacheStorageWrapper wrapper allows to have first layer of in-memory cache,
+before accessing to LocalDiskCacheStorage itself.
+
+Declares the LocalDiskCacheStorage class, which is used to store cached
+values on disk.
+
+How these classes work together
+-------------------------------
+
+- LocalDiskCacheStorageManager : each instance of this is able
+to create LocalDiskCacheStorage instances wrapped by InMemoryCacheStorageWrapper,
+and to clear data from cache storage folder. It is also LocalDiskCacheStorageManager
+responsibility to check if the context is valid for the storage, and to log warning
+if the context is not valid.
+
+- LocalDiskCacheStorage : each instance of this is able to get, set, delete, and clear
+entries from disk for a single `@st.cache_data` decorated function if `persist="disk"`
+is used in CacheStorageContext.
+
+
+    ┌───────────────────────────────┐
+    │  LocalDiskCacheStorageManager │
+    │                               │
+    │     - clear_all               │
+    │     - check_context           │
+    │                               │
+    └──┬────────────────────────────┘
+       │
+       │                ┌──────────────────────────────┐
+       │                │                              │
+       │ create(context)│  InMemoryCacheStorageWrapper │
+       └────────────────►                              │
+                        │    ┌─────────────────────┐   │
+                        │    │                     │   │
+                        │    │   LocalDiskStorage  │   │
+                        │    │                     │   │
+                        │    └─────────────────────┘   │
+                        │                              │
+                        └──────────────────────────────┘
+
+"""
+
+from __future__ import annotations
+
+import math
+import os
+import shutil
+from typing import Final
+
+from streamlit import errors
+from streamlit.file_util import get_streamlit_file_path, streamlit_read, streamlit_write
+from streamlit.logger import get_logger
+from streamlit.runtime.caching.storage.cache_storage_protocol import (
+    CacheStorage,
+    CacheStorageContext,
+    CacheStorageError,
+    CacheStorageKeyNotFoundError,
+    CacheStorageManager,
+)
+from streamlit.runtime.caching.storage.in_memory_cache_storage_wrapper import (
+    InMemoryCacheStorageWrapper,
+)
+
+_LOGGER: Final = get_logger(__name__)
+
+# Streamlit directory where persisted @st.cache_data objects live.
+# (This is the same directory that @st.cache persisted objects live.
+# But @st.cache_data uses a different extension, so they don't overlap.)
+_CACHE_DIR_NAME: Final = "cache"
+
+# The extension for our persisted @st.cache_data objects.
+# (`@st.cache_data` was originally called `@st.memo`)
+_CACHED_FILE_EXTENSION: Final = "memo"
+
+
+class LocalDiskCacheStorageManager(CacheStorageManager):
+    def create(self, context: CacheStorageContext) -> CacheStorage:
+        """Creates a new cache storage instance wrapped with in-memory cache layer"""
+        persist_storage = LocalDiskCacheStorage(context)
+        return InMemoryCacheStorageWrapper(
+            persist_storage=persist_storage, context=context
+        )
+
+    def clear_all(self) -> None:
+        cache_path = get_cache_folder_path()
+        if os.path.isdir(cache_path):
+            shutil.rmtree(cache_path)
+
+    def check_context(self, context: CacheStorageContext) -> None:
+        if (
+            context.persist == "disk"
+            and context.ttl_seconds is not None
+            and not math.isinf(context.ttl_seconds)
+        ):
+            _LOGGER.warning(
+                f"The cached function '{context.function_display_name}' has a TTL "
+                "that will be ignored. Persistent cached functions currently don't "
+                "support TTL."
+            )
+
+
+class LocalDiskCacheStorage(CacheStorage):
+    """Cache storage that persists data to disk
+    This is the default cache persistence layer for `@st.cache_data`
+    """
+
+    def __init__(self, context: CacheStorageContext):
+        self.function_key = context.function_key
+        self.persist = context.persist
+        self._ttl_seconds = context.ttl_seconds
+        self._max_entries = context.max_entries
+
+    @property
+    def ttl_seconds(self) -> float:
+        return self._ttl_seconds if self._ttl_seconds is not None else math.inf
+
+    @property
+    def max_entries(self) -> float:
+        return float(self._max_entries) if self._max_entries is not None else math.inf
+
+    def get(self, key: str) -> bytes:
+        """
+        Returns the stored value for the key if persisted,
+        raise CacheStorageKeyNotFoundError if not found, or not configured
+        with persist="disk"
+        """
+        if self.persist == "disk":
+            path = self._get_cache_file_path(key)
+            try:
+                with streamlit_read(path, binary=True) as input:
+                    value = input.read()
+                    _LOGGER.debug("Disk cache HIT: %s", key)
+                    return bytes(value)
+            except FileNotFoundError:
+                raise CacheStorageKeyNotFoundError("Key not found in disk cache")
+            except Exception as ex:
+                _LOGGER.exception("Error reading from cache")
+                raise CacheStorageError("Unable to read from cache") from ex
+        else:
+            raise CacheStorageKeyNotFoundError(
+                f"Local disk cache storage is disabled (persist={self.persist})"
+            )
+
+    def set(self, key: str, value: bytes) -> None:
+        """Sets the value for a given key"""
+        if self.persist == "disk":
+            path = self._get_cache_file_path(key)
+            try:
+                with streamlit_write(path, binary=True) as output:
+                    output.write(value)
+            except errors.Error as ex:
+                _LOGGER.debug("Unable to write to cache", exc_info=ex)
+                # Clean up file so we don't leave zero byte files.
+                try:
+                    os.remove(path)
+                except (FileNotFoundError, OSError):
+                    # If we can't remove the file, it's not a big deal.
+                    pass
+                raise CacheStorageError("Unable to write to cache") from ex
+
+    def delete(self, key: str) -> None:
+        """Delete a cache file from disk. If the file does not exist on disk,
+        return silently. If another exception occurs, log it. Does not throw.
+        """
+        if self.persist == "disk":
+            path = self._get_cache_file_path(key)
+            try:
+                os.remove(path)
+            except FileNotFoundError:
+                # The file is already removed.
+                pass
+            except Exception as ex:
+                _LOGGER.exception(
+                    "Unable to remove a file from the disk cache", exc_info=ex
+                )
+
+    def clear(self) -> None:
+        """Delete all keys for the current storage"""
+        cache_dir = get_cache_folder_path()
+
+        if os.path.isdir(cache_dir):
+            # We try to remove all files in the cache directory that start with
+            # the function key, whether `clear` called for `self.persist`
+            # storage or not, to avoid leaving orphaned files in the cache directory.
+            for file_name in os.listdir(cache_dir):
+                if self._is_cache_file(file_name):
+                    os.remove(os.path.join(cache_dir, file_name))
+
+    def close(self) -> None:
+        """Dummy implementation of close, we don't need to actually "close" anything"""
+
+    def _get_cache_file_path(self, value_key: str) -> str:
+        """Return the path of the disk cache file for the given value."""
+        cache_dir = get_cache_folder_path()
+        return os.path.join(
+            cache_dir, f"{self.function_key}-{value_key}.{_CACHED_FILE_EXTENSION}"
+        )
+
+    def _is_cache_file(self, fname: str) -> bool:
+        """Return true if the given file name is a cache file for this storage."""
+        return fname.startswith(f"{self.function_key}-") and fname.endswith(
+            f".{_CACHED_FILE_EXTENSION}"
+        )
+
+
+def get_cache_folder_path() -> str:
+    return get_streamlit_file_path(_CACHE_DIR_NAME)

streamlit/runtime/connection_factory.py

@@ -0,0 +1,436 @@
+# 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.
+
+from __future__ import annotations
+
+import os
+import re
+from typing import TYPE_CHECKING, Any, Final, Literal, TypeVar, overload
+
+from streamlit.connections import (
+    BaseConnection,
+    SnowflakeConnection,
+    SnowparkConnection,
+    SQLConnection,
+)
+from streamlit.deprecation_util import deprecate_obj_name
+from streamlit.errors import StreamlitAPIException
+from streamlit.runtime.caching import cache_resource
+from streamlit.runtime.metrics_util import gather_metrics
+from streamlit.runtime.secrets import secrets_singleton
+
+if TYPE_CHECKING:
+    from datetime import timedelta
+
+# NOTE: Adding support for a new first party connection requires:
+#   1. Adding the new connection name and class to this dict.
+#   2. Writing two new @overloads for connection_factory (one for the case where the
+#      only the connection name is specified and another when both name and type are).
+#   3. Updating test_get_first_party_connection_helper in connection_factory_test.py.
+FIRST_PARTY_CONNECTIONS = {
+    "snowflake": SnowflakeConnection,
+    "snowpark": SnowparkConnection,
+    "sql": SQLConnection,
+}
+MODULE_EXTRACTION_REGEX = re.compile(r"No module named \'(.+)\'")
+MODULES_TO_PYPI_PACKAGES: Final[dict[str, str]] = {
+    "MySQLdb": "mysqlclient",
+    "psycopg2": "psycopg2-binary",
+    "sqlalchemy": "sqlalchemy",
+    "snowflake": "snowflake-connector-python",
+    "snowflake.connector": "snowflake-connector-python",
+    "snowflake.snowpark": "snowflake-snowpark-python",
+}
+
+# The BaseConnection bound is parameterized to `Any` below as subclasses of
+# BaseConnection are responsible for binding the type parameter of BaseConnection to a
+# concrete type, but the type it gets bound to isn't important to us here.
+ConnectionClass = TypeVar("ConnectionClass", bound=BaseConnection[Any])
+
+
+@gather_metrics("connection")
+def _create_connection(
+    name: str,
+    connection_class: type[ConnectionClass],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    **kwargs,
+) -> ConnectionClass:
+    """Create an instance of connection_class with the given name and kwargs.
+
+    The weird implementation of this function with the @cache_resource annotated
+    function defined internally is done to:
+      - Always @gather_metrics on the call even if the return value is a cached one.
+      - Allow the user to specify ttl and max_entries when calling st.connection.
+    """
+
+    def __create_connection(
+        name: str, connection_class: type[ConnectionClass], **kwargs
+    ) -> ConnectionClass:
+        return connection_class(connection_name=name, **kwargs)
+
+    if not issubclass(connection_class, BaseConnection):
+        raise StreamlitAPIException(
+            f"{connection_class} is not a subclass of BaseConnection!"
+        )
+
+    # We modify our helper function's `__qualname__` here to work around default
+    # `@st.cache_resource` behavior. Otherwise, `st.connection` being called with
+    # different `ttl` or `max_entries` values will reset the cache with each call.
+    ttl_str = str(ttl).replace(  # Avoid adding extra `.` characters to `__qualname__`
+        ".", "_"
+    )
+    __create_connection.__qualname__ = (
+        f"{__create_connection.__qualname__}_{ttl_str}_{max_entries}"
+    )
+    __create_connection = cache_resource(
+        max_entries=max_entries,
+        show_spinner="Running `st.connection(...)`.",
+        ttl=ttl,
+    )(__create_connection)
+
+    return __create_connection(name, connection_class, **kwargs)
+
+
+def _get_first_party_connection(connection_class: str):
+    if connection_class in FIRST_PARTY_CONNECTIONS:
+        return FIRST_PARTY_CONNECTIONS[connection_class]
+
+    raise StreamlitAPIException(
+        f"Invalid connection '{connection_class}'. "
+        f"Supported connection classes: {FIRST_PARTY_CONNECTIONS}"
+    )
+
+
+@overload
+def connection_factory(
+    name: Literal["sql"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    autocommit: bool = False,
+    **kwargs,
+) -> SQLConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: str,
+    type: Literal["sql"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    autocommit: bool = False,
+    **kwargs,
+) -> SQLConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: Literal["snowflake"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    autocommit: bool = False,
+    **kwargs,
+) -> SnowflakeConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: str,
+    type: Literal["snowflake"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    autocommit: bool = False,
+    **kwargs,
+) -> SnowflakeConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: Literal["snowpark"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    **kwargs,
+) -> SnowparkConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: str,
+    type: Literal["snowpark"],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    **kwargs,
+) -> SnowparkConnection:
+    pass
+
+
+@overload
+def connection_factory(
+    name: str,
+    type: type[ConnectionClass],
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    **kwargs,
+) -> ConnectionClass:
+    pass
+
+
+@overload
+def connection_factory(
+    name: str,
+    type: str | None = None,
+    max_entries: int | None = None,
+    ttl: float | timedelta | None = None,
+    **kwargs,
+) -> BaseConnection[Any]:
+    pass
+
+
+def connection_factory(
+    name,
+    type=None,
+    max_entries=None,
+    ttl=None,
+    **kwargs,
+):
+    """Create a new connection to a data store or API, or return an existing one.
+
+    Configuration options, credentials, and secrets for connections are
+    combined from the following sources:
+
+    - The keyword arguments passed to this command.
+    - The app's ``secrets.toml`` files.
+    - Any connection-specific configuration files.
+
+    The connection returned from ``st.connection`` is internally cached with
+    ``st.cache_resource`` and is therefore shared between sessions.
+
+    Parameters
+    ----------
+    name : str
+        The connection name used for secrets lookup in ``secrets.toml``.
+        Streamlit uses secrets under ``[connections.<name>]`` for the
+        connection. ``type`` will be inferred if ``name`` is one of the
+        following: ``"snowflake"``, ``"snowpark"``, or ``"sql"``.
+
+    type : str, connection class, or None
+        The type of connection to create. This can be one of the following:
+
+        - ``None`` (default): Streamlit will infer the connection type from
+          ``name``. If the type is not inferrable from ``name``, the type must
+          be specified in ``secrets.toml`` instead.
+        - ``"snowflake"``: Streamlit will initialize a connection with
+          |SnowflakeConnection|_.
+        - ``"snowpark"``: Streamlit will initialize a connection with
+          |SnowparkConnection|_. This is deprecated.
+        - ``"sql"``: Streamlit will initialize a connection with
+          |SQLConnection|_.
+        - A string path to an importable class: This must be a dot-separated
+          module path ending in the importable class. Streamlit will import the
+          class and initialize a connection with it. The class must extend
+          ``st.connections.BaseConnection``.
+        - An imported class reference: Streamlit will initialize a connection
+          with the referenced class, which must extend
+          ``st.connections.BaseConnection``.
+
+        .. |SnowflakeConnection| replace:: ``SnowflakeConnection``
+        .. _SnowflakeConnection: https://docs.streamlit.io/develop/api-reference/connections/st.connections.snowflakeconnection
+        .. |SnowparkConnection| replace:: ``SnowparkConnection``
+        .. _SnowparkConnection: https://docs.streamlit.io/develop/api-reference/connections/st.connections.snowparkconnection
+        .. |SQLConnection| replace:: ``SQLConnection``
+        .. _SQLConnection: https://docs.streamlit.io/develop/api-reference/connections/st.connections.sqlconnection
+
+    max_entries : int or None
+        The maximum number of connections to keep in the cache.
+        If this is ``None`` (default), the cache is unbounded. Otherwise, when
+        a new entry is added to a full cache, the oldest cached entry is
+        removed.
+    ttl : float, timedelta, or None
+        The maximum number of seconds to keep results in the cache.
+        If this is ``None`` (default), cached results do not expire with time.
+    **kwargs : any
+        Connection-specific keyword arguments that are passed to the
+        connection's ``._connect()`` method. ``**kwargs`` are typically
+        combined with (and take precendence over) key-value pairs in
+        ``secrets.toml``. To learn more, see the specific connection's
+        documentation.
+
+    Returns
+    -------
+    Subclass of BaseConnection
+        An initialized connection object of the specified ``type``.
+
+    Examples
+    --------
+
+    **Example 1: Inferred connection type**
+
+    The easiest way to create a first-party (SQL, Snowflake, or Snowpark) connection is
+    to use their default names and define corresponding sections in your ``secrets.toml``
+    file. The following example creates a ``"sql"``-type connection.
+
+    ``.streamlit/secrets.toml``:
+
+    >>> [connections.sql]
+    >>> dialect = "xxx"
+    >>> host = "xxx"
+    >>> username = "xxx"
+    >>> password = "xxx"
+
+    Your app code:
+
+    >>> import streamlit as st
+    >>> conn = st.connection("sql")
+
+    **Example 2: Named connections**
+
+    Creating a connection with a custom name requires you to explicitly
+    specify the type. If ``type`` is not passed as a keyword argument, it must
+    be set in the appropriate section of ``secrets.toml``. The following
+    example creates two ``"sql"``-type connections, each with their own
+    custom name. The first defines ``type`` in the ``st.connection`` command;
+    the second defines ``type`` in ``secrets.toml``.
+
+    ``.streamlit/secrets.toml``:
+
+    >>> [connections.first_connection]
+    >>> dialect = "xxx"
+    >>> host = "xxx"
+    >>> username = "xxx"
+    >>> password = "xxx"
+    >>>
+    >>> [connections.second_connection]
+    >>> type = "sql"
+    >>> dialect = "yyy"
+    >>> host = "yyy"
+    >>> username = "yyy"
+    >>> password = "yyy"
+
+    Your app code:
+
+    >>> import streamlit as st
+    >>> conn1 = st.connection("first_connection", type="sql")
+    >>> conn2 = st.connection("second_connection")
+
+    **Example 3: Using a path to the connection class**
+
+    Passing the full module path to the connection class can be useful,
+    especially when working with a custom connection. Although this is not the
+    typical way to create first party connections, the following example
+    creates the same type of connection as one with ``type="sql"``. Note that
+    ``type`` is a string path.
+
+    ``.streamlit/secrets.toml``:
+
+    >>> [connections.my_sql_connection]
+    >>> url = "xxx+xxx://xxx:xxx@xxx:xxx/xxx"
+
+    Your app code:
+
+    >>> import streamlit as st
+    >>> conn = st.connection(
+    ...     "my_sql_connection", type="streamlit.connections.SQLConnection"
+    ... )
+
+    **Example 4: Importing the connection class**
+
+    You can pass the connection class directly to the ``st.connection``
+    command. Doing so allows static type checking tools such as ``mypy`` to
+    infer the exact return type of ``st.connection``. The following example
+    creates the same connection as in Example 3.
+
+    ``.streamlit/secrets.toml``:
+
+    >>> [connections.my_sql_connection]
+    >>> url = "xxx+xxx://xxx:xxx@xxx:xxx/xxx"
+
+    Your app code:
+
+    >>> import streamlit as st
+    >>> from streamlit.connections import SQLConnection
+    >>> conn = st.connection("my_sql_connection", type=SQLConnection)
+
+    """
+    USE_ENV_PREFIX = "env:"
+
+    if name.startswith(USE_ENV_PREFIX):
+        # It'd be nice to use str.removeprefix() here, but we won't be able to do that
+        # until the minimium Python version we support is 3.9.
+        envvar_name = name[len(USE_ENV_PREFIX) :]
+        name = os.environ[envvar_name]
+
+    if type is None:
+        if name in FIRST_PARTY_CONNECTIONS:
+            # We allow users to simply write `st.connection("sql")` instead of
+            # `st.connection("sql", type="sql")`.
+            type = _get_first_party_connection(name)
+        else:
+            # The user didn't specify a type, so we try to pull it out from their
+            # secrets.toml file. NOTE: we're okay with any of the dict lookups below
+            # exploding with a KeyError since, if type isn't explicitly specified here,
+            # it must be the case that it's defined in secrets.toml and should raise an
+            # Exception otherwise.
+            secrets_singleton.load_if_toml_exists()
+            type = secrets_singleton["connections"][name]["type"]
+
+    # type is a nice kwarg name for the st.connection user but is annoying to work with
+    # since it conflicts with the builtin function name and thus gets syntax
+    # highlighted.
+    connection_class = type
+
+    if isinstance(connection_class, str):
+        # We assume that a connection_class specified via string is either the fully
+        # qualified name of a class (its module and exported classname) or the string
+        # literal shorthand for one of our first party connections. In the former case,
+        # connection_class will always contain a "." in its name.
+        if "." in connection_class:
+            parts = connection_class.split(".")
+            classname = parts.pop()
+
+            import importlib
+
+            connection_module = importlib.import_module(".".join(parts))
+            connection_class = getattr(connection_module, classname)
+        else:
+            connection_class = _get_first_party_connection(connection_class)
+
+    # At this point, connection_class should be of type Type[ConnectionClass].
+    try:
+        conn = _create_connection(
+            name, connection_class, max_entries=max_entries, ttl=ttl, **kwargs
+        )
+        if isinstance(conn, SnowparkConnection):
+            conn = deprecate_obj_name(
+                conn,
+                'connection("snowpark")',
+                'connection("snowflake")',
+                "2024-04-01",
+            )
+        return conn
+    except ModuleNotFoundError as e:
+        err_string = str(e)
+        missing_module = re.search(MODULE_EXTRACTION_REGEX, err_string)
+
+        extra_info = "You may be missing a dependency required to use this connection."
+        if missing_module:
+            pypi_package = MODULES_TO_PYPI_PACKAGES.get(missing_module.group(1))
+            if pypi_package:
+                extra_info = f"You need to install the '{pypi_package}' package to use this connection."
+
+        raise ModuleNotFoundError(f"{str(e)}. {extra_info}")