sqlspec 0.16.0__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

sqlspec/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.builder._parsing_utils import parse_order_expression
+from sqlspec.exceptions import SQLBuilderError
+
+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.convert(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.convert(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/builder/mixins/_pivot_operations.py

@@ -0,0 +1,144 @@
+"""Pivot and Unpivot operations mixins for SQL builders."""
+
+from typing import TYPE_CHECKING, Optional, Union, cast
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.builder._select import Select
+
+__all__ = ("PivotClauseMixin", "UnpivotClauseMixin")
+
+
+class PivotClauseMixin:
+    """Mixin class to add PIVOT functionality to a Select."""
+
+    _expression: "Optional[exp.Expression]" = None
+    dialect: "DialectType" = None
+
+    def pivot(
+        self: "PivotClauseMixin",
+        aggregate_function: Union[str, exp.Expression],
+        aggregate_column: Union[str, exp.Expression],
+        pivot_column: Union[str, exp.Expression],
+        pivot_values: list[Union[str, int, float, exp.Expression]],
+        alias: Optional[str] = None,
+    ) -> "Select":
+        """Adds a PIVOT clause to the SELECT statement.
+
+        Example:
+            `query.pivot(aggregate_function="SUM", aggregate_column="Sales", pivot_column="Quarter", pivot_values=["Q1", "Q2", "Q3", "Q4"], alias="PivotTable")`
+
+        Args:
+            aggregate_function: The aggregate function to use (e.g., "SUM", "AVG").
+            aggregate_column: The column to be aggregated.
+            pivot_column: The column whose unique values will become new column headers.
+            pivot_values: A list of specific values from the pivot_column to be turned into columns.
+            alias: Optional alias for the pivoted table/subquery.
+
+        Returns:
+            The SelectBuilder instance for chaining.
+        """
+        current_expr = self._expression
+        if not isinstance(current_expr, exp.Select):
+            msg = "Pivot can only be applied to a Select expression managed by SelectBuilder."
+            raise TypeError(msg)
+
+        agg_func_name = aggregate_function if isinstance(aggregate_function, str) else aggregate_function.name
+        agg_col_expr = exp.column(aggregate_column) if isinstance(aggregate_column, str) else aggregate_column
+        pivot_col_expr = exp.column(pivot_column) if isinstance(pivot_column, str) else pivot_column
+
+        pivot_agg_expr = exp.func(agg_func_name, agg_col_expr)
+
+        pivot_value_exprs: list[exp.Expression] = []
+        for val in pivot_values:
+            if isinstance(val, exp.Expression):
+                pivot_value_exprs.append(val)
+            elif isinstance(val, (str, int, float)):
+                pivot_value_exprs.append(exp.convert(val))
+            else:
+                pivot_value_exprs.append(exp.convert(str(val)))
+
+        in_expr = exp.In(this=pivot_col_expr, expressions=pivot_value_exprs)
+
+        pivot_node = exp.Pivot(expressions=[pivot_agg_expr], fields=[in_expr], unpivot=False)
+
+        if alias:
+            pivot_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(pivot_node)
+                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):
+            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:
+                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)