"""
This type stub file was generated by pyright.
"""

from contextlib import contextmanager
from typing import TYPE_CHECKING
from streamlit.connections import BaseConnection
from collections.abc import Iterator
from datetime import timedelta
from pandas import DataFrame
from snowflake.snowpark.session import Session

if TYPE_CHECKING:
    ...
_REQUIRED_CONNECTION_PARAMS = ...
class SnowparkConnection(BaseConnection["Session"]):
    """A connection to Snowpark using snowflake.snowpark.session.Session. Initialize using
    ``st.connection("<name>", type="snowpark")``.

    In addition to providing access to the Snowpark Session, SnowparkConnection supports
    direct SQL querying using ``query("...")`` and thread safe access using
    ``with conn.safe_session():``. See methods below for more information.
    SnowparkConnections should always be created using ``st.connection()``, **not**
    initialized directly.

    .. note::
        We don't expect this iteration of SnowparkConnection to be able to scale
        well in apps with many concurrent users due to the lock contention that will occur
        over the single underlying Session object under high load.
    """
    def __init__(self, connection_name: str, **kwargs) -> None:
        ...
    
    def query(self, sql: str, ttl: float | int | timedelta | None = ...) -> DataFrame:
        """Run a read-only SQL query.

        This method implements both query result caching (with caching behavior
        identical to that of using ``@st.cache_data``) as well as simple error handling/retries.

        .. note::
            Queries that are run without a specified ttl are cached indefinitely.

        Parameters
        ----------
        sql : str
            The read-only SQL query to execute.
        ttl : float, int, 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.

        Returns
        -------
        pandas.DataFrame
            The result of running the query, formatted as a pandas DataFrame.

        Example
        -------
        >>> import streamlit as st
        >>>
        >>> conn = st.connection("snowpark")
        >>> df = conn.query("SELECT * FROM pet_owners")
        >>> st.dataframe(df)
        """
        ...
    
    @property
    def session(self) -> Session:
        """Access the underlying Snowpark session.

        .. note::
            Snowpark sessions are **not** thread safe. Users of this method are
            responsible for ensuring that access to the session returned by this method is
            done in a thread-safe manner. For most users, we recommend using the thread-safe
            safe_session() method and a ``with`` block.

        Information on how to use Snowpark sessions can be found in the `Snowpark documentation
        <https://docs.snowflake.com/en/developer-guide/snowpark/python/working-with-dataframes>`_.

        Example
        -------
        >>> import streamlit as st
        >>>
        >>> session = st.connection("snowpark").session
        >>> df = session.table("mytable").limit(10).to_pandas()
        >>> st.dataframe(df)
        """
        ...
    
    @contextmanager
    def safe_session(self) -> Iterator[Session]:
        """Grab the underlying Snowpark session in a thread-safe manner.

        As operations on a Snowpark session are not thread safe, we need to take care
        when using a session in the context of a Streamlit app where each script run
        occurs in its own thread. Using the contextmanager pattern to do this ensures
        that access on this connection's underlying Session is done in a thread-safe
        manner.

        Information on how to use Snowpark sessions can be found in the `Snowpark documentation
        <https://docs.snowflake.com/en/developer-guide/snowpark/python/working-with-dataframes>`_.

        Example
        -------
        >>> import streamlit as st
        >>>
        >>> conn = st.connection("snowpark")
        >>> with conn.safe_session() as session:
        ...     df = session.table("mytable").limit(10).to_pandas()
        >>>
        >>> st.dataframe(df)
        """
        ...
    


