wip

Author: cofin

sqlspec/extensions/litestar/session.py

@@ -3,18 +3,62 @@
 from typing import TYPE_CHECKING, Any, Optional, Union
 
 from litestar.middleware.session.base import BaseSessionBackend
+from litestar.types import Scopes
 
-from sqlspec.extensions.litestar.store import SessionStore
+from sqlspec.extensions.litestar.store import SQLSpecSessionStore
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
     from litestar.connection import ASGIConnection
+    from litestar.types import Message, ScopeSession
 
     from sqlspec.config import AsyncConfigT, DatabaseConfigProtocol, SyncConfigT
 
 logger = get_logger("extensions.litestar.session")
 
-__all__ = ("SQLSpecSessionBackend",)
+__all__ = ("SQLSpecSessionBackend", "SQLSpecSessionConfig")
+
+
+class SQLSpecSessionConfig:
+    """Configuration for SQLSpec session backend."""
+
+    def __init__(
+        self,
+        key: str = "session",
+        max_age: int = 1209600,  # 14 days
+        path: str = "/",
+        domain: Optional[str] = None,
+        secure: bool = False,
+        httponly: bool = True,
+        samesite: str = "lax",
+        exclude: Optional[Union[str, list[str]]] = None,
+        exclude_opt_key: str = "skip_session",
+        scopes: Scopes = frozenset({"http", "websocket"}),
+    ) -> None:
+        """Initialize session configuration.
+
+        Args:
+            key: Cookie key name
+            max_age: Cookie max age in seconds
+            path: Cookie path
+            domain: Cookie domain
+            secure: Require HTTPS for cookie
+            httponly: Make cookie HTTP-only
+            samesite: SameSite policy for cookie
+            exclude: Patterns to exclude from session middleware
+            exclude_opt_key: Key to opt out of session middleware
+            scopes: Scopes where session middleware applies
+        """
+        self.key = key
+        self.max_age = max_age
+        self.path = path
+        self.domain = domain
+        self.secure = secure
+        self.httponly = httponly
+        self.samesite = samesite
+        self.exclude = exclude
+        self.exclude_opt_key = exclude_opt_key
+        self.scopes = scopes
 
 
 class SQLSpecSessionBackend(BaseSessionBackend):
@@ -24,7 +68,7 @@ class SQLSpecSessionBackend(BaseSessionBackend):
     middleware, providing transparent session management with database persistence.
     """
 
-    __slots__ = ("_session_id_generator", "_session_lifetime", "_store")
+    __slots__ = ("_session_id_generator", "_session_lifetime", "_store", "config")
 
     def __init__(
         self,
@@ -36,6 +80,7 @@ def __init__(
         expires_at_column: str = "expires_at",
         created_at_column: str = "created_at",
         session_lifetime: int = 24 * 60 * 60,  # 24 hours
+        session_config: Optional[SQLSpecSessionConfig] = None,
     ) -> None:
         """Initialize the session backend.
 
@@ -47,17 +92,19 @@ def __init__(
             expires_at_column: Name of the expires at column
             created_at_column: Name of the created at column
             session_lifetime: Default session lifetime in seconds
+            session_config: Session configuration for middleware
         """
-        self._store = SessionStore(
+        self._store = SQLSpecSessionStore(
             config,
             table_name=table_name,
             session_id_column=session_id_column,
             data_column=data_column,
             expires_at_column=expires_at_column,
             created_at_column=created_at_column,
         )
-        self._session_id_generator = SessionStore.generate_session_id
+        self._session_id_generator = SQLSpecSessionStore.generate_session_id
         self._session_lifetime = session_lifetime
+        self.config = session_config or SQLSpecSessionConfig()
 
     async def load_from_connection(self, connection: "ASGIConnection[Any, Any, Any, Any]") -> dict[str, Any]:
         """Load session data from the connection.
@@ -110,9 +157,112 @@ def get_session_id(self, connection: "ASGIConnection[Any, Any, Any, Any]") -> Op
         Returns:
             Session identifier if found
         """
-        # Look for session ID in cookies
-        session_cookie_name = getattr(connection.app.session_config, "session_cookie_name", "session")  # type: ignore[union-attr]
-        return connection.cookies.get(session_cookie_name)
+        # Try to get session ID from cookies using the config key
+        session_id = connection.cookies.get(self.config.key)
+        if session_id and session_id != "null":
+            return session_id
+
+        # Fallback to getting session ID from connection state
+        session_id = connection.get_session_id()
+        if session_id:
+            return session_id
+
+        return None
+
+    async def store_in_message(
+        self, scope_session: "ScopeSession", message: "Message", connection: "ASGIConnection[Any, Any, Any, Any]"
+    ) -> None:
+        """Store session information in the outgoing message.
+
+        For server-side sessions, this method sets a cookie containing the session ID.
+        If the session is empty, a null-cookie will be set to clear any existing session.
+
+        Args:
+            scope_session: Current session data to store
+            message: Outgoing ASGI message to modify
+            connection: ASGI connection instance
+        """
+        if message["type"] != "http.response.start":
+            return
+
+        cookie_key = self.config.key
+
+        # If session is empty, set a null cookie to clear any existing session
+        if not scope_session:
+            cookie_value = self._build_cookie_value(
+                key=cookie_key,
+                value="null",
+                max_age=0,
+                path=self.config.path,
+                domain=self.config.domain,
+                secure=self.config.secure,
+                httponly=self.config.httponly,
+                samesite=self.config.samesite,
+            )
+            self._add_cookie_to_message(message, cookie_value)
+            return
+
+        # Get or generate session ID
+        session_id = self.get_session_id(connection)
+        if not session_id:
+            session_id = self._session_id_generator()
+
+        # Store session data in the backend
+        try:
+            await self._store.set(session_id, scope_session, expires_in=self._session_lifetime)
+        except Exception:
+            logger.exception("Failed to store session data for session %s", session_id)
+            # Don't set the cookie if we failed to store the data
+            return
+
+        # Set the session ID cookie
+        cookie_value = self._build_cookie_value(
+            key=cookie_key,
+            value=session_id,
+            max_age=self.config.max_age,
+            path=self.config.path,
+            domain=self.config.domain,
+            secure=self.config.secure,
+            httponly=self.config.httponly,
+            samesite=self.config.samesite,
+        )
+        self._add_cookie_to_message(message, cookie_value)
+
+    def _build_cookie_value(
+        self,
+        key: str,
+        value: str,
+        max_age: Optional[int] = None,
+        path: Optional[str] = None,
+        domain: Optional[str] = None,
+        secure: bool = False,
+        httponly: bool = False,
+        samesite: Optional[str] = None,
+    ) -> str:
+        """Build a cookie value string with attributes."""
+        cookie_parts = [f"{key}={value}"]
+
+        if path:
+            cookie_parts.append(f"Path={path}")
+        if domain:
+            cookie_parts.append(f"Domain={domain}")
+        if max_age is not None:
+            cookie_parts.append(f"Max-Age={max_age}")
+        if secure:
+            cookie_parts.append("Secure")
+        if httponly:
+            cookie_parts.append("HttpOnly")
+        if samesite:
+            cookie_parts.append(f"SameSite={samesite}")
+
+        return "; ".join(cookie_parts)
+
+    def _add_cookie_to_message(self, message: "Message", cookie_value: str) -> None:
+        """Add a Set-Cookie header to the ASGI message."""
+        if message["type"] == "http.response.start":
+            headers = list(message.get("headers", []))
+            headers.append([b"set-cookie", cookie_value.encode()])
+            message["headers"] = headers
 
     async def delete_session(self, session_id: str) -> None:
         """Delete a session.
@@ -152,7 +302,7 @@ async def get_all_session_ids(self) -> list[str]:
         return session_ids
 
     @property
-    def store(self) -> SessionStore:
+    def store(self) -> SQLSpecSessionStore:
         """Get the underlying session store.
 
         Returns:

sqlspec/extensions/litestar/store.py

@@ -22,14 +22,14 @@
 
 logger = get_logger("extensions.litestar.store")
 
-__all__ = ("SessionStore", "SessionStoreError")
+__all__ = ("SQLSpecSessionStore", "SQLSpecSessionStoreError")
 
 
-class SessionStoreError(SQLSpecError):
+class SQLSpecSessionStoreError(SQLSpecError):
     """Exception raised by session store operations."""
 
 
-class SessionStore(Store):
+class SQLSpecSessionStore(Store):
     """SQLSpec-based session store for Litestar.
 
     This store uses SQLSpec's builder API to create dialect-aware SQL operations
@@ -129,7 +129,7 @@ async def _ensure_table_exists(self, driver: Union[SyncDriverAdapterBase, AsyncD
         except Exception as e:
             msg = f"Failed to create session table: {e}"
             logger.exception("Failed to create session table %s", self._table_name)
-            raise SessionStoreError(msg) from e
+            raise SQLSpecSessionStoreError(msg) from e
 
     def _get_dialect_upsert_sql(self, dialect: str, session_id: str, data: str, expires_at: datetime) -> Any:
         """Generate dialect-specific upsert SQL using SQL builder API.
@@ -152,12 +152,10 @@ def _get_dialect_upsert_sql(self, dialect: str, session_id: str, data: str, expi
                 .columns(self._session_id_column, self._data_column, self._expires_at_column, self._created_at_column)
                 .values(session_id, data, expires_at, current_time)
                 .on_conflict(self._session_id_column)
-                .do_update(
-                    **{
-                        self._data_column: sql.raw("EXCLUDED." + self._data_column),
-                        self._expires_at_column: sql.raw("EXCLUDED." + self._expires_at_column),
-                    }
-                )
+                .do_update(**{
+                    self._data_column: sql.raw("EXCLUDED." + self._data_column),
+                    self._expires_at_column: sql.raw("EXCLUDED." + self._expires_at_column),
+                })
             )
 
         if dialect in {"mysql", "mariadb"}:
@@ -166,12 +164,10 @@ def _get_dialect_upsert_sql(self, dialect: str, session_id: str, data: str, expi
                 sql.insert(self._table_name)
                 .columns(self._session_id_column, self._data_column, self._expires_at_column, self._created_at_column)
                 .values(session_id, data, expires_at, current_time)
-                .on_duplicate_key_update(
-                    **{
-                        self._data_column: sql.raw(f"VALUES({self._data_column})"),
-                        self._expires_at_column: sql.raw(f"VALUES({self._expires_at_column})"),
-                    }
-                )
+                .on_duplicate_key_update(**{
+                    self._data_column: sql.raw(f"VALUES({self._data_column})"),
+                    self._expires_at_column: sql.raw(f"VALUES({self._expires_at_column})"),
+                })
             )
 
         if dialect == "sqlite":
@@ -181,12 +177,10 @@ def _get_dialect_upsert_sql(self, dialect: str, session_id: str, data: str, expi
                 .columns(self._session_id_column, self._data_column, self._expires_at_column, self._created_at_column)
                 .values(session_id, data, expires_at, current_time)
                 .on_conflict(self._session_id_column)
-                .do_update(
-                    **{
-                        self._data_column: sql.raw("EXCLUDED." + self._data_column),
-                        self._expires_at_column: sql.raw("EXCLUDED." + self._expires_at_column),
-                    }
-                )
+                .do_update(**{
+                    self._data_column: sql.raw("EXCLUDED." + self._data_column),
+                    self._expires_at_column: sql.raw("EXCLUDED." + self._expires_at_column),
+                })
             )
 
         if dialect == "oracle":
@@ -363,7 +357,7 @@ async def _set_session_data(
         except Exception as e:
             msg = f"Failed to store session: {e}"
             logger.exception("Failed to store session %s", key)
-            raise SessionStoreError(msg) from e
+            raise SQLSpecSessionStoreError(msg) from e
 
     async def delete(self, key: str) -> None:
         """Delete session data.
@@ -392,7 +386,7 @@ async def _delete_session_data(
         except Exception as e:
             msg = f"Failed to delete session: {e}"
             logger.exception("Failed to delete session %s", key)
-            raise SessionStoreError(msg) from e
+            raise SQLSpecSessionStoreError(msg) from e
 
     async def exists(self, key: str) -> bool:
         """Check if a session exists and is not expired.
@@ -469,11 +463,9 @@ async def expires_in(self, key: str) -> int:
                 delta = expires_at - current_time
                 return max(0, int(delta.total_seconds()))
 
-            return 0
-
         except Exception:
             logger.exception("Failed to get expires_in for session %s", key)
-            return 0
+        return 0
 
     async def delete_all(self, pattern: str = "*") -> None:
         """Delete all sessions matching pattern.
@@ -499,7 +491,7 @@ async def _delete_all_sessions(self, driver: Union[SyncDriverAdapterBase, AsyncD
         except Exception as e:
             msg = f"Failed to delete all sessions: {e}"
             logger.exception("Failed to delete all sessions")
-            raise SessionStoreError(msg) from e
+            raise SQLSpecSessionStoreError(msg) from e
 
     async def delete_expired(self) -> None:
         """Delete expired sessions."""

tests/integration/test_adapters/test_adbc/test_extensions/test_litestar/__init__.py

@@ -0,0 +1,3 @@
+import pytest
+
+pytestmark = [pytest.mark.adbc, pytest.mark.postgres]