sqlspec 0.26.0__py3-none-any.whl → 0.28.0__py3-none-any.whl

sqlspec/extensions/starlette/extension.py

@@ -0,0 +1,254 @@
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any
+
+from sqlspec.base import SQLSpec
+from sqlspec.exceptions import ImproperConfigurationError
+from sqlspec.extensions.starlette._state import SQLSpecConfigState
+from sqlspec.extensions.starlette._utils import get_or_create_session
+from sqlspec.extensions.starlette.middleware import SQLSpecAutocommitMiddleware, SQLSpecManualMiddleware
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from collections.abc import AsyncGenerator
+
+    from starlette.applications import Starlette
+    from starlette.requests import Request
+
+__all__ = ("SQLSpecPlugin",)
+
+logger = get_logger("extensions.starlette")
+
+DEFAULT_COMMIT_MODE = "manual"
+DEFAULT_CONNECTION_KEY = "db_connection"
+DEFAULT_POOL_KEY = "db_pool"
+DEFAULT_SESSION_KEY = "db_session"
+
+
+class SQLSpecPlugin:
+    """SQLSpec integration for Starlette applications.
+
+    Provides middleware-based session management, automatic transaction handling,
+    and connection pooling lifecycle management.
+
+    Example:
+        from starlette.applications import Starlette
+        from sqlspec import SQLSpec
+        from sqlspec.adapters.asyncpg import AsyncpgConfig
+        from sqlspec.extensions.starlette import SQLSpecPlugin
+
+        sqlspec = SQLSpec()
+        sqlspec.add_config(AsyncpgConfig(
+            bind_key="default",
+            pool_config={"dsn": "postgresql://localhost/mydb"},
+            extension_config={
+                "starlette": {
+                    "commit_mode": "autocommit",
+                    "session_key": "db"
+                }
+            }
+        ))
+
+        app = Starlette()
+        db_ext = SQLSpecPlugin(sqlspec, app)
+
+        @app.route("/users")
+        async def list_users(request):
+            db = db_ext.get_session(request)
+            result = await db.execute("SELECT * FROM users")
+            return JSONResponse({"users": result.all()})
+    """
+
+    __slots__ = ("_config_states", "_sqlspec")
+
+    def __init__(self, sqlspec: SQLSpec, app: "Starlette | None" = None) -> None:
+        """Initialize SQLSpec Starlette extension.
+
+        Args:
+            sqlspec: Pre-configured SQLSpec instance with registered configs.
+            app: Optional Starlette application to initialize immediately.
+        """
+        self._sqlspec = sqlspec
+        self._config_states: list[SQLSpecConfigState] = []
+
+        for cfg in self._sqlspec.configs.values():
+            settings = self._extract_starlette_settings(cfg)
+            state = self._create_config_state(cfg, settings)
+            self._config_states.append(state)
+
+        if app is not None:
+            self.init_app(app)
+
+    def _extract_starlette_settings(self, config: Any) -> "dict[str, Any]":
+        """Extract Starlette settings from config.extension_config.
+
+        Args:
+            config: Database configuration instance.
+
+        Returns:
+            Dictionary of Starlette-specific settings.
+        """
+        starlette_config = config.extension_config.get("starlette", {})
+
+        connection_key = starlette_config.get("connection_key", DEFAULT_CONNECTION_KEY)
+        pool_key = starlette_config.get("pool_key", DEFAULT_POOL_KEY)
+        session_key = starlette_config.get("session_key", DEFAULT_SESSION_KEY)
+        commit_mode = starlette_config.get("commit_mode", DEFAULT_COMMIT_MODE)
+
+        if not config.supports_connection_pooling and pool_key == DEFAULT_POOL_KEY:
+            pool_key = f"_{DEFAULT_POOL_KEY}_{id(config)}"
+
+        return {
+            "connection_key": connection_key,
+            "pool_key": pool_key,
+            "session_key": session_key,
+            "commit_mode": commit_mode,
+            "extra_commit_statuses": starlette_config.get("extra_commit_statuses"),
+            "extra_rollback_statuses": starlette_config.get("extra_rollback_statuses"),
+        }
+
+    def _create_config_state(self, config: Any, settings: "dict[str, Any]") -> SQLSpecConfigState:
+        """Create configuration state object.
+
+        Args:
+            config: Database configuration instance.
+            settings: Extracted Starlette settings.
+
+        Returns:
+            Configuration state instance.
+        """
+        return SQLSpecConfigState(
+            config=config,
+            connection_key=settings["connection_key"],
+            pool_key=settings["pool_key"],
+            session_key=settings["session_key"],
+            commit_mode=settings["commit_mode"],
+            extra_commit_statuses=settings["extra_commit_statuses"],
+            extra_rollback_statuses=settings["extra_rollback_statuses"],
+        )
+
+    def init_app(self, app: "Starlette") -> None:
+        """Initialize Starlette application with SQLSpec.
+
+        Validates configuration, wraps lifespan, and adds middleware.
+
+        Args:
+            app: Starlette application instance.
+        """
+        self._validate_unique_keys()
+
+        original_lifespan = app.router.lifespan_context
+
+        @asynccontextmanager
+        async def combined_lifespan(app: "Starlette") -> "AsyncGenerator[None, None]":
+            async with self.lifespan(app), original_lifespan(app):
+                yield
+
+        app.router.lifespan_context = combined_lifespan
+
+        for config_state in self._config_states:
+            self._add_middleware(app, config_state)
+
+    def _validate_unique_keys(self) -> None:
+        """Validate that all state keys are unique across configs.
+
+        Raises:
+            ImproperConfigurationError: If duplicate keys found.
+        """
+        all_keys: set[str] = set()
+
+        for state in self._config_states:
+            keys = {state.connection_key, state.pool_key, state.session_key}
+            duplicates = all_keys & keys
+
+            if duplicates:
+                msg = f"Duplicate state keys found: {duplicates}"
+                raise ImproperConfigurationError(msg)
+
+            all_keys.update(keys)
+
+    def _add_middleware(self, app: "Starlette", config_state: SQLSpecConfigState) -> None:
+        """Add transaction middleware for configuration.
+
+        Args:
+            app: Starlette application instance.
+            config_state: Configuration state.
+        """
+        if config_state.commit_mode == "manual":
+            app.add_middleware(SQLSpecManualMiddleware, config_state=config_state)
+        elif config_state.commit_mode == "autocommit":
+            app.add_middleware(SQLSpecAutocommitMiddleware, config_state=config_state, include_redirect=False)
+        elif config_state.commit_mode == "autocommit_include_redirect":
+            app.add_middleware(SQLSpecAutocommitMiddleware, config_state=config_state, include_redirect=True)
+
+    @asynccontextmanager
+    async def lifespan(self, app: "Starlette") -> "AsyncGenerator[None, None]":
+        """Manage connection pool lifecycle.
+
+        Args:
+            app: Starlette application instance.
+
+        Yields:
+            None
+        """
+        for config_state in self._config_states:
+            if config_state.config.supports_connection_pooling:
+                pool = await config_state.config.create_pool()
+                setattr(app.state, config_state.pool_key, pool)
+
+        try:
+            yield
+        finally:
+            for config_state in self._config_states:
+                if config_state.config.supports_connection_pooling:
+                    close_result = config_state.config.close_pool()
+                    if close_result is not None:
+                        await close_result
+
+    def get_session(self, request: "Request", key: "str | None" = None) -> Any:
+        """Get or create database session for request.
+
+        Sessions are cached per request to ensure consistency.
+
+        Args:
+            request: Starlette request instance.
+            key: Optional session key to retrieve specific database session.
+
+        Returns:
+            Database session (driver instance).
+        """
+        config_state = self._config_states[0] if key is None else self._get_config_state_by_key(key)
+
+        return get_or_create_session(request, config_state)
+
+    def get_connection(self, request: "Request", key: "str | None" = None) -> Any:
+        """Get database connection from request state.
+
+        Args:
+            request: Starlette request instance.
+            key: Optional session key to retrieve specific database connection.
+
+        Returns:
+            Database connection object.
+        """
+        config_state = self._config_states[0] if key is None else self._get_config_state_by_key(key)
+
+        return getattr(request.state, config_state.connection_key)
+
+    def _get_config_state_by_key(self, key: str) -> SQLSpecConfigState:
+        """Get configuration state by session key.
+
+        Args:
+            key: Session key to search for.
+
+        Returns:
+            Configuration state matching the key.
+
+        Raises:
+            ValueError: If no configuration found with the specified key.
+        """
+        for state in self._config_states:
+            if state.session_key == key:
+                return state
+
+        msg = f"No configuration found with session_key: {key}"
+        raise ValueError(msg)

sqlspec/extensions/starlette/middleware.py

@@ -0,0 +1,154 @@
+from typing import TYPE_CHECKING, Any
+
+from starlette.middleware.base import BaseHTTPMiddleware
+
+from sqlspec.utils.logging import get_logger
+
+if TYPE_CHECKING:
+    from starlette.requests import Request
+
+    from sqlspec.extensions.starlette._state import SQLSpecConfigState
+
+__all__ = ("SQLSpecAutocommitMiddleware", "SQLSpecManualMiddleware")
+
+logger = get_logger("extensions.starlette.middleware")
+
+HTTP_200_OK = 200
+HTTP_300_MULTIPLE_CHOICES = 300
+HTTP_400_BAD_REQUEST = 400
+
+
+class SQLSpecManualMiddleware(BaseHTTPMiddleware):
+    """Middleware for manual transaction mode.
+
+    Acquires connection from pool, stores in request.state, releases after request.
+    No automatic commit or rollback - user code must handle transactions.
+    """
+
+    def __init__(self, app: Any, config_state: "SQLSpecConfigState") -> None:
+        """Initialize middleware.
+
+        Args:
+            app: Starlette application instance.
+            config_state: Configuration state for this database.
+        """
+        super().__init__(app)
+        self.config_state = config_state
+
+    async def dispatch(self, request: "Request", call_next: Any) -> Any:
+        """Process request with manual transaction mode.
+
+        Args:
+            request: Incoming HTTP request.
+            call_next: Next middleware or route handler.
+
+        Returns:
+            HTTP response.
+        """
+        config = self.config_state.config
+        connection_key = self.config_state.connection_key
+
+        if config.supports_connection_pooling:
+            pool = getattr(request.app.state, self.config_state.pool_key)
+            async with config.provide_connection(pool) as connection:  # type: ignore[union-attr]
+                setattr(request.state, connection_key, connection)
+                try:
+                    return await call_next(request)
+                finally:
+                    delattr(request.state, connection_key)
+        else:
+            connection = await config.create_connection()
+            setattr(request.state, connection_key, connection)
+            try:
+                return await call_next(request)
+            finally:
+                await connection.close()
+
+
+class SQLSpecAutocommitMiddleware(BaseHTTPMiddleware):
+    """Middleware for autocommit transaction mode.
+
+    Acquires connection, commits on success status codes, rollbacks on error status codes.
+    """
+
+    def __init__(self, app: Any, config_state: "SQLSpecConfigState", include_redirect: bool = False) -> None:
+        """Initialize middleware.
+
+        Args:
+            app: Starlette application instance.
+            config_state: Configuration state for this database.
+            include_redirect: If True, commit on 3xx status codes as well.
+        """
+        super().__init__(app)
+        self.config_state = config_state
+        self.include_redirect = include_redirect
+
+    async def dispatch(self, request: "Request", call_next: Any) -> Any:
+        """Process request with autocommit transaction mode.
+
+        Args:
+            request: Incoming HTTP request.
+            call_next: Next middleware or route handler.
+
+        Returns:
+            HTTP response.
+        """
+        config = self.config_state.config
+        connection_key = self.config_state.connection_key
+
+        if config.supports_connection_pooling:
+            pool = getattr(request.app.state, self.config_state.pool_key)
+            async with config.provide_connection(pool) as connection:  # type: ignore[union-attr]
+                setattr(request.state, connection_key, connection)
+                try:
+                    response = await call_next(request)
+
+                    if self._should_commit(response.status_code):
+                        await connection.commit()
+                    else:
+                        await connection.rollback()
+                except Exception:
+                    await connection.rollback()
+                    raise
+                else:
+                    return response
+                finally:
+                    delattr(request.state, connection_key)
+        else:
+            connection = await config.create_connection()
+            setattr(request.state, connection_key, connection)
+            try:
+                response = await call_next(request)
+
+                if self._should_commit(response.status_code):
+                    await connection.commit()
+                else:
+                    await connection.rollback()
+            except Exception:
+                await connection.rollback()
+                raise
+            else:
+                return response
+            finally:
+                await connection.close()
+
+    def _should_commit(self, status_code: int) -> bool:
+        """Determine if response status code should trigger commit.
+
+        Args:
+            status_code: HTTP status code.
+
+        Returns:
+            True if should commit, False if should rollback.
+        """
+        extra_commit = self.config_state.extra_commit_statuses or set()
+        extra_rollback = self.config_state.extra_rollback_statuses or set()
+
+        if status_code in extra_commit:
+            return True
+        if status_code in extra_rollback:
+            return False
+
+        if HTTP_200_OK <= status_code < HTTP_300_MULTIPLE_CHOICES:
+            return True
+        return bool(self.include_redirect and HTTP_300_MULTIPLE_CHOICES <= status_code < HTTP_400_BAD_REQUEST)

sqlspec/loader.py

@@ -9,7 +9,7 @@ import re
 import time
 from datetime import datetime, timezone
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Final, Optional, Union
+from typing import TYPE_CHECKING, Any, Final
 from urllib.parse import unquote, urlparse
 
 from sqlspec.core.cache import get_cache, get_cache_config
@@ -95,7 +95,7 @@ class NamedStatement:
 
     __slots__ = ("dialect", "name", "sql", "start_line")
 
-    def __init__(self, name: str, sql: str, dialect: "Optional[str]" = None, start_line: int = 0) -> None:
+    def __init__(self, name: str, sql: str, dialect: "str | None" = None, start_line: int = 0) -> None:
         self.name = name
         self.sql = sql
         self.dialect = dialect
@@ -112,11 +112,7 @@ class SQLFile:
     __slots__ = ("checksum", "content", "loaded_at", "metadata", "path")
 
     def __init__(
-        self,
-        content: str,
-        path: str,
-        metadata: "Optional[dict[str, Any]]" = None,
-        loaded_at: "Optional[datetime]" = None,
+        self, content: str, path: str, metadata: "dict[str, Any] | None" = None, loaded_at: "datetime | None" = None
     ) -> None:
         """Initialize SQLFile.
 
@@ -163,7 +159,7 @@ class SQLFileLoader:
 
     __slots__ = ("_files", "_queries", "_query_to_file", "encoding", "storage_registry")
 
-    def __init__(self, *, encoding: str = "utf-8", storage_registry: "Optional[StorageRegistry]" = None) -> None:
+    def __init__(self, *, encoding: str = "utf-8", storage_registry: "StorageRegistry | None" = None) -> None:
         """Initialize the SQL file loader.
 
         Args:
@@ -188,7 +184,7 @@ class SQLFileLoader:
         """
         raise SQLFileNotFoundError(path)
 
-    def _generate_file_cache_key(self, path: Union[str, Path]) -> str:
+    def _generate_file_cache_key(self, path: str | Path) -> str:
         """Generate cache key for a file path.
 
         Args:
@@ -201,7 +197,7 @@ class SQLFileLoader:
         path_hash = hashlib.md5(path_str.encode(), usedforsecurity=False).hexdigest()
         return f"file:{path_hash[:16]}"
 
-    def _calculate_file_checksum(self, path: Union[str, Path]) -> str:
+    def _calculate_file_checksum(self, path: str | Path) -> str:
         """Calculate checksum for file content validation.
 
         Args:
@@ -218,7 +214,7 @@ class SQLFileLoader:
         except Exception as e:
             raise SQLFileParseError(str(path), str(path), e) from e
 
-    def _is_file_unchanged(self, path: Union[str, Path], cached_file: CachedSQLFile) -> bool:
+    def _is_file_unchanged(self, path: str | Path, cached_file: CachedSQLFile) -> bool:
         """Check if file has changed since caching.
 
         Args:
@@ -235,7 +231,7 @@ class SQLFileLoader:
         else:
             return current_checksum == cached_file.sql_file.checksum
 
-    def _read_file_content(self, path: Union[str, Path]) -> str:
+    def _read_file_content(self, path: str | Path) -> str:
         """Read file content using storage backend.
 
         Args:
@@ -349,7 +345,7 @@ class SQLFileLoader:
 
         return statements
 
-    def load_sql(self, *paths: Union[str, Path]) -> None:
+    def load_sql(self, *paths: str | Path) -> None:
         """Load SQL files and parse named queries.
 
         Args:
@@ -358,43 +354,20 @@ class SQLFileLoader:
         correlation_id = CorrelationContext.get()
         start_time = time.perf_counter()
 
-        logger.info("Loading SQL files", extra={"file_count": len(paths), "correlation_id": correlation_id})
-
-        loaded_count = 0
-        query_count_before = len(self._queries)
-
         try:
             for path in paths:
                 path_str = str(path)
                 if "://" in path_str:
                     self._load_single_file(path, None)
-                    loaded_count += 1
                 else:
                     path_obj = Path(path)
                     if path_obj.is_dir():
-                        loaded_count += self._load_directory(path_obj)
+                        self._load_directory(path_obj)
                     elif path_obj.exists():
                         self._load_single_file(path_obj, None)
-                        loaded_count += 1
                     elif path_obj.suffix:
                         self._raise_file_not_found(str(path))
 
-            duration = time.perf_counter() - start_time
-            new_queries = len(self._queries) - query_count_before
-
-            logger.info(
-                "Loaded %d SQL files with %d new queries in %.3fms",
-                loaded_count,
-                new_queries,
-                duration * 1000,
-                extra={
-                    "files_loaded": loaded_count,
-                    "new_queries": new_queries,
-                    "duration_ms": duration * 1000,
-                    "correlation_id": correlation_id,
-                },
-            )
-
         except Exception as e:
             duration = time.perf_counter() - start_time
             logger.exception(
@@ -408,34 +381,40 @@ class SQLFileLoader:
             )
             raise
 
-    def _load_directory(self, dir_path: Path) -> int:
-        """Load all SQL files from a directory."""
+    def _load_directory(self, dir_path: Path) -> None:
+        """Load all SQL files from a directory.
+
+        Args:
+            dir_path: Directory path to load SQL files from.
+        """
         sql_files = list(dir_path.rglob("*.sql"))
         if not sql_files:
-            return 0
+            return
 
         for file_path in sql_files:
             relative_path = file_path.relative_to(dir_path)
             namespace_parts = relative_path.parent.parts
             self._load_single_file(file_path, ".".join(namespace_parts) if namespace_parts else None)
-        return len(sql_files)
 
-    def _load_single_file(self, file_path: Union[str, Path], namespace: Optional[str]) -> None:
+    def _load_single_file(self, file_path: str | Path, namespace: str | None) -> bool:
         """Load a single SQL file with optional namespace.
 
         Args:
             file_path: Path to the SQL file.
             namespace: Optional namespace prefix for queries.
+
+        Returns:
+            True if file was newly loaded, False if already cached.
         """
         path_str = str(file_path)
 
         if path_str in self._files:
-            return
+            return False
 
         cache_config = get_cache_config()
         if not cache_config.compiled_cache_enabled:
             self._load_file_without_cache(file_path, namespace)
-            return
+            return True
 
         cache_key_str = self._generate_file_cache_key(file_path)
         cache = get_cache()
@@ -459,7 +438,7 @@ class SQLFileLoader:
                         )
                 self._queries[namespaced_name] = statement
                 self._query_to_file[namespaced_name] = path_str
-            return
+            return True
 
         self._load_file_without_cache(file_path, namespace)
 
@@ -476,7 +455,9 @@ class SQLFileLoader:
             cached_file_data = CachedSQLFile(sql_file=sql_file, parsed_statements=file_statements)
             cache.put("file", cache_key_str, cached_file_data)
 
-    def _load_file_without_cache(self, file_path: Union[str, Path], namespace: Optional[str]) -> None:
+        return True
+
+    def _load_file_without_cache(self, file_path: str | Path, namespace: str | None) -> None:
         """Load a single SQL file without using cache.
 
         Args:
@@ -503,7 +484,7 @@ class SQLFileLoader:
             self._queries[namespaced_name] = statement
             self._query_to_file[namespaced_name] = path_str
 
-    def add_named_sql(self, name: str, sql: str, dialect: "Optional[str]" = None) -> None:
+    def add_named_sql(self, name: str, sql: str, dialect: "str | None" = None) -> None:
         """Add a named SQL query directly without loading from a file.
 
         Args:
@@ -529,7 +510,7 @@ class SQLFileLoader:
         self._queries[normalized_name] = statement
         self._query_to_file[normalized_name] = "<directly added>"
 
-    def get_file(self, path: Union[str, Path]) -> "Optional[SQLFile]":
+    def get_file(self, path: str | Path) -> "SQLFile | None":
         """Get a loaded SQLFile object by path.
 
         Args:
@@ -540,7 +521,7 @@ class SQLFileLoader:
         """
         return self._files.get(str(path))
 
-    def get_file_for_query(self, name: str) -> "Optional[SQLFile]":
+    def get_file_for_query(self, name: str) -> "SQLFile | None":
         """Get the SQLFile object containing a query.
 
         Args: