Repository: h2oai/wave
File: py/h2o_wave/server.py (Python)

Code before:
def _session_for(sessions: dict, session_id: str):
    session = sessions.get(session_id, None)
    if session is None:
        session = Expando()
        sessions[session_id] = session
    return session


class Auth:
    """
    Represents authentication information for a given query context. Carries valid information only if single sign on is enabled.
    """

    def __init__(self, username: str, subject: str, access_token: str, refresh_token: str, session_id: str):
        self.username = username
        """The username of the user."""
        self.subject = subject
        """A unique identifier for the user."""
        self.access_token = access_token
        """The access token of the user."""
        self.refresh_token = refresh_token
        """The refresh token of the user."""
        self._session_id = session_id
        """Session identifier. Do not access, internal use only."""
    
    async def force_token_refresh(self):
        """
        Implictly refresh OIDC tokens when needed, e.g. during long-running background jobs.
        """
        async with httpx.AsyncClient(auth=(_config.hub_access_key_id, _config.hub_access_key_secret), verify=False) as http:
            res = await http.get(_config.hub_address + '_auth/refresh', headers={'Wave-Session-ID': self._session_id}) 

            access_token = res.headers.get('Wave-Access-Token', None)
            refresh_token = res.headers.get('Wave-Refresh-Token', None)
            if access_token and refresh_token:
                self.access_token = access_token
                self.refresh_token = refresh_token


class Query:
    """
    Represents the query context.
    The query context is passed to the `@app` handler function whenever a query
    arrives from the browser (page load, user interaction events, etc.).
    The query context contains useful information about the query, including arguments
    `args` (equivalent to URL query strings) and app-level, user-level and client-level state.
    """

    def __init__(
            self,
            site: AsyncSite,

Review comment (suggestion):
Plea: can we also return the obtained `access_token`? It would make this method usable as the async token provider used by client libraries.

e.g.:
```python
from typing_extensions import Protocol

class TokenProvider(Protocol):
    async def __call__(self) -> str:
        pass
```

Code after:
def _session_for(sessions: dict, session_id: str):
    session = sessions.get(session_id, None)
    if session is None:
        session = Expando()
        sessions[session_id] = session
    return session


class Auth:
    """
    Represents authentication information for a given query context. Carries valid information only if single sign on is enabled.
    """

    def __init__(self, username: str, subject: str, access_token: str, refresh_token: str, session_id: str):
        self.username = username
        """The username of the user."""
        self.subject = subject
        """A unique identifier for the user."""
        self.access_token = access_token
        """The access token of the user."""
        self.refresh_token = refresh_token
        """The refresh token of the user."""
        self._session_id = session_id
        """Session identifier. Do not access, internal use only."""
    
    async def ensure_fresh_token(self) -> Optional[str]:
        """
        Explicitly refresh OIDC tokens when needed, e.g. during long-running background jobs.
        """
        async with httpx.AsyncClient(auth=(_config.hub_access_key_id, _config.hub_access_key_secret), verify=False) as http:
            res = await http.get(_config.hub_address + '_auth/refresh', headers={'Wave-Session-ID': self._session_id}) 

            access_token = res.headers.get('Wave-Access-Token', None)
            refresh_token = res.headers.get('Wave-Refresh-Token', None)
            if access_token and refresh_token:
                self.access_token = access_token
                self.refresh_token = refresh_token
            return access_token


class Query:
    """
    Represents the query context.
    The query context is passed to the `@app` handler function whenever a query
    arrives from the browser (page load, user interaction events, etc.).
    The query context contains useful information about the query, including arguments
    `args` (equivalent to URL query strings) and app-level, user-level and client-level state.
    """

    def __init__(
            self,