sqlspec 0.17.0__py3-none-any.whl → 0.18.0__py3-none-any.whl

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

@@ -173,7 +173,7 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
                 else:
                     param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
                 _, param_name = self.add_parameter(value, name=param_name)
-                value_placeholders.append(exp.var(param_name))
+                value_placeholders.append(exp.Placeholder(this=param_name))
 
         tuple_expr = exp.Tuple(expressions=value_placeholders)
         if self._values_added_count == 0:
@@ -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/_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/_insert_operations.py

@@ -75,14 +75,34 @@ class InsertValuesMixin:
         if not isinstance(self._expression, exp.Insert):
             msg = "Cannot set columns on a non-INSERT expression."
             raise SQLBuilderError(msg)
-        column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns]
-        self._expression.set("columns", column_exprs)
+
+        # Get the current table from the expression
+        current_this = self._expression.args.get("this")
+        if current_this is None:
+            msg = "Table must be set using .into() before setting columns."
+            raise SQLBuilderError(msg)
+
+        if columns:
+            # Create identifiers for columns
+            column_identifiers = [exp.to_identifier(col) if isinstance(col, str) else col for col in columns]
+
+            # Get table name from current this
+            table_name = current_this.this
+
+            # Create Schema object with table and columns
+            schema = exp.Schema(this=table_name, expressions=column_identifiers)
+            self._expression.set("this", schema)
+        # No columns specified - ensure we have just a Table object
+        elif isinstance(current_this, exp.Schema):
+            table_name = current_this.this
+            self._expression.set("this", exp.Table(this=table_name))
+
         try:
             cols = self._columns
             if not columns:
                 cols.clear()
             else:
-                cols[:] = [col.name if isinstance(col, exp.Column) else str(col) for col in columns]
+                cols[:] = [col if isinstance(col, str) else str(col) for col in columns]
         except AttributeError:
             pass
         return self
@@ -128,7 +148,7 @@ class InsertValuesMixin:
                         column_name = column_name.split(".")[-1]
                     param_name = self._generate_unique_parameter_name(column_name)
                     _, param_name = self.add_parameter(val, name=param_name)
-                    row_exprs.append(exp.var(param_name))
+                    row_exprs.append(exp.Placeholder(this=param_name))
         elif len(values) == 1 and hasattr(values[0], "items"):
             mapping = values[0]
             try:
@@ -147,7 +167,7 @@ class InsertValuesMixin:
                         column_name = column_name.split(".")[-1]
                     param_name = self._generate_unique_parameter_name(column_name)
                     _, param_name = self.add_parameter(val, name=param_name)
-                    row_exprs.append(exp.var(param_name))
+                    row_exprs.append(exp.Placeholder(this=param_name))
         else:
             try:
                 _columns = self._columns
@@ -173,7 +193,7 @@ class InsertValuesMixin:
                     except AttributeError:
                         param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
                     _, param_name = self.add_parameter(v, name=param_name)
-                    row_exprs.append(exp.var(param_name))
+                    row_exprs.append(exp.Placeholder(this=param_name))
 
         values_expr = exp.Values(expressions=[row_exprs])
         self._expression.set("expression", values_expr)

sqlspec/builder/mixins/_merge_operations.py

@@ -365,7 +365,7 @@ class MergeNotMatchedClauseMixin:
                     column_name = column_name.split(".")[-1]
                 param_name = self._generate_unique_parameter_name(column_name)
                 param_name = self.add_parameter(val, name=param_name)[1]
-                parameterized_values.append(exp.var(param_name))
+                parameterized_values.append(exp.Placeholder())
 
             insert_args["this"] = exp.Tuple(expressions=[exp.column(c) for c in columns])
             insert_args["expression"] = exp.Tuple(expressions=parameterized_values)

sqlspec/builder/mixins/_order_limit_operations.py

@@ -10,6 +10,9 @@ from sqlspec.builder._parsing_utils import parse_order_expression
 from sqlspec.exceptions import SQLBuilderError
 
 if TYPE_CHECKING:
+    from sqlspec.builder._column import Column
+    from sqlspec.builder._expression_wrappers import ExpressionWrapper
+    from sqlspec.builder.mixins._select_operations import Case
     from sqlspec.protocols import SQLBuilderProtocol
 
 __all__ = ("LimitOffsetClauseMixin", "OrderByClauseMixin", "ReturningClauseMixin")
@@ -24,7 +27,7 @@ class OrderByClauseMixin:
     # Type annotation for PyRight - this will be provided by the base class
     _expression: Optional[exp.Expression]
 
-    def order_by(self, *items: Union[str, exp.Ordered], desc: bool = False) -> Self:
+    def order_by(self, *items: Union[str, exp.Ordered, "Column"], desc: bool = False) -> Self:
         """Add ORDER BY clause.
 
         Args:
@@ -49,7 +52,13 @@ class OrderByClauseMixin:
                 if desc:
                     order_item = order_item.desc()
             else:
-                order_item = item
+                # Extract expression from Column objects or use as-is for sqlglot expressions
+                from sqlspec._sql import SQLFactory
+
+                extracted_item = SQLFactory._extract_expression(item)
+                order_item = extracted_item
+                if desc and not isinstance(item, exp.Ordered):
+                    order_item = order_item.desc()
             current_expr = current_expr.order_by(order_item, copy=False)
         builder._expression = current_expr
         return cast("Self", builder)
@@ -111,7 +120,7 @@ class ReturningClauseMixin:
     # Type annotation for PyRight - this will be provided by the base class
     _expression: Optional[exp.Expression]
 
-    def returning(self, *columns: Union[str, exp.Expression]) -> Self:
+    def returning(self, *columns: Union[str, exp.Expression, "Column", "ExpressionWrapper", "Case"]) -> Self:
         """Add RETURNING clause to the statement.
 
         Args:
@@ -130,6 +139,9 @@ class ReturningClauseMixin:
         if not isinstance(self._expression, valid_types):
             msg = "RETURNING is only supported for INSERT, UPDATE, and DELETE statements."
             raise SQLBuilderError(msg)
-        returning_exprs = [exp.column(c) if isinstance(c, str) else c for c in columns]
+        # Extract expressions from various wrapper types
+        from sqlspec._sql import SQLFactory
+
+        returning_exprs = [SQLFactory._extract_expression(c) for c in columns]
         self._expression.set("returning", exp.Returning(expressions=returning_exprs))
         return self

sqlspec/builder/mixins/_select_operations.py

@@ -1,6 +1,5 @@
 """SELECT clause mixins consolidated into a single module."""
 
-from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
 from mypy_extensions import trait
@@ -538,13 +537,10 @@ class SelectClauseMixin:
         return CaseBuilder(builder, alias)
 
 
-@dataclass
 class CaseBuilder:
     """Builder for CASE expressions."""
 
-    _parent: "SelectBuilderProtocol"
-    _alias: Optional[str]
-    _case_expr: exp.Case
+    __slots__ = ("_alias", "_case_expr", "_parent")
 
     def __init__(self, parent: "SelectBuilderProtocol", alias: "Optional[str]" = None) -> None:
         """Initialize CaseBuilder.
@@ -858,7 +854,7 @@ class Case:
         from sqlspec._sql import SQLFactory
 
         cond_expr = exp.maybe_parse(condition) or exp.column(condition) if isinstance(condition, str) else condition
-        val_expr = SQLFactory._to_literal(value)
+        val_expr = SQLFactory._to_expression(value)
 
         # SQLGlot uses exp.If for CASE WHEN clauses, not exp.When
         when_clause = exp.If(this=cond_expr, true=val_expr)
@@ -876,7 +872,7 @@ class Case:
         """
         from sqlspec._sql import SQLFactory
 
-        self._default = SQLFactory._to_literal(value)
+        self._default = SQLFactory._to_expression(value)
         return self
 
     def end(self) -> Self:

sqlspec/builder/mixins/_update_operations.py

@@ -111,10 +111,10 @@ class UpdateSetClauseMixin:
         """Set columns and values for the UPDATE statement.
 
         Supports:
-        - set(column, value)
-        - set(mapping)
-        - set(**kwargs)
-        - set(mapping, **kwargs)
+        - set_(column, value)
+        - set_(mapping)
+        - set_(**kwargs)
+        - set_(mapping, **kwargs)
 
         Args:
             *args: Either (column, value) or a mapping.

sqlspec/config.py

@@ -5,6 +5,7 @@ from typing_extensions import NotRequired, TypedDict
 
 from sqlspec.core.parameters import ParameterStyle, ParameterStyleConfig
 from sqlspec.core.statement import StatementConfig
+from sqlspec.migrations.tracker import AsyncMigrationTracker, SyncMigrationTracker
 from sqlspec.utils.logging import get_logger
 
 if TYPE_CHECKING:
@@ -21,6 +22,7 @@ __all__ = (
     "DatabaseConfigProtocol",
     "DriverT",
     "LifecycleConfig",
+    "MigrationConfig",
     "NoPoolAsyncConfig",
     "NoPoolSyncConfig",
     "SyncConfigT",
@@ -43,7 +45,7 @@ logger = get_logger("config")
 
 
 class LifecycleConfig(TypedDict, total=False):
-    """Universal lifecycle hooks for all adapters.
+    """Lifecycle hooks for all adapters.
 
     Each hook accepts a list of callables to support multiple handlers.
     """
@@ -59,6 +61,22 @@ class LifecycleConfig(TypedDict, total=False):
     on_error: NotRequired[list[Callable[[Exception, str, dict], None]]]
 
 
+class MigrationConfig(TypedDict, total=False):
+    """Configuration options for database migrations.
+
+    All fields are optional with sensible defaults.
+    """
+
+    script_location: NotRequired[str]
+    """Path to the migrations directory. Defaults to 'migrations'."""
+
+    version_table_name: NotRequired[str]
+    """Name of the table used to track applied migrations. Defaults to 'sqlspec_migrations'."""
+
+    project_root: NotRequired[str]
+    """Path to the project root directory. Used for relative path resolution."""
+
+
 class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
     """Protocol defining the interface for database configurations."""
 
@@ -73,7 +91,7 @@ class DatabaseConfigProtocol(ABC, Generic[ConnectionT, PoolT, DriverT]):
     supports_native_parquet_export: "ClassVar[bool]" = False
     statement_config: "StatementConfig"
     pool_instance: "Optional[PoolT]"
-    migration_config: "dict[str, Any]"
+    migration_config: "Union[dict[str, Any], MigrationConfig]"
 
     def __hash__(self) -> int:
         return id(self)
@@ -142,18 +160,19 @@ class NoPoolSyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
     __slots__ = ("connection_config",)
     is_async: "ClassVar[bool]" = False
     supports_connection_pooling: "ClassVar[bool]" = False
+    migration_tracker_type: "ClassVar[type[Any]]" = SyncMigrationTracker
 
     def __init__(
         self,
         *,
         connection_config: Optional[dict[str, Any]] = None,
-        migration_config: "Optional[dict[str, Any]]" = None,
+        migration_config: "Optional[Union[dict[str, Any], MigrationConfig]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         self.pool_instance = None
         self.connection_config = connection_config or {}
-        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+        self.migration_config: Union[dict[str, Any], MigrationConfig] = migration_config or {}
 
         if statement_config is None:
             default_parameter_config = ParameterStyleConfig(
@@ -192,21 +211,21 @@ class NoPoolAsyncConfig(DatabaseConfigProtocol[ConnectionT, None, DriverT]):
     """Base class for an async database configurations that do not implement a pool."""
 
     __slots__ = ("connection_config",)
-
     is_async: "ClassVar[bool]" = True
     supports_connection_pooling: "ClassVar[bool]" = False
+    migration_tracker_type: "ClassVar[type[Any]]" = AsyncMigrationTracker
 
     def __init__(
         self,
         *,
         connection_config: "Optional[dict[str, Any]]" = None,
-        migration_config: "Optional[dict[str, Any]]" = None,
+        migration_config: "Optional[Union[dict[str, Any], MigrationConfig]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         self.pool_instance = None
         self.connection_config = connection_config or {}
-        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+        self.migration_config: Union[dict[str, Any], MigrationConfig] = migration_config or {}
 
         if statement_config is None:
             default_parameter_config = ParameterStyleConfig(
@@ -245,22 +264,22 @@ class SyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
     """Generic Sync Database Configuration."""
 
     __slots__ = ("pool_config",)
-
     is_async: "ClassVar[bool]" = False
     supports_connection_pooling: "ClassVar[bool]" = True
+    migration_tracker_type: "ClassVar[type[Any]]" = SyncMigrationTracker
 
     def __init__(
         self,
         *,
         pool_config: "Optional[dict[str, Any]]" = None,
         pool_instance: "Optional[PoolT]" = None,
-        migration_config: "Optional[dict[str, Any]]" = None,
+        migration_config: "Optional[Union[dict[str, Any], MigrationConfig]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         self.pool_instance = pool_instance
         self.pool_config = pool_config or {}
-        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+        self.migration_config: Union[dict[str, Any], MigrationConfig] = migration_config or {}
 
         if statement_config is None:
             default_parameter_config = ParameterStyleConfig(
@@ -321,22 +340,22 @@ class AsyncDatabaseConfig(DatabaseConfigProtocol[ConnectionT, PoolT, DriverT]):
     """Generic Async Database Configuration."""
 
     __slots__ = ("pool_config",)
-
     is_async: "ClassVar[bool]" = True
     supports_connection_pooling: "ClassVar[bool]" = True
+    migration_tracker_type: "ClassVar[type[Any]]" = AsyncMigrationTracker
 
     def __init__(
         self,
         *,
         pool_config: "Optional[dict[str, Any]]" = None,
         pool_instance: "Optional[PoolT]" = None,
-        migration_config: "Optional[dict[str, Any]]" = None,
+        migration_config: "Optional[Union[dict[str, Any], MigrationConfig]]" = None,
         statement_config: "Optional[StatementConfig]" = None,
         driver_features: "Optional[dict[str, Any]]" = None,
     ) -> None:
         self.pool_instance = pool_instance
         self.pool_config = pool_config or {}
-        self.migration_config: dict[str, Any] = migration_config if migration_config is not None else {}
+        self.migration_config: Union[dict[str, Any], MigrationConfig] = migration_config or {}
 
         if statement_config is None:
             self.statement_config = StatementConfig(

sqlspec/core/__init__.py

@@ -1,17 +1,92 @@
-"""SQLSpec Core Module - SQL Processing System.
-
-This module provides the core SQL processing components including statement handling,
-parameter processing, compilation, and result management.
-
-Components:
-- statement.py: SQL class with StatementConfig
-- parameters.py: Parameter processing pipeline
-- compiler.py: SQL compilation with caching
-- result.py: Result classes for query execution
-- filters.py: Statement filter system
-- cache.py: Unified caching system
-- splitter.py: SQL statement splitter
-- hashing.py: Cache key generation
+"""SQLSpec Core Module - High-Performance SQL Processing System.
+
+This module provides the core SQL processing infrastructure for SQLSpec, implementing
+a complete pipeline for SQL statement compilation, parameter processing, caching,
+and result management. All components are optimized for MyPyC compilation and
+designed for maximum performance with minimal overhead.
+
+Architecture Overview:
+    The core module implements a single-pass processing pipeline where SQL statements
+    are parsed once, transformed once, and validated once. The SQL object serves as
+    the single source of truth throughout the system.
+
+Key Components:
+    statement.py: SQL statement representation and configuration management
+        - SQL class for statement encapsulation with lazy compilation
+        - StatementConfig for processing pipeline configuration
+        - ProcessedState for cached compilation results
+        - Support for execute_many and script execution modes
+
+    parameters.py: Type-safe parameter processing and style conversion
+        - Automatic parameter style detection and conversion
+        - Support for QMARK (?), NAMED (:name), NUMERIC ($1), FORMAT (%s) styles
+        - Parameter validation and type coercion
+        - Batch parameter handling for execute_many operations
+
+    compiler.py: SQL compilation with validation and optimization
+        - SQLProcessor for statement compilation and validation
+        - Operation type detection (SELECT, INSERT, UPDATE, DELETE, etc.)
+        - AST-based SQL analysis using SQLGlot
+        - Support for multiple SQL dialects
+        - Compiled result caching for performance
+
+    result.py: Comprehensive result handling for all SQL operations
+        - SQLResult for standard query results with metadata
+        - ArrowResult for Apache Arrow format integration
+        - Support for DML operations with RETURNING clauses
+        - Script execution result aggregation
+        - Iterator protocol support for result rows
+
+    filters.py: Composable SQL statement filters
+        - BeforeAfterFilter for date range filtering
+        - InCollectionFilter for IN clause generation
+        - LimitOffsetFilter for pagination
+        - OrderByFilter for dynamic sorting
+        - SearchFilter for text search operations
+        - Parameter conflict resolution
+
+    cache.py: Unified caching system with LRU eviction
+        - UnifiedCache with configurable TTL and size limits
+        - StatementCache for compiled SQL statements
+        - ExpressionCache for parsed SQLGlot expressions
+        - ParameterCache for processed parameters
+        - Thread-safe operations with fine-grained locking
+        - Cache statistics and monitoring
+
+    splitter.py: Dialect-aware SQL script splitting
+        - Support for Oracle PL/SQL, T-SQL, PostgreSQL, MySQL
+        - Proper handling of block structures (BEGIN/END)
+        - Dollar-quoted string support for PostgreSQL
+        - Batch separator recognition (GO for T-SQL)
+        - Comment and string literal preservation
+
+    hashing.py: Efficient cache key generation
+        - SQL statement hashing with parameter consideration
+        - Expression tree hashing for AST caching
+        - Parameter set hashing for batch operations
+        - Optimized hash computation with caching
+
+Performance Optimizations:
+    - MyPyC compilation support with proper annotations
+    - __slots__ usage for memory efficiency
+    - Final annotations for constant folding
+    - Lazy evaluation and compilation
+    - Comprehensive result caching
+    - Minimal object allocation in hot paths
+
+Thread Safety:
+    All caching components are thread-safe with RLock protection.
+    The processing pipeline is stateless and safe for concurrent use.
+
+Example Usage:
+    >>> from sqlspec.core import SQL, StatementConfig
+    >>> config = StatementConfig(dialect="postgresql")
+    >>> stmt = SQL(
+    ...     "SELECT * FROM users WHERE id = ?",
+    ...     1,
+    ...     statement_config=config,
+    ... )
+    >>> compiled_sql, params = stmt.compile()
 """
 
 from sqlspec.core import filters