sqlspec 0.13.1__py3-none-any.whl → 0.14.0__py3-none-any.whl

sqlspec/statement/builder/mixins/__init__.py

@@ -1,18 +1,14 @@
 """SQL statement builder mixins."""
 
-from sqlspec.statement.builder.mixins._aggregate_functions import AggregateFunctionsMixin
-from sqlspec.statement.builder.mixins._case_builder import CaseBuilderMixin
-from sqlspec.statement.builder.mixins._common_table_expr import CommonTableExpressionMixin
-from sqlspec.statement.builder.mixins._delete_from import DeleteFromClauseMixin
-from sqlspec.statement.builder.mixins._from import FromClauseMixin
-from sqlspec.statement.builder.mixins._group_by import GroupByClauseMixin
-from sqlspec.statement.builder.mixins._having import HavingClauseMixin
-from sqlspec.statement.builder.mixins._insert_from_select import InsertFromSelectMixin
-from sqlspec.statement.builder.mixins._insert_into import InsertIntoClauseMixin
-from sqlspec.statement.builder.mixins._insert_values import InsertValuesMixin
-from sqlspec.statement.builder.mixins._join import JoinClauseMixin
-from sqlspec.statement.builder.mixins._limit_offset import LimitOffsetClauseMixin
-from sqlspec.statement.builder.mixins._merge_clauses import (
+from sqlspec.statement.builder.mixins._cte_and_set_ops import CommonTableExpressionMixin, SetOperationMixin
+from sqlspec.statement.builder.mixins._delete_operations import DeleteFromClauseMixin
+from sqlspec.statement.builder.mixins._insert_operations import (
+    InsertFromSelectMixin,
+    InsertIntoClauseMixin,
+    InsertValuesMixin,
+)
+from sqlspec.statement.builder.mixins._join_operations import JoinClauseMixin
+from sqlspec.statement.builder.mixins._merge_operations import (
     MergeIntoClauseMixin,
     MergeMatchedClauseMixin,
     MergeNotMatchedBySourceClauseMixin,
@@ -20,25 +16,24 @@ from sqlspec.statement.builder.mixins._merge_clauses import (
     MergeOnClauseMixin,
     MergeUsingClauseMixin,
 )
-from sqlspec.statement.builder.mixins._order_by import OrderByClauseMixin
-from sqlspec.statement.builder.mixins._pivot import PivotClauseMixin
-from sqlspec.statement.builder.mixins._returning import ReturningClauseMixin
-from sqlspec.statement.builder.mixins._select_columns import SelectColumnsMixin
-from sqlspec.statement.builder.mixins._set_ops import SetOperationMixin
-from sqlspec.statement.builder.mixins._unpivot import UnpivotClauseMixin
-from sqlspec.statement.builder.mixins._update_from import UpdateFromClauseMixin
-from sqlspec.statement.builder.mixins._update_set import UpdateSetClauseMixin
-from sqlspec.statement.builder.mixins._update_table import UpdateTableClauseMixin
-from sqlspec.statement.builder.mixins._where import WhereClauseMixin
-from sqlspec.statement.builder.mixins._window_functions import WindowFunctionsMixin
+from sqlspec.statement.builder.mixins._order_limit_operations import (
+    LimitOffsetClauseMixin,
+    OrderByClauseMixin,
+    ReturningClauseMixin,
+)
+from sqlspec.statement.builder.mixins._pivot_operations import PivotClauseMixin, UnpivotClauseMixin
+from sqlspec.statement.builder.mixins._select_operations import CaseBuilder, SelectClauseMixin
+from sqlspec.statement.builder.mixins._update_operations import (
+    UpdateFromClauseMixin,
+    UpdateSetClauseMixin,
+    UpdateTableClauseMixin,
+)
+from sqlspec.statement.builder.mixins._where_clause import HavingClauseMixin, WhereClauseMixin
 
 __all__ = (
-    "AggregateFunctionsMixin",
-    "CaseBuilderMixin",
+    "CaseBuilder",
     "CommonTableExpressionMixin",
     "DeleteFromClauseMixin",
-    "FromClauseMixin",
-    "GroupByClauseMixin",
     "HavingClauseMixin",
     "InsertFromSelectMixin",
     "InsertIntoClauseMixin",
@@ -54,12 +49,11 @@ __all__ = (
     "OrderByClauseMixin",
     "PivotClauseMixin",
     "ReturningClauseMixin",
-    "SelectColumnsMixin",
+    "SelectClauseMixin",
     "SetOperationMixin",
     "UnpivotClauseMixin",
     "UpdateFromClauseMixin",
     "UpdateSetClauseMixin",
     "UpdateTableClauseMixin",
     "WhereClauseMixin",
-    "WindowFunctionsMixin",
 )

sqlspec/statement/builder/mixins/{_set_ops.py → _cte_and_set_ops.py}

@@ -1,11 +1,95 @@
-from typing import Any, Optional
+"""CTE (Common Table Expression) and Set Operations mixins for SQL builders."""
+
+from typing import Any, Optional, Union
 
 from sqlglot import exp
 from typing_extensions import Self
 
 from sqlspec.exceptions import SQLBuilderError
 
-__all__ = ("SetOperationMixin",)
+__all__ = ("CommonTableExpressionMixin", "SetOperationMixin")
+
+
+class CommonTableExpressionMixin:
+    """Mixin providing WITH clause (Common Table Expressions) support for SQL builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def with_(
+        self, name: str, query: Union[Any, str], recursive: bool = False, columns: Optional[list[str]] = None
+    ) -> Self:
+        """Add WITH clause (Common Table Expression).
+
+        Args:
+            name: The name of the CTE.
+            query: The query for the CTE (builder instance or SQL string).
+            recursive: Whether this is a recursive CTE.
+            columns: Optional column names for the CTE.
+
+        Raises:
+            SQLBuilderError: If the query type is unsupported.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            msg = "Cannot add WITH clause: expression not initialized."
+            raise SQLBuilderError(msg)
+
+        if not hasattr(self._expression, "with_") and not isinstance(
+            self._expression, (exp.Select, exp.Insert, exp.Update, exp.Delete)
+        ):
+            msg = f"Cannot add WITH clause to {type(self._expression).__name__} expression."
+            raise SQLBuilderError(msg)
+
+        cte_expr: Optional[exp.Expression] = None
+        if hasattr(query, "to_statement"):
+            # Query is a builder instance
+            built_query = query.to_statement()  # pyright: ignore
+            cte_sql = built_query.to_sql()
+            cte_expr = exp.maybe_parse(cte_sql, dialect=getattr(self, "dialect", None))
+
+            # Merge parameters
+            if hasattr(self, "add_parameter"):
+                parameters = getattr(built_query, "parameters", None) or {}
+                for param_name, param_value in parameters.items():
+                    self.add_parameter(param_value, name=param_name)  # pyright: ignore
+        elif isinstance(query, str):
+            cte_expr = exp.maybe_parse(query, dialect=getattr(self, "dialect", None))
+        elif isinstance(query, exp.Expression):
+            cte_expr = query
+
+        if not cte_expr:
+            msg = f"Could not parse CTE query: {query}"
+            raise SQLBuilderError(msg)
+
+        if columns:
+            # CTE with explicit column list: name(col1, col2, ...)
+            cte_alias_expr = exp.alias_(cte_expr, name, table=[exp.to_identifier(col) for col in columns])
+        else:
+            # Simple CTE alias: name
+            cte_alias_expr = exp.alias_(cte_expr, name)
+
+        # Different handling for different expression types
+        if hasattr(self._expression, "with_"):
+            existing_with = self._expression.args.get("with")  # pyright: ignore
+            if existing_with:
+                existing_with.expressions.append(cte_alias_expr)
+                if recursive:
+                    existing_with.set("recursive", recursive)
+            else:
+                self._expression = self._expression.with_(cte_alias_expr, as_=name, copy=False)  # pyright: ignore
+                if recursive:
+                    with_clause = self._expression.find(exp.With)
+                    if with_clause:
+                        with_clause.set("recursive", recursive)
+        else:
+            # Store CTEs for later application during build
+            if not hasattr(self, "_with_ctes"):
+                setattr(self, "_with_ctes", {})
+            self._with_ctes[name] = exp.CTE(this=cte_expr, alias=exp.to_table(name))  # type: ignore[attr-defined]
+
+        return self
 
 
 class SetOperationMixin:

sqlspec/statement/builder/mixins/{_delete_from.py → _delete_operations.py}

@@ -1,3 +1,5 @@
+"""Delete operation mixins for SQL builders."""
+
 from typing import Optional
 
 from sqlglot import exp

sqlspec/statement/builder/mixins/{_insert_values.py → _insert_operations.py}

@@ -1,3 +1,5 @@
+"""Insert operation mixins for SQL builders."""
+
 from collections.abc import Sequence
 from typing import Any, Optional, Union
 
@@ -6,7 +8,35 @@ from typing_extensions import Self
 
 from sqlspec.exceptions import SQLBuilderError
 
-__all__ = ("InsertValuesMixin",)
+__all__ = ("InsertFromSelectMixin", "InsertIntoClauseMixin", "InsertValuesMixin")
+
+
+class InsertIntoClauseMixin:
+    """Mixin providing INTO clause for INSERT builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def into(self, table: str) -> Self:
+        """Set the target table for the INSERT statement.
+
+        Args:
+            table: The name of the table to insert data into.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an INSERT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot set target table on a non-INSERT expression."
+            raise SQLBuilderError(msg)
+
+        setattr(self, "_table", table)
+        self._expression.set("this", exp.to_table(table))
+        return self
 
 
 class InsertValuesMixin:
@@ -65,3 +95,42 @@ class InsertValuesMixin:
             The current builder instance for method chaining.
         """
         return self.values(*values)
+
+
+class InsertFromSelectMixin:
+    """Mixin providing INSERT ... SELECT support for INSERT builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def from_select(self, select_builder: Any) -> Self:
+        """Sets the INSERT source to a SELECT statement.
+
+        Args:
+            select_builder: A SelectBuilder instance representing the SELECT query.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Raises:
+            SQLBuilderError: If the table is not set or the select_builder is invalid.
+        """
+        if not getattr(self, "_table", None):
+            msg = "The target table must be set using .into() before adding values."
+            raise SQLBuilderError(msg)
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot set INSERT source on a non-INSERT expression."
+            raise SQLBuilderError(msg)
+        # Merge parameters from the SELECT builder
+        subquery_params = getattr(select_builder, "_parameters", None)
+        if subquery_params:
+            for p_name, p_value in subquery_params.items():
+                self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+        select_expr = getattr(select_builder, "_expression", None)
+        if select_expr and isinstance(select_expr, exp.Select):
+            self._expression.set("expression", select_expr.copy())
+        else:
+            msg = "SelectBuilder must have a valid SELECT expression."
+            raise SQLBuilderError(msg)
+        return self

sqlspec/statement/builder/mixins/{_merge_clauses.py → _merge_operations.py}

@@ -1,3 +1,5 @@
+"""Merge operation mixins for SQL builders."""
+
 from typing import Any, Optional, Union
 
 from sqlglot import exp

sqlspec/statement/builder/mixins/_order_limit_operations.py

@@ -0,0 +1,123 @@
+"""Order, Limit, Offset and Returning operations mixins for SQL builders."""
+
+from typing import TYPE_CHECKING, Optional, Union, cast
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+from sqlspec.statement.builder._parsing_utils import parse_order_expression
+
+if TYPE_CHECKING:
+    from sqlspec.protocols import SQLBuilderProtocol
+
+__all__ = ("LimitOffsetClauseMixin", "OrderByClauseMixin", "ReturningClauseMixin")
+
+
+class OrderByClauseMixin:
+    """Mixin providing ORDER BY clause."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def order_by(self, *items: Union[str, exp.Ordered], desc: bool = False) -> Self:
+        """Add ORDER BY clause.
+
+        Args:
+            *items: Columns to order by. Can be strings (column names) or sqlglot.exp.Ordered instances for specific directions (e.g., exp.column("name").desc()).
+            desc: Whether to order in descending order (applies to all items if they are strings).
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement or if the item type is unsupported.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SQLBuilderProtocol", self)
+        if not isinstance(builder._expression, exp.Select):
+            msg = "ORDER BY is only supported for SELECT statements."
+            raise SQLBuilderError(msg)
+
+        current_expr = builder._expression
+        for item in items:
+            if isinstance(item, str):
+                order_item = parse_order_expression(item)
+                if desc:
+                    order_item = order_item.desc()
+            else:
+                order_item = item
+            current_expr = current_expr.order_by(order_item, copy=False)
+        builder._expression = current_expr
+        return cast("Self", builder)
+
+
+class LimitOffsetClauseMixin:
+    """Mixin providing LIMIT and OFFSET clauses."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def limit(self, value: int) -> Self:
+        """Add LIMIT clause.
+
+        Args:
+            value: The maximum number of rows to return.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SQLBuilderProtocol", self)
+        if not isinstance(builder._expression, exp.Select):
+            msg = "LIMIT is only supported for SELECT statements."
+            raise SQLBuilderError(msg)
+        builder._expression = builder._expression.limit(exp.Literal.number(value), copy=False)
+        return cast("Self", builder)
+
+    def offset(self, value: int) -> Self:
+        """Add OFFSET clause.
+
+        Args:
+            value: The number of rows to skip before starting to return rows.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SQLBuilderProtocol", self)
+        if not isinstance(builder._expression, exp.Select):
+            msg = "OFFSET is only supported for SELECT statements."
+            raise SQLBuilderError(msg)
+        builder._expression = builder._expression.offset(exp.Literal.number(value), copy=False)
+        return cast("Self", builder)
+
+
+class ReturningClauseMixin:
+    """Mixin providing RETURNING clause."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def returning(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add RETURNING clause to the statement.
+
+        Args:
+            *columns: Columns to return. Can be strings or sqlglot expressions.
+
+        Raises:
+            SQLBuilderError: If the current expression is not INSERT, UPDATE, or DELETE.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            msg = "Cannot add RETURNING: expression is not initialized."
+            raise SQLBuilderError(msg)
+        valid_types = (exp.Insert, exp.Update, exp.Delete)
+        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]
+        self._expression.set("returning", exp.Returning(expressions=returning_exprs))
+        return self

sqlspec/statement/builder/mixins/{_pivot.py → _pivot_operations.py}

@@ -1,3 +1,5 @@
+"""Pivot and Unpivot operations mixins for SQL builders."""
+
 from typing import TYPE_CHECKING, Optional, Union, cast
 
 from sqlglot import exp
@@ -5,9 +7,9 @@ from sqlglot import exp
 if TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
 
-    from sqlspec.statement.builder.select import Select
+    from sqlspec.statement.builder._select import Select
 
-__all__ = ("PivotClauseMixin",)
+__all__ = ("PivotClauseMixin", "UnpivotClauseMixin")
 
 
 class PivotClauseMixin:
@@ -77,3 +79,70 @@ class PivotClauseMixin:
                 table.set("pivots", existing_pivots)
 
         return cast("Select", self)
+
+
+class UnpivotClauseMixin:
+    """Mixin class to add UNPIVOT functionality to a Select."""
+
+    _expression: "Optional[exp.Expression]" = None
+    dialect: "DialectType" = None
+
+    def unpivot(
+        self: "UnpivotClauseMixin",
+        value_column_name: str,
+        name_column_name: str,
+        columns_to_unpivot: list[Union[str, exp.Expression]],
+        alias: Optional[str] = None,
+    ) -> "Select":
+        """Adds an UNPIVOT clause to the SELECT statement.
+
+        Example:
+            `query.unpivot(value_column_name="Sales", name_column_name="Quarter", columns_to_unpivot=["Q1Sales", "Q2Sales"], alias="UnpivotTable")`
+
+        Args:
+            value_column_name: The name for the new column that will hold the values from the unpivoted columns.
+            name_column_name: The name for the new column that will hold the names of the original unpivoted columns.
+            columns_to_unpivot: A list of columns to be unpivoted into rows.
+            alias: Optional alias for the unpivoted table/subquery.
+
+        Raises:
+            TypeError: If the current expression is not a Select expression.
+
+        Returns:
+            The Select instance for chaining.
+        """
+        current_expr = self._expression
+        if not isinstance(current_expr, exp.Select):
+            # SelectBuilder's __init__ ensures _expression is exp.Select.
+            msg = "Unpivot can only be applied to a Select expression managed by Select."
+            raise TypeError(msg)
+
+        value_col_ident = exp.to_identifier(value_column_name)
+        name_col_ident = exp.to_identifier(name_column_name)
+
+        unpivot_cols_exprs: list[exp.Expression] = []
+        for col_name_or_expr in columns_to_unpivot:
+            if isinstance(col_name_or_expr, exp.Expression):
+                unpivot_cols_exprs.append(col_name_or_expr)
+            elif isinstance(col_name_or_expr, str):
+                unpivot_cols_exprs.append(exp.column(col_name_or_expr))
+            else:
+                # Fallback for other types, should ideally be an error or more specific handling
+                unpivot_cols_exprs.append(exp.column(str(col_name_or_expr)))
+
+        in_expr = exp.In(this=name_col_ident, expressions=unpivot_cols_exprs)
+
+        unpivot_node = exp.Pivot(expressions=[value_col_ident], fields=[in_expr], unpivot=True)
+
+        if alias:
+            unpivot_node.set("alias", exp.TableAlias(this=exp.to_identifier(alias)))
+
+        from_clause = current_expr.args.get("from")
+        if from_clause and isinstance(from_clause, exp.From):
+            table = from_clause.this
+            if isinstance(table, exp.Table):
+                existing_pivots = table.args.get("pivots", [])
+                existing_pivots.append(unpivot_node)
+                table.set("pivots", existing_pivots)
+
+        return cast("Select", self)