sqlspec 0.16.2__py3-none-any.whl → 0.17.1__py3-none-any.whl

sqlspec/adapters/sqlite/pool.py

@@ -0,0 +1,140 @@
+"""SQLite database configuration with thread-local connections."""
+
+import sqlite3
+import threading
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Optional, TypedDict, cast
+
+from typing_extensions import NotRequired
+
+from sqlspec.adapters.sqlite._types import SqliteConnection
+
+if TYPE_CHECKING:
+    from collections.abc import Generator
+
+
+class SqliteConnectionParams(TypedDict, total=False):
+    """SQLite connection parameters."""
+
+    database: NotRequired[str]
+    timeout: NotRequired[float]
+    detect_types: NotRequired[int]
+    isolation_level: "NotRequired[Optional[str]]"
+    check_same_thread: NotRequired[bool]
+    factory: "NotRequired[Optional[type[SqliteConnection]]]"
+    cached_statements: NotRequired[int]
+    uri: NotRequired[bool]
+
+
+__all__ = ("SqliteConnectionPool",)
+
+
+class SqliteConnectionPool:
+    """Thread-local connection manager for SQLite.
+
+    SQLite connections aren't thread-safe, so we use thread-local storage
+    to ensure each thread has its own connection. This is simpler and more
+    efficient than a traditional pool for SQLite's constraints.
+    """
+
+    __slots__ = ("_connection_parameters", "_enable_optimizations", "_thread_local")
+
+    def __init__(
+        self,
+        connection_parameters: "dict[str, Any]",
+        enable_optimizations: bool = True,
+        **kwargs: Any,  # Accept and ignore pool parameters for compatibility
+    ) -> None:
+        """Initialize the thread-local connection manager.
+
+        Args:
+            connection_parameters: SQLite connection parameters
+            enable_optimizations: Whether to apply performance PRAGMAs
+            **kwargs: Ignored pool parameters for compatibility
+        """
+        self._connection_parameters = connection_parameters
+        self._thread_local = threading.local()
+        self._enable_optimizations = enable_optimizations
+
+    def _create_connection(self) -> SqliteConnection:
+        """Create a new SQLite connection with optimizations."""
+        connection = sqlite3.connect(**self._connection_parameters)
+
+        # Only apply optimizations if requested and not in-memory
+        if self._enable_optimizations:
+            database = self._connection_parameters.get("database", ":memory:")
+            is_memory = database == ":memory:" or database.startswith("file::memory:")
+
+            if not is_memory:
+                # WAL mode doesn't work with in-memory databases
+                connection.execute("PRAGMA journal_mode = WAL")
+                # Set busy timeout for better concurrent access
+                connection.execute("PRAGMA busy_timeout = 5000")
+                connection.execute("PRAGMA optimize")
+            # These work for all database types
+            connection.execute("PRAGMA foreign_keys = ON")
+            connection.execute("PRAGMA synchronous = NORMAL")
+
+        return connection  # type: ignore[no-any-return]
+
+    def _get_thread_connection(self) -> SqliteConnection:
+        """Get or create a connection for the current thread."""
+        try:
+            return cast("SqliteConnection", self._thread_local.connection)
+        except AttributeError:
+            # Connection doesn't exist for this thread yet
+            connection = self._create_connection()
+            self._thread_local.connection = connection
+            return connection
+
+    def _close_thread_connection(self) -> None:
+        """Close the connection for the current thread."""
+        try:
+            connection = self._thread_local.connection
+            connection.close()
+            del self._thread_local.connection
+        except AttributeError:
+            # No connection for this thread
+            pass
+
+    @contextmanager
+    def get_connection(self) -> "Generator[SqliteConnection, None, None]":
+        """Get a thread-local connection.
+
+        Yields:
+            SqliteConnection: A thread-local connection.
+        """
+        yield self._get_thread_connection()
+
+    def close(self) -> None:
+        """Close the thread-local connection if it exists."""
+        self._close_thread_connection()
+
+    def acquire(self) -> SqliteConnection:
+        """Acquire a thread-local connection.
+
+        Returns:
+            SqliteConnection: A thread-local connection
+        """
+        return self._get_thread_connection()
+
+    def release(self, connection: SqliteConnection) -> None:
+        """Release a connection (no-op for thread-local connections).
+
+        Args:
+            connection: The connection to release (ignored)
+        """
+        # No-op: thread-local connections are managed per-thread
+
+    # Compatibility methods that return dummy values
+    def size(self) -> int:
+        """Get pool size (always 1 for thread-local)."""
+        try:
+            _ = self._thread_local.connection
+        except AttributeError:
+            return 0
+        return 1
+
+    def checked_out(self) -> int:
+        """Get number of checked out connections (always 0)."""
+        return 0

sqlspec/base.py

@@ -26,7 +26,10 @@ from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
     from contextlib import AbstractAsyncContextManager, AbstractContextManager
+    from pathlib import Path
 
+    from sqlspec.core.statement import SQL
+    from sqlspec.loader import SQLFileLoader
     from sqlspec.typing import ConnectionT, PoolT
 
 
@@ -38,54 +41,77 @@ logger = get_logger()
 class SQLSpec:
     """Configuration manager and registry for database connections and pools."""
 
-    __slots__ = ("_cleanup_tasks", "_configs", "_instance_cache_config")
+    __slots__ = ("_configs", "_instance_cache_config", "_sql_loader")
 
-    def __init__(self) -> None:
+    def __init__(self, *, loader: "Optional[SQLFileLoader]" = None) -> None:
         self._configs: dict[Any, DatabaseConfigProtocol[Any, Any, Any]] = {}
-        atexit.register(self._cleanup_pools)
+        # Register sync cleanup only for sync resources
+        atexit.register(self._cleanup_sync_pools)
         self._instance_cache_config: Optional[CacheConfig] = None
-        self._cleanup_tasks: list[asyncio.Task[None]] = []
+        self._sql_loader: Optional[SQLFileLoader] = loader
 
     @staticmethod
     def _get_config_name(obj: Any) -> str:
         """Get display name for configuration object."""
         return getattr(obj, "__name__", str(obj))
 
-    def _cleanup_pools(self) -> None:
-        """Clean up all registered connection pools."""
+    def _cleanup_sync_pools(self) -> None:
+        """Clean up only synchronous connection pools at exit."""
         cleaned_count = 0
 
+        for config_type, config in self._configs.items():
+            if config.supports_connection_pooling and not config.is_async:
+                try:
+                    config.close_pool()
+                    cleaned_count += 1
+                except Exception as e:
+                    logger.warning("Failed to clean up sync pool for config %s: %s", config_type.__name__, e)
+
+        if cleaned_count > 0:
+            logger.debug("Sync pool cleanup completed. Cleaned %d pools.", cleaned_count)
+
+    async def close_all_pools(self) -> None:
+        """Explicitly close all connection pools (async and sync).
+
+        This method should be called before application shutdown for proper cleanup.
+        """
+        cleanup_tasks = []
+        sync_configs = []
+
         for config_type, config in self._configs.items():
             if config.supports_connection_pooling:
                 try:
                     if config.is_async:
                         close_pool_awaitable = config.close_pool()
                         if close_pool_awaitable is not None:
-                            try:
-                                loop = asyncio.get_running_loop()
-                                if loop.is_running():
-                                    task = asyncio.create_task(cast("Coroutine[Any, Any, None]", close_pool_awaitable))
-                                    self._cleanup_tasks.append(task)
-                                else:
-                                    asyncio.run(cast("Coroutine[Any, Any, None]", close_pool_awaitable))
-                            except RuntimeError:
-                                asyncio.run(cast("Coroutine[Any, Any, None]", close_pool_awaitable))
+                            cleanup_tasks.append(cast("Coroutine[Any, Any, None]", close_pool_awaitable))
                     else:
-                        config.close_pool()
-                    cleaned_count += 1
+                        sync_configs.append((config_type, config))
                 except Exception as e:
-                    logger.warning("Failed to clean up pool for config %s: %s", config_type.__name__, e)
+                    logger.warning("Failed to prepare cleanup for config %s: %s", config_type.__name__, e)
 
-        if self._cleanup_tasks:
+        # Close async pools concurrently
+        if cleanup_tasks:
             try:
-                loop = asyncio.get_running_loop()
-                if loop.is_running():
-                    asyncio.gather(*self._cleanup_tasks, return_exceptions=True)
-            except RuntimeError:
-                pass
+                await asyncio.gather(*cleanup_tasks, return_exceptions=True)
+                logger.debug("Async pool cleanup completed. Cleaned %d pools.", len(cleanup_tasks))
+            except Exception as e:
+                logger.warning("Failed to complete async pool cleanup: %s", e)
+
+        # Close sync pools
+        for _config_type, config in sync_configs:
+            config.close_pool()  # Let exceptions propagate for proper logging
+
+        if sync_configs:
+            logger.debug("Sync pool cleanup completed. Cleaned %d pools.", len(sync_configs))
 
-        self._configs.clear()
-        logger.info("Pool cleanup completed. Cleaned %d pools.", cleaned_count)
+    async def __aenter__(self) -> "SQLSpec":
+        """Async context manager entry."""
+        return self
+
+    async def __aexit__(self, _exc_type: Any, _exc_val: Any, _exc_tb: Any) -> None:
+        """Async context manager exit with automatic cleanup."""
+        await self.close_all_pools()
 
     @overload
     def add_config(self, config: "SyncConfigT") -> "type[SyncConfigT]":  # pyright: ignore[reportInvalidTypeVarUse]
@@ -569,3 +595,98 @@ class SQLSpec:
                 else current_config.optimized_cache_enabled,
             )
         )
+
+    # SQL File Loading Integration
+
+    def _ensure_sql_loader(self) -> "SQLFileLoader":
+        """Ensure SQL loader is initialized lazily."""
+        if self._sql_loader is None:
+            # Import here to avoid circular imports
+            from sqlspec.loader import SQLFileLoader
+
+            self._sql_loader = SQLFileLoader()
+        return self._sql_loader
+
+    def load_sql_files(self, *paths: "Union[str, Path]") -> None:
+        """Load SQL files from paths or directories.
+
+        Args:
+            *paths: One or more file paths or directory paths to load.
+        """
+        loader = self._ensure_sql_loader()
+        loader.load_sql(*paths)
+        logger.debug("Loaded SQL files: %s", paths)
+
+    def add_named_sql(self, name: str, sql: str, dialect: "Optional[str]" = None) -> None:
+        """Add a named SQL query directly.
+
+        Args:
+            name: Name for the SQL query.
+            sql: Raw SQL content.
+            dialect: Optional dialect for the SQL statement.
+        """
+        loader = self._ensure_sql_loader()
+        loader.add_named_sql(name, sql, dialect)
+        logger.debug("Added named SQL: %s", name)
+
+    def get_sql(self, name: str) -> "SQL":
+        """Get a SQL object by name.
+
+        Args:
+            name: Name of the statement (from -- name: in SQL file).
+                  Hyphens in names are converted to underscores.
+
+        Returns:
+            SQL object ready for execution.
+        """
+        loader = self._ensure_sql_loader()
+        return loader.get_sql(name)
+
+    def list_sql_queries(self) -> "list[str]":
+        """List all available query names.
+
+        Returns:
+            Sorted list of query names.
+        """
+        if self._sql_loader is None:
+            return []
+        return self._sql_loader.list_queries()
+
+    def has_sql_query(self, name: str) -> bool:
+        """Check if a SQL query exists.
+
+        Args:
+            name: Query name to check.
+
+        Returns:
+            True if query exists.
+        """
+        if self._sql_loader is None:
+            return False
+        return self._sql_loader.has_query(name)
+
+    def clear_sql_cache(self) -> None:
+        """Clear the SQL file cache."""
+        if self._sql_loader is not None:
+            self._sql_loader.clear_cache()
+            logger.debug("Cleared SQL cache")
+
+    def reload_sql_files(self) -> None:
+        """Reload all SQL files.
+
+        Note: This clears the cache and requires calling load_sql_files again.
+        """
+        if self._sql_loader is not None:
+            # Clear cache to force reload
+            self._sql_loader.clear_cache()
+            logger.debug("Cleared SQL cache for reload")
+
+    def get_sql_files(self) -> "list[str]":
+        """Get list of loaded SQL files.
+
+        Returns:
+            Sorted list of file paths.
+        """
+        if self._sql_loader is None:
+            return []
+        return self._sql_loader.list_files()

sqlspec/builder/__init__.py

@@ -29,10 +29,13 @@ from sqlspec.builder._merge import Merge
 from sqlspec.builder._select import Select
 from sqlspec.builder._update import Update
 from sqlspec.builder.mixins import WhereClauseMixin
+from sqlspec.builder.mixins._join_operations import JoinBuilder
+from sqlspec.builder.mixins._select_operations import Case, SubqueryBuilder, WindowFunctionBuilder
 from sqlspec.exceptions import SQLBuilderError
 
 __all__ = (
     "AlterTable",
+    "Case",
     "Column",
     "ColumnExpression",
     "CommentOn",
@@ -50,13 +53,16 @@ __all__ = (
     "DropView",
     "FunctionColumn",
     "Insert",
+    "JoinBuilder",
     "Merge",
     "QueryBuilder",
     "RenameTable",
     "SQLBuilderError",
     "SafeQuery",
     "Select",
+    "SubqueryBuilder",
     "Truncate",
     "Update",
     "WhereClauseMixin",
+    "WindowFunctionBuilder",
 )

sqlspec/builder/_column.py

@@ -5,7 +5,7 @@ for building SQL conditions with type safety and parameter binding.
 """
 
 from collections.abc import Iterable
-from typing import Any, Optional
+from typing import Any, Optional, cast
 
 from sqlglot import exp
 
@@ -241,6 +241,10 @@ class Column:
         """Create a DESC ordering expression."""
         return exp.Ordered(this=self._expression, desc=True)
 
+    def as_(self, alias: str) -> exp.Alias:
+        """Create an aliased expression."""
+        return cast("exp.Alias", exp.alias_(self._expression, alias))
+
     def __repr__(self) -> str:
         if self.table:
             return f"Column<{self.table}.{self.name}>"

sqlspec/builder/_expression_wrappers.py

@@ -0,0 +1,46 @@
+"""Expression wrapper classes for proper type annotations."""
+
+from typing import cast
+
+from sqlglot import exp
+
+__all__ = ("AggregateExpression", "ConversionExpression", "FunctionExpression", "MathExpression", "StringExpression")
+
+
+class ExpressionWrapper:
+    """Base wrapper for SQLGlot expressions."""
+
+    def __init__(self, expression: exp.Expression) -> None:
+        self._expression = expression
+
+    def as_(self, alias: str) -> exp.Alias:
+        """Create an aliased expression."""
+        return cast("exp.Alias", exp.alias_(self._expression, alias))
+
+    @property
+    def expression(self) -> exp.Expression:
+        """Get the underlying SQLGlot expression."""
+        return self._expression
+
+    def __str__(self) -> str:
+        return str(self._expression)
+
+
+class AggregateExpression(ExpressionWrapper):
+    """Aggregate functions like COUNT, SUM, AVG."""
+
+
+class FunctionExpression(ExpressionWrapper):
+    """General SQL functions."""
+
+
+class MathExpression(ExpressionWrapper):
+    """Mathematical functions like ROUND."""
+
+
+class StringExpression(ExpressionWrapper):
+    """String functions like UPPER, LOWER, LENGTH."""
+
+
+class ConversionExpression(ExpressionWrapper):
+    """Conversion functions like CAST, COALESCE."""

sqlspec/builder/_insert.py

@@ -412,9 +412,7 @@ class ConflictBuilder:
         # Create ON CONFLICT with proper structure
         conflict_keys = [exp.to_identifier(col) for col in self._columns] if self._columns else None
         on_conflict = exp.OnConflict(
-            conflict_keys=conflict_keys,
-            action=exp.var("DO UPDATE"),
-            expressions=set_expressions if set_expressions else None,
+            conflict_keys=conflict_keys, action=exp.var("DO UPDATE"), expressions=set_expressions or None
         )
 
         insert_expr.set("conflict", on_conflict)

sqlspec/builder/_parsing_utils.py

@@ -9,6 +9,7 @@ from typing import Any, Final, Optional, Union, cast
 
 from sqlglot import exp, maybe_parse, parse_one
 
+from sqlspec.core.parameters import ParameterStyle
 from sqlspec.utils.type_guards import has_expression_attr, has_parameter_builder
 
 
@@ -151,6 +152,32 @@ def parse_condition_expression(
     if not isinstance(condition_input, str):
         condition_input = str(condition_input)
 
+    # Convert database-specific parameter styles to SQLGlot-compatible format
+    # This ensures that placeholders like $1, %s, :1 are properly recognized as parameters
+    from sqlspec.core.parameters import ParameterValidator
+
+    validator = ParameterValidator()
+    param_info = validator.extract_parameters(condition_input)
+
+    # If we found parameters, convert incompatible ones to SQLGlot-compatible format
+    if param_info:
+        # Convert problematic parameter styles to :param_N format for SQLGlot
+        converted_condition = condition_input
+        for param in reversed(param_info):  # Reverse to preserve positions
+            if param.style in {
+                ParameterStyle.NUMERIC,
+                ParameterStyle.POSITIONAL_PYFORMAT,
+                ParameterStyle.POSITIONAL_COLON,
+            }:
+                # Convert $1, %s, :1 to :param_0, :param_1, etc.
+                placeholder = f":param_{param.ordinal}"
+                converted_condition = (
+                    converted_condition[: param.position]
+                    + placeholder
+                    + converted_condition[param.position + len(param.placeholder_text) :]
+                )
+        condition_input = converted_condition
+
     try:
         return exp.condition(condition_input)
     except Exception:

sqlspec/builder/_update.py

@@ -44,26 +44,26 @@ class Update(
         update_query = (
             Update()
             .table("users")
-            .set(name="John Doe")
-            .set(email="john@example.com")
+            .set_(name="John Doe")
+            .set_(email="john@example.com")
             .where("id = 1")
         )
 
         update_query = (
-            Update("users").set(name="John Doe").where("id = 1")
+            Update("users").set_(name="John Doe").where("id = 1")
         )
 
         update_query = (
             Update()
             .table("users")
-            .set(status="active")
+            .set_(status="active")
             .where_eq("id", 123)
         )
 
         update_query = (
             Update()
             .table("users", "u")
-            .set(name="Updated Name")
+            .set_(name="Updated Name")
             .from_("profiles", "p")
             .where("u.id = p.user_id AND p.is_verified = true")
         )

sqlspec/builder/mixins/_join_operations.py

@@ -9,10 +9,11 @@ from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.type_guards import has_query_builder_parameters
 
 if TYPE_CHECKING:
+    from sqlspec.builder._column import ColumnExpression
     from sqlspec.core.statement import SQL
     from sqlspec.protocols import SQLBuilderProtocol
 
-__all__ = ("JoinClauseMixin",)
+__all__ = ("JoinBuilder", "JoinClauseMixin")
 
 
 @trait
@@ -147,3 +148,116 @@ class JoinClauseMixin:
         join_expr = exp.Join(this=table_expr, kind="CROSS")
         builder._expression = builder._expression.join(join_expr, copy=False)
         return cast("Self", builder)
+
+
+@trait
+class JoinBuilder:
+    """Builder for JOIN operations with fluent syntax.
+
+    Example:
+        ```python
+        from sqlspec import sql
+
+        # sql.left_join_("posts").on("users.id = posts.user_id")
+        join_clause = sql.left_join_("posts").on(
+            "users.id = posts.user_id"
+        )
+
+        # Or with query builder
+        query = (
+            sql.select("users.name", "posts.title")
+            .from_("users")
+            .join(
+                sql.left_join_("posts").on(
+                    "users.id = posts.user_id"
+                )
+            )
+        )
+        ```
+    """
+
+    def __init__(self, join_type: str) -> None:
+        """Initialize the join builder.
+
+        Args:
+            join_type: Type of join (inner, left, right, full, cross)
+        """
+        self._join_type = join_type.upper()
+        self._table: Optional[Union[str, exp.Expression]] = None
+        self._condition: Optional[exp.Expression] = None
+        self._alias: Optional[str] = None
+
+    def __eq__(self, other: object) -> "ColumnExpression":  # type: ignore[override]
+        """Equal to (==) - not typically used but needed for type consistency."""
+        from sqlspec.builder._column import ColumnExpression
+
+        # JoinBuilder doesn't have a direct expression, so this is a placeholder
+        # In practice, this shouldn't be called as joins are used differently
+        placeholder_expr = exp.Literal.string(f"join_{self._join_type.lower()}")
+        if other is None:
+            return ColumnExpression(exp.Is(this=placeholder_expr, expression=exp.Null()))
+        return ColumnExpression(exp.EQ(this=placeholder_expr, expression=exp.convert(other)))
+
+    def __hash__(self) -> int:
+        """Make JoinBuilder hashable."""
+        return hash(id(self))
+
+    def __call__(self, table: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Set the table to join.
+
+        Args:
+            table: Table name or expression to join
+            alias: Optional alias for the table
+
+        Returns:
+            Self for method chaining
+        """
+        self._table = table
+        self._alias = alias
+        return self
+
+    def on(self, condition: Union[str, exp.Expression]) -> exp.Expression:
+        """Set the join condition and build the JOIN expression.
+
+        Args:
+            condition: JOIN condition (e.g., "users.id = posts.user_id")
+
+        Returns:
+            Complete JOIN expression
+        """
+        if not self._table:
+            msg = "Table must be set before calling .on()"
+            raise SQLBuilderError(msg)
+
+        # Parse the condition
+        condition_expr: exp.Expression
+        if isinstance(condition, str):
+            parsed: Optional[exp.Expression] = exp.maybe_parse(condition)
+            condition_expr = parsed or exp.condition(condition)
+        else:
+            condition_expr = condition
+
+        # Build table expression
+        table_expr: exp.Expression
+        if isinstance(self._table, str):
+            table_expr = exp.to_table(self._table)
+            if self._alias:
+                table_expr = exp.alias_(table_expr, self._alias)
+        else:
+            table_expr = self._table
+            if self._alias:
+                table_expr = exp.alias_(table_expr, self._alias)
+
+        # Create the appropriate join type using same pattern as existing JoinClauseMixin
+        if self._join_type == "INNER JOIN":
+            return exp.Join(this=table_expr, on=condition_expr)
+        if self._join_type == "LEFT JOIN":
+            return exp.Join(this=table_expr, on=condition_expr, side="LEFT")
+        if self._join_type == "RIGHT JOIN":
+            return exp.Join(this=table_expr, on=condition_expr, side="RIGHT")
+        if self._join_type == "FULL JOIN":
+            return exp.Join(this=table_expr, on=condition_expr, side="FULL", kind="OUTER")
+        if self._join_type == "CROSS JOIN":
+            # CROSS JOIN doesn't use ON condition
+            return exp.Join(this=table_expr, kind="CROSS")
+        return exp.Join(this=table_expr, on=condition_expr)