sqlspec 0.16.1__py3-none-any.whl → 0.17.0__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/_insert.py

@@ -142,6 +142,29 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
         for i, value in enumerate(values):
             if isinstance(value, exp.Expression):
                 value_placeholders.append(value)
+            elif hasattr(value, "expression") and hasattr(value, "sql"):
+                # Handle SQL objects (from sql.raw with parameters)
+                expression = getattr(value, "expression", None)
+                if expression is not None and isinstance(expression, exp.Expression):
+                    # Merge parameters from SQL object into builder
+                    if hasattr(value, "parameters"):
+                        sql_parameters = getattr(value, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            self.add_parameter(param_value, name=param_name)
+                    value_placeholders.append(expression)
+                else:
+                    # If expression is None, fall back to parsing the raw SQL
+                    sql_text = getattr(value, "sql", "")
+                    # Merge parameters even when parsing raw SQL
+                    if hasattr(value, "parameters"):
+                        sql_parameters = getattr(value, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            self.add_parameter(param_value, name=param_name)
+                    # Check if sql_text is callable (like Expression.sql method)
+                    if callable(sql_text):
+                        sql_text = str(value)
+                    value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
+                    value_placeholders.append(value_expr)
             else:
                 if self._columns and i < len(self._columns):
                     column_str = str(self._columns[i])
@@ -228,29 +251,171 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
 
         return self
 
-    def on_conflict_do_nothing(self) -> "Self":
-        """Adds an ON CONFLICT DO NOTHING clause (PostgreSQL syntax).
+    def on_conflict(self, *columns: str) -> "ConflictBuilder":
+        """Adds an ON CONFLICT clause with specified columns.
+
+        Args:
+            *columns: Column names that define the conflict. If no columns provided,
+                     creates an ON CONFLICT without specific columns (catches all conflicts).
+
+        Returns:
+            A ConflictBuilder instance for chaining conflict resolution methods.
+
+        Example:
+            ```python
+            # ON CONFLICT (id) DO NOTHING
+            sql.insert("users").values(id=1, name="John").on_conflict(
+                "id"
+            ).do_nothing()
+
+            # ON CONFLICT (email, username) DO UPDATE SET updated_at = NOW()
+            sql.insert("users").values(...).on_conflict(
+                "email", "username"
+            ).do_update(updated_at=sql.raw("NOW()"))
+
+            # ON CONFLICT DO NOTHING (catches all conflicts)
+            sql.insert("users").values(...).on_conflict().do_nothing()
+            ```
+        """
+        return ConflictBuilder(self, columns)
+
+    def on_conflict_do_nothing(self, *columns: str) -> "Insert":
+        """Adds an ON CONFLICT DO NOTHING clause (convenience method).
 
-        This is used to ignore rows that would cause a conflict.
+        Args:
+            *columns: Column names that define the conflict. If no columns provided,
+                     creates an ON CONFLICT without specific columns.
 
         Returns:
             The current builder instance for method chaining.
 
         Note:
-            This is PostgreSQL-specific syntax. Different databases have different syntax.
-            For a more general solution, you might need dialect-specific handling.
+            This is a convenience method. For more control, use on_conflict().do_nothing().
         """
-        insert_expr = self._get_insert_expression()
-        insert_expr.set("on", exp.OnConflict(this=None, expressions=[]))
-        return self
+        return self.on_conflict(*columns).do_nothing()
 
-    def on_duplicate_key_update(self, **_: Any) -> "Self":
-        """Adds an ON DUPLICATE KEY UPDATE clause (MySQL syntax).
+    def on_duplicate_key_update(self, **kwargs: Any) -> "Insert":
+        """Adds conflict resolution using the ON CONFLICT syntax (cross-database compatible).
 
         Args:
-            **_: Column-value pairs to update on duplicate key.
+            **kwargs: Column-value pairs to update on conflict.
 
         Returns:
             The current builder instance for method chaining.
+
+        Note:
+            This method uses PostgreSQL-style ON CONFLICT syntax but SQLGlot will
+            transpile it to the appropriate syntax for each database (MySQL's
+            ON DUPLICATE KEY UPDATE, etc.).
         """
-        return self
+        if not kwargs:
+            return self
+        return self.on_conflict().do_update(**kwargs)
+
+
+class ConflictBuilder:
+    """Builder for ON CONFLICT clauses in INSERT statements.
+
+    This builder provides a fluent interface for constructing conflict resolution
+    clauses using PostgreSQL-style syntax, which SQLGlot can transpile to other dialects.
+    """
+
+    __slots__ = ("_columns", "_insert_builder")
+
+    def __init__(self, insert_builder: "Insert", columns: tuple[str, ...]) -> None:
+        """Initialize ConflictBuilder.
+
+        Args:
+            insert_builder: The parent Insert builder
+            columns: Column names that define the conflict
+        """
+        self._insert_builder = insert_builder
+        self._columns = columns
+
+    def do_nothing(self) -> "Insert":
+        """Add DO NOTHING conflict resolution.
+
+        Returns:
+            The parent Insert builder for method chaining.
+
+        Example:
+            ```python
+            sql.insert("users").values(id=1, name="John").on_conflict(
+                "id"
+            ).do_nothing()
+            ```
+        """
+        insert_expr = self._insert_builder._get_insert_expression()
+
+        # 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 NOTHING"))
+
+        insert_expr.set("conflict", on_conflict)
+        return self._insert_builder
+
+    def do_update(self, **kwargs: Any) -> "Insert":
+        """Add DO UPDATE conflict resolution with SET clauses.
+
+        Args:
+            **kwargs: Column-value pairs to update on conflict.
+
+        Returns:
+            The parent Insert builder for method chaining.
+
+        Example:
+            ```python
+            sql.insert("users").values(id=1, name="John").on_conflict(
+                "id"
+            ).do_update(
+                name="Updated Name", updated_at=sql.raw("NOW()")
+            )
+            ```
+        """
+        insert_expr = self._insert_builder._get_insert_expression()
+
+        # Create SET expressions for the UPDATE
+        set_expressions = []
+        for col, val in kwargs.items():
+            if hasattr(val, "expression") and hasattr(val, "sql"):
+                # Handle SQL objects (from sql.raw with parameters)
+                expression = getattr(val, "expression", None)
+                if expression is not None and isinstance(expression, exp.Expression):
+                    # Merge parameters from SQL object into builder
+                    if hasattr(val, "parameters"):
+                        sql_parameters = getattr(val, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            self._insert_builder.add_parameter(param_value, name=param_name)
+                    value_expr = expression
+                else:
+                    # If expression is None, fall back to parsing the raw SQL
+                    sql_text = getattr(val, "sql", "")
+                    # Merge parameters even when parsing raw SQL
+                    if hasattr(val, "parameters"):
+                        sql_parameters = getattr(val, "parameters", {})
+                        for param_name, param_value in sql_parameters.items():
+                            self._insert_builder.add_parameter(param_value, name=param_name)
+                    # Check if sql_text is callable (like Expression.sql method)
+                    if callable(sql_text):
+                        sql_text = str(val)
+                    value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
+            elif isinstance(val, exp.Expression):
+                value_expr = val
+            else:
+                # Create parameter for regular values
+                param_name = self._insert_builder._generate_unique_parameter_name(col)
+                _, param_name = self._insert_builder.add_parameter(val, name=param_name)
+                value_expr = exp.Placeholder(this=param_name)
+
+            set_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
+
+        # 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,
+        )
+
+        insert_expr.set("conflict", on_conflict)
+        return self._insert_builder