File: C:/Users/fred/anaconda3/Lib/site-packages/streamlit/runtime/connection_factory.py
# Copyright (c) Streamlit Inc. (2018-2022) Snowflake Inc. (2022-2024)
#
# 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.
Config options, credentials, secrets, etc. for connections are taken from various
sources:
- Any connection-specific configuration files.
- An app's ``secrets.toml`` files.
- The kwargs passed to this function.
Parameters
----------
name : str
The connection name used for secrets lookup in ``[connections.<name>]``.
Type will be inferred from passing ``"sql"``, ``"snowflake"``, or ``"snowpark"``.
type : str, connection class, or None
The type of connection to create. It can be a keyword (``"sql"``, ``"snowflake"``,
or ``"snowpark"``), a path to an importable class, or an imported class reference.
All classes must extend ``st.connections.BaseConnection`` and implement the
``_connect()`` method. If the type kwarg is None, a ``type`` field must be set in
the connection's section in ``secrets.toml``.
max_entries : int or None
The maximum number of connections to keep in the cache, or None
for an unbounded cache. (When a new entry is added to a full cache,
the oldest cached entry will be removed.) The default is None.
ttl : float, timedelta, or None
The maximum number of seconds to keep results in the cache, or
None if cached results should not expire. The default is None.
**kwargs : any
Additional connection specific kwargs that are passed to the Connection's
``_connect()`` method. Learn more from the specific Connection's documentation.
Returns
-------
Connection object
An initialized Connection object of the specified type.
Examples
--------
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.
>>> import streamlit as st
>>> conn = st.connection("sql") # [connections.sql] section in secrets.toml.
Creating a SQLConnection with a custom name requires you to explicitly specify the
type. If type is not passed as a kwarg, it must be set in the appropriate section of
``secrets.toml``.
>>> import streamlit as st
>>> conn1 = st.connection(
... "my_sql_connection", type="sql"
... ) # Config section defined in [connections.my_sql_connection].
>>> conn2 = st.connection(
... "my_other_sql_connection"
... ) # type must be set in [connections.my_other_sql_connection].
Passing the full module path to the connection class that you want to use can be
useful, especially when working with a custom connection:
>>> import streamlit as st
>>> conn = st.connection(
... "my_sql_connection", type="streamlit.connections.SQLConnection"
... )
Finally, you can pass the connection class to use directly to this function. Doing
so allows static type checking tools such as ``mypy`` to infer the exact return
type of ``st.connection``.
>>> 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}")