sqlspec 0.11.1__py3-none-any.whl → 0.12.1__py3-none-any.whl

sqlspec/statement/builder/mixins/_pivot.py

@@ -0,0 +1,82 @@
+from typing import TYPE_CHECKING, Optional, Union, cast
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.statement.builder.select import SelectBuilder
+
+__all__ = ("PivotClauseMixin",)
+
+
+class PivotClauseMixin:
+    """Mixin class to add PIVOT functionality to a SelectBuilder."""
+
+    _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,
+    ) -> "SelectBuilder":
+        """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):
+                pivot_value_exprs.append(exp.Literal.string(val))
+            elif isinstance(val, (int, float)):
+                pivot_value_exprs.append(exp.Literal.number(val))
+            else:
+                pivot_value_exprs.append(exp.Literal.string(str(val)))
+
+        # Create the pivot expression with proper fields structure
+        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)))
+
+        # Add pivot to the table in the FROM clause
+        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):
+                # Add to pivots array
+                existing_pivots = table.args.get("pivots", [])
+                existing_pivots.append(pivot_node)
+                table.set("pivots", existing_pivots)
+
+        return cast("SelectBuilder", self)

sqlspec/statement/builder/mixins/_returning.py

@@ -0,0 +1,37 @@
+from typing import Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("ReturningClauseMixin",)
+
+
+class ReturningClauseMixin:
+    """Mixin providing RETURNING clause for INSERT, UPDATE, and DELETE builders."""
+
+    _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/_select_columns.py

@@ -0,0 +1,60 @@
+from typing import TYPE_CHECKING, Union, cast
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+from sqlspec.statement.builder._parsing_utils import parse_column_expression
+
+if TYPE_CHECKING:
+    from sqlspec.statement.builder.protocols import BuilderProtocol
+
+__all__ = ("SelectColumnsMixin",)
+
+
+class SelectColumnsMixin:
+    """Mixin providing SELECT column and DISTINCT clauses for SELECT builders."""
+
+    def select(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add columns to SELECT clause.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("BuilderProtocol", self)
+        if builder._expression is None:
+            builder._expression = exp.Select()
+        if not isinstance(builder._expression, exp.Select):
+            msg = "Cannot add select columns to a non-SELECT expression."
+            raise SQLBuilderError(msg)
+        for column in columns:
+            builder._expression = builder._expression.select(parse_column_expression(column), copy=False)
+        return cast("Self", builder)
+
+    def distinct(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add DISTINCT clause to SELECT.
+
+        Args:
+            *columns: Optional columns to make distinct. If none provided, applies DISTINCT to all selected columns.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("BuilderProtocol", self)
+        if builder._expression is None:
+            builder._expression = exp.Select()
+        if not isinstance(builder._expression, exp.Select):
+            msg = "Cannot add DISTINCT to a non-SELECT expression."
+            raise SQLBuilderError(msg)
+        if not columns:
+            builder._expression.set("distinct", exp.Distinct())
+        else:
+            distinct_columns = [parse_column_expression(column) for column in columns]
+            builder._expression.set("distinct", exp.Distinct(expressions=distinct_columns))
+        return cast("Self", builder)

sqlspec/statement/builder/mixins/_set_ops.py

@@ -0,0 +1,122 @@
+from typing import Any, Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("SetOperationMixin",)
+
+
+class SetOperationMixin:
+    """Mixin providing set operations (UNION, INTERSECT, EXCEPT) for SELECT builders."""
+
+    _expression: Any = None
+    _parameters: dict[str, Any] = {}
+    dialect: Any = None
+
+    def union(self, other: Any, all_: bool = False) -> Self:
+        """Combine this query with another using UNION.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to union with.
+            all_: If True, use UNION ALL instead of UNION.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the union query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=getattr(self, "dialect", None))
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=getattr(self, "dialect", None))
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for UNION operation"
+            raise SQLBuilderError(msg)
+        union_expr = exp.union(left_expr, right_expr, distinct=not all_)
+        new_builder = type(self)()
+        new_builder.dialect = getattr(self, "dialect", None)
+        new_builder._expression = union_expr
+        merged_params = dict(left_query.parameters)
+        for param_name, param_value in right_query.parameters.items():
+            if param_name in merged_params:
+                counter = 1
+                new_param_name = f"{param_name}_right_{counter}"
+                while new_param_name in merged_params:
+                    counter += 1
+                    new_param_name = f"{param_name}_right_{counter}"
+
+                # Use AST transformation instead of string manipulation
+                def rename_parameter(node: exp.Expression) -> exp.Expression:
+                    if isinstance(node, exp.Placeholder) and node.name == param_name:  # noqa: B023
+                        return exp.Placeholder(this=new_param_name)  # noqa: B023
+                    return node
+
+                right_expr = right_expr.transform(rename_parameter)
+                union_expr = exp.union(left_expr, right_expr, distinct=not all_)
+                new_builder._expression = union_expr
+                merged_params[new_param_name] = param_value
+            else:
+                merged_params[param_name] = param_value
+        new_builder._parameters = merged_params
+        return new_builder
+
+    def intersect(self, other: Any) -> Self:
+        """Add INTERSECT clause.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to intersect with.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the intersect query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=getattr(self, "dialect", None))
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=getattr(self, "dialect", None))
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for INTERSECT operation"
+            raise SQLBuilderError(msg)
+        intersect_expr = exp.intersect(left_expr, right_expr, distinct=True)
+        new_builder = type(self)()
+        new_builder.dialect = getattr(self, "dialect", None)
+        new_builder._expression = intersect_expr
+        # Merge parameters
+        merged_params = dict(left_query.parameters)
+        merged_params.update(right_query.parameters)
+        new_builder._parameters = merged_params
+        return new_builder
+
+    def except_(self, other: Any) -> Self:
+        """Combine this query with another using EXCEPT.
+
+        Args:
+            other: Another SelectBuilder or compatible builder to except with.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The new builder instance for the except query.
+        """
+        left_query = self.build()  # type: ignore[attr-defined]
+        right_query = other.build()
+        left_expr: Optional[exp.Expression] = exp.maybe_parse(left_query.sql, dialect=getattr(self, "dialect", None))
+        right_expr: Optional[exp.Expression] = exp.maybe_parse(right_query.sql, dialect=getattr(self, "dialect", None))
+        if not left_expr or not right_expr:
+            msg = "Could not parse queries for EXCEPT operation"
+            raise SQLBuilderError(msg)
+        except_expr = exp.except_(left_expr, right_expr)
+        new_builder = type(self)()
+        new_builder.dialect = getattr(self, "dialect", None)
+        new_builder._expression = except_expr
+        # Merge parameters
+        merged_params = dict(left_query.parameters)
+        merged_params.update(right_query.parameters)
+        new_builder._parameters = merged_params
+        return new_builder

sqlspec/statement/builder/mixins/_unpivot.py

@@ -0,0 +1,80 @@
+from typing import TYPE_CHECKING, Optional, Union, cast
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlglot.dialects.dialect import DialectType
+
+    from sqlspec.statement.builder.select import SelectBuilder
+
+__all__ = ("UnpivotClauseMixin",)
+
+
+class UnpivotClauseMixin:
+    """Mixin class to add UNPIVOT functionality to a SelectBuilder."""
+
+    _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,
+    ) -> "SelectBuilder":
+        """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 SelectBuilder 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 SelectBuilder."
+            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)))
+
+        # Create the unpivot expression (stored as Pivot with unpivot=True)
+        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)))
+
+        # Add unpivot to the table in the FROM clause
+        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):
+                # Add to pivots array
+                existing_pivots = table.args.get("pivots", [])
+                existing_pivots.append(unpivot_node)
+                table.set("pivots", existing_pivots)
+
+        return cast("SelectBuilder", self)

sqlspec/statement/builder/mixins/_update_from.py

@@ -0,0 +1,54 @@
+from typing import Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("UpdateFromClauseMixin",)
+
+
+class UpdateFromClauseMixin:
+    """Mixin providing FROM clause for UPDATE builders (e.g., PostgreSQL style)."""
+
+    def from_(self, table: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
+        """Add a FROM clause to the UPDATE statement.
+
+        Args:
+            table: The table name, expression, or subquery to add to the FROM clause.
+            alias: Optional alias for the table in the FROM clause.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an UPDATE statement.
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Update):  # type: ignore[attr-defined]
+            msg = "Cannot add FROM clause to non-UPDATE expression. Set the main table first."
+            raise SQLBuilderError(msg)
+        table_expr: exp.Expression
+        if isinstance(table, str):
+            table_expr = exp.to_table(table, alias=alias)
+        elif hasattr(table, "build"):
+            subquery_builder_params = getattr(table, "_parameters", None)
+            if subquery_builder_params:
+                for p_name, p_value in subquery_builder_params.items():
+                    self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+            subquery_exp = exp.paren(getattr(table, "_expression", exp.select()))
+            table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
+        elif isinstance(table, exp.Expression):
+            table_expr = exp.alias_(table, alias) if alias else table
+        else:
+            msg = f"Unsupported table type for FROM clause: {type(table)}"
+            raise SQLBuilderError(msg)
+        if self._expression.args.get("from") is None:  # type: ignore[attr-defined]
+            self._expression.set("from", exp.From(expressions=[]))  # type: ignore[attr-defined]
+        from_clause = self._expression.args["from"]  # type: ignore[attr-defined]
+        if hasattr(from_clause, "append"):
+            from_clause.append("expressions", table_expr)
+        else:
+            if not from_clause.expressions:
+                from_clause.expressions = []
+            from_clause.expressions.append(table_expr)
+        return self

sqlspec/statement/builder/mixins/_update_set.py

@@ -0,0 +1,91 @@
+from collections.abc import Mapping
+from typing import Any, Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("UpdateSetClauseMixin",)
+
+MIN_SET_ARGS = 2
+
+
+class UpdateSetClauseMixin:
+    """Mixin providing SET clause for UPDATE builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def set(self, *args: Any, **kwargs: Any) -> Self:
+        """Set columns and values for the UPDATE statement.
+
+        Supports:
+        - set(column, value)
+        - set(mapping)
+        - set(**kwargs)
+        - set(mapping, **kwargs)
+
+        Args:
+            *args: Either (column, value) or a mapping.
+            **kwargs: Column-value pairs to set.
+
+        Raises:
+            SQLBuilderError: If the current expression is not an UPDATE statement or usage is invalid.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+
+        if self._expression is None:
+            self._expression = exp.Update()
+        if not isinstance(self._expression, exp.Update):
+            msg = "Cannot add SET clause to non-UPDATE expression."
+            raise SQLBuilderError(msg)
+        assignments = []
+        # (column, value) signature
+        if len(args) == MIN_SET_ARGS and not kwargs:
+            col, val = args
+            col_expr = col if isinstance(col, exp.Column) else exp.column(col)
+            # If value is an expression, use it directly
+            if isinstance(val, exp.Expression):
+                value_expr = val
+            elif hasattr(val, "_expression") and hasattr(val, "build"):
+                # It's a builder (like SelectBuilder), convert to subquery
+                subquery = val.build()
+                # Parse the SQL and use as expression
+                value_expr = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(self, "dialect", None)))
+                # Merge parameters from subquery
+                if hasattr(val, "_parameters"):
+                    for p_name, p_value in val._parameters.items():
+                        self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+            else:
+                param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
+                value_expr = exp.Placeholder(this=param_name)
+            assignments.append(exp.EQ(this=col_expr, expression=value_expr))
+        # mapping and/or kwargs
+        elif (len(args) == 1 and isinstance(args[0], Mapping)) or kwargs:
+            all_values = dict(args[0] if args else {}, **kwargs)
+            for col, val in all_values.items():
+                # If value is an expression, use it directly
+                if isinstance(val, exp.Expression):
+                    value_expr = val
+                elif hasattr(val, "_expression") and hasattr(val, "build"):
+                    # It's a builder (like SelectBuilder), convert to subquery
+                    subquery = val.build()
+                    # Parse the SQL and use as expression
+                    value_expr = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(self, "dialect", None)))
+                    # Merge parameters from subquery
+                    if hasattr(val, "_parameters"):
+                        for p_name, p_value in val._parameters.items():
+                            self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+                else:
+                    param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
+                    value_expr = exp.Placeholder(this=param_name)
+                assignments.append(exp.EQ(this=exp.column(col), expression=value_expr))
+        else:
+            msg = "Invalid arguments for set(): use (column, value), mapping, or kwargs."
+            raise SQLBuilderError(msg)
+        # Append to existing expressions instead of replacing
+        existing = self._expression.args.get("expressions", [])
+        self._expression.set("expressions", existing + assignments)
+        return self

sqlspec/statement/builder/mixins/_update_table.py

@@ -0,0 +1,29 @@
+from typing import Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+__all__ = ("UpdateTableClauseMixin",)
+
+
+class UpdateTableClauseMixin:
+    """Mixin providing TABLE clause for UPDATE builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def table(self, table_name: str, alias: Optional[str] = None) -> Self:
+        """Set the table to update.
+
+        Args:
+            table_name: The name of the table.
+            alias: Optional alias for the table.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Update):
+            self._expression = exp.Update(this=None, expressions=[], joins=[])
+        table_expr: exp.Expression = exp.to_table(table_name, alias=alias)
+        self._expression.set("this", table_expr)
+        setattr(self, "_table", table_name)
+        return self