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

sqlspec/statement/builder/mixins/_case_builder.py

@@ -0,0 +1,91 @@
+# mypy: disable-error-code="valid-type,type-var"
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any, Optional, Union, cast
+
+from sqlglot import exp
+
+if TYPE_CHECKING:
+    from sqlspec.statement.builder.base import QueryBuilder
+    from sqlspec.typing import RowT
+
+__all__ = ("CaseBuilder", "CaseBuilderMixin")
+
+
+class CaseBuilderMixin:
+    """Mixin providing CASE expression functionality for SQL builders."""
+
+    def case_(self, alias: "Optional[str]" = None) -> "CaseBuilder":
+        """Create a CASE expression for the SELECT clause.
+
+        Args:
+            alias: Optional alias for the CASE expression.
+
+        Returns:
+            CaseBuilder: A CaseBuilder instance for building the CASE expression.
+        """
+        builder = cast("QueryBuilder[RowT]", self)  # pyright: ignore
+        return CaseBuilder(builder, alias)
+
+
+@dataclass
+class CaseBuilder:
+    """Builder for CASE expressions."""
+
+    _parent: "QueryBuilder[RowT]"  # pyright: ignore
+    _alias: Optional[str]
+    _case_expr: exp.Case
+
+    def __init__(self, parent: "QueryBuilder[RowT]", alias: "Optional[str]" = None) -> None:
+        """Initialize CaseBuilder.
+
+        Args:
+            parent: The parent builder.
+            alias: Optional alias for the CASE expression.
+        """
+        self._parent = parent
+        self._alias = alias
+        self._case_expr = exp.Case()
+
+    def when(self, condition: "Union[str, exp.Expression]", value: "Any") -> "CaseBuilder":
+        """Add WHEN clause to CASE expression.
+
+        Args:
+            condition: The condition to test.
+            value: The value to return if condition is true.
+
+        Returns:
+            CaseBuilder: The current builder instance for method chaining.
+        """
+        cond_expr = exp.condition(condition) if isinstance(condition, str) else condition
+        param_name = self._parent.add_parameter(value)[1]
+        value_expr = exp.Placeholder(this=param_name)
+
+        when_clause = exp.When(this=cond_expr, then=value_expr)
+
+        if not self._case_expr.args.get("ifs"):
+            self._case_expr.set("ifs", [])
+        self._case_expr.args["ifs"].append(when_clause)
+        return self
+
+    def else_(self, value: "Any") -> "CaseBuilder":
+        """Add ELSE clause to CASE expression.
+
+        Args:
+            value: The value to return if no conditions match.
+
+        Returns:
+            CaseBuilder: The current builder instance for method chaining.
+        """
+        param_name = self._parent.add_parameter(value)[1]
+        value_expr = exp.Placeholder(this=param_name)
+        self._case_expr.set("default", value_expr)
+        return self
+
+    def end(self) -> "QueryBuilder[RowT]":
+        """Finalize the CASE expression and add it to the SELECT clause.
+
+        Returns:
+            The parent builder instance.
+        """
+        select_expr = exp.alias_(self._case_expr, self._alias) if self._alias else self._case_expr
+        return cast("QueryBuilder[RowT]", self._parent.select(select_expr))  # type: ignore[attr-defined]

sqlspec/statement/builder/mixins/_common_table_expr.py

@@ -0,0 +1,91 @@
+from typing import Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("CommonTableExpressionMixin",)
+
+
+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)
+
+        # Create a proper CTE with table alias
+        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

sqlspec/statement/builder/mixins/_delete_from.py

@@ -0,0 +1,34 @@
+from typing import Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("DeleteFromClauseMixin",)
+
+
+class DeleteFromClauseMixin:
+    """Mixin providing FROM clause for DELETE builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def from_(self, table: str) -> Self:
+        """Set the target table for the DELETE statement.
+
+        Args:
+            table: The table name to delete from.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Delete()
+        if not isinstance(self._expression, exp.Delete):
+            current_expr_type = type(self._expression).__name__
+            msg = f"Base expression for DeleteBuilder is {current_expr_type}, expected Delete."
+            raise SQLBuilderError(msg)
+
+        setattr(self, "_table", table)
+        self._expression.set("this", exp.to_table(table))
+        return self

sqlspec/statement/builder/mixins/_from.py

@@ -0,0 +1,61 @@
+from typing import TYPE_CHECKING, Any, 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_table_expression
+from sqlspec.typing import is_expression
+
+if TYPE_CHECKING:
+    from sqlspec.statement.builder.protocols import BuilderProtocol
+
+__all__ = ("FromClauseMixin",)
+
+
+class FromClauseMixin:
+    """Mixin providing FROM clause for SELECT builders."""
+
+    def from_(self, table: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
+        """Add FROM clause.
+
+        Args:
+            table: The table name, expression, or subquery to select from.
+            alias: Optional alias for the table.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement or if the table type is unsupported.
+
+        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 = "FROM clause is only supported for SELECT statements."
+            raise SQLBuilderError(msg)
+        from_expr: exp.Expression
+        if isinstance(table, str):
+            from_expr = parse_table_expression(table, alias)
+        elif is_expression(table):
+            # Direct sqlglot expression - use as is
+            from_expr = exp.alias_(table, alias) if alias else table
+        elif hasattr(table, "build"):
+            # Query builder with build() method
+            subquery = table.build()  # pyright: ignore
+            subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(builder, "dialect", None)))
+            from_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
+            current_params = getattr(builder, "_parameters", None)
+            merged_params = getattr(type(builder), "ParameterConverter", None)
+            if merged_params:
+                merged_params = merged_params.merge_parameters(
+                    parameters=subquery.parameters,
+                    args=current_params if isinstance(current_params, list) else None,
+                    kwargs=current_params if isinstance(current_params, dict) else {},
+                )
+                setattr(builder, "_parameters", merged_params)
+        else:
+            from_expr = table
+        builder._expression = builder._expression.from_(from_expr, copy=False)
+        return cast("Self", builder)

sqlspec/statement/builder/mixins/_group_by.py

@@ -0,0 +1,119 @@
+from typing import Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+__all__ = ("GroupByClauseMixin",)
+
+
+class GroupByClauseMixin:
+    """Mixin providing GROUP BY clause functionality for SQL builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def group_by(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add GROUP BY clause.
+
+        Args:
+            *columns: Columns to group by. Can be column names, expressions,
+                     or special grouping expressions like ROLLUP, CUBE, etc.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Select):
+            return self
+
+        for column in columns:
+            self._expression = self._expression.group_by(
+                exp.column(column) if isinstance(column, str) else column, copy=False
+            )
+        return self
+
+    def group_by_rollup(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add GROUP BY ROLLUP clause.
+
+        ROLLUP generates subtotals and grand totals for a hierarchical set of columns.
+
+        Args:
+            *columns: Columns to include in the rollup hierarchy.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Example:
+            ```python
+            # GROUP BY ROLLUP(product, region)
+            query = (
+                sql.select("product", "region", sql.sum("sales"))
+                .from_("sales_data")
+                .group_by_rollup("product", "region")
+            )
+            ```
+        """
+        column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns]
+        rollup_expr = exp.Rollup(expressions=column_exprs)
+        return self.group_by(rollup_expr)
+
+    def group_by_cube(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Add GROUP BY CUBE clause.
+
+        CUBE generates subtotals for all possible combinations of the specified columns.
+
+        Args:
+            *columns: Columns to include in the cube.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Example:
+            ```python
+            # GROUP BY CUBE(product, region)
+            query = (
+                sql.select("product", "region", sql.sum("sales"))
+                .from_("sales_data")
+                .group_by_cube("product", "region")
+            )
+            ```
+        """
+        column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns]
+        cube_expr = exp.Cube(expressions=column_exprs)
+        return self.group_by(cube_expr)
+
+    def group_by_grouping_sets(self, *column_sets: Union[tuple[str, ...], list[str]]) -> Self:
+        """Add GROUP BY GROUPING SETS clause.
+
+        GROUPING SETS allows you to specify multiple grouping sets in a single query.
+
+        Args:
+            *column_sets: Sets of columns to group by. Each set can be a tuple or list.
+                         Empty tuple/list creates a grand total grouping.
+
+        Returns:
+            The current builder instance for method chaining.
+
+        Example:
+            ```python
+            # GROUP BY GROUPING SETS ((product), (region), ())
+            query = (
+                sql.select("product", "region", sql.sum("sales"))
+                .from_("sales_data")
+                .group_by_grouping_sets(("product",), ("region",), ())
+            )
+            ```
+        """
+        set_expressions = []
+        for column_set in column_sets:
+            if isinstance(column_set, (tuple, list)):
+                if len(column_set) == 0:
+                    # Empty set for grand total
+                    set_expressions.append(exp.Tuple(expressions=[]))
+                else:
+                    columns = [exp.column(col) for col in column_set]
+                    set_expressions.append(exp.Tuple(expressions=columns))
+            else:
+                # Single column
+                set_expressions.append(exp.column(column_set))
+
+        grouping_sets_expr = exp.GroupingSets(expressions=set_expressions)
+        return self.group_by(grouping_sets_expr)

sqlspec/statement/builder/mixins/_having.py

@@ -0,0 +1,35 @@
+from typing import Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("HavingClauseMixin",)
+
+
+class HavingClauseMixin:
+    """Mixin providing HAVING clause for SELECT builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def having(self, condition: Union[str, exp.Expression]) -> Self:
+        """Add HAVING clause.
+
+        Args:
+            condition: The condition for the HAVING clause.
+
+        Raises:
+            SQLBuilderError: If the current expression is not a SELECT statement.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        if self._expression is None:
+            self._expression = exp.Select()
+        if not isinstance(self._expression, exp.Select):
+            msg = "Cannot add HAVING to a non-SELECT expression."
+            raise SQLBuilderError(msg)
+        having_expr = exp.condition(condition) if isinstance(condition, str) else condition
+        self._expression = self._expression.having(having_expr, copy=False)
+        return self

sqlspec/statement/builder/mixins/_insert_from_select.py

@@ -0,0 +1,48 @@
+from typing import Any, Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("InsertFromSelectMixin",)
+
+
+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]
+        # Set the SELECT expression as the source
+        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/_insert_into.py

@@ -0,0 +1,36 @@
+from typing import Optional
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("InsertIntoClauseMixin",)
+
+
+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

sqlspec/statement/builder/mixins/_insert_values.py

@@ -0,0 +1,69 @@
+from collections.abc import Sequence
+from typing import Any, Optional, Union
+
+from sqlglot import exp
+from typing_extensions import Self
+
+from sqlspec.exceptions import SQLBuilderError
+
+__all__ = ("InsertValuesMixin",)
+
+
+class InsertValuesMixin:
+    """Mixin providing VALUES and columns methods for INSERT builders."""
+
+    _expression: Optional[exp.Expression] = None
+
+    def columns(self, *columns: Union[str, exp.Expression]) -> Self:
+        """Set the columns for the INSERT statement and synchronize the _columns attribute on the builder."""
+        if self._expression is None:
+            self._expression = exp.Insert()
+        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)
+        # Synchronize the _columns attribute on the builder (if present)
+        if hasattr(self, "_columns"):
+            # If no columns, clear the list
+            if not columns:
+                self._columns.clear()  # pyright: ignore
+            else:
+                self._columns[:] = [col.name if isinstance(col, exp.Column) else str(col) for col in columns]  # pyright: ignore
+        return self
+
+    def values(self, *values: Any) -> Self:
+        """Add a row of values to the INSERT statement, validating against _columns if set."""
+        if self._expression is None:
+            self._expression = exp.Insert()
+        if not isinstance(self._expression, exp.Insert):
+            msg = "Cannot add values to a non-INSERT expression."
+            raise SQLBuilderError(msg)
+        # Validate value count if _columns is present and non-empty
+        if (
+            hasattr(self, "_columns") and getattr(self, "_columns", []) and len(values) != len(self._columns)  # pyright: ignore
+        ):
+            msg = f"Number of values ({len(values)}) does not match the number of specified columns ({len(self._columns)})."  # pyright: ignore
+            raise SQLBuilderError(msg)
+        row_exprs = []
+        for v in values:
+            if isinstance(v, exp.Expression):
+                row_exprs.append(v)
+            else:
+                # Add as parameter
+                _, param_name = self.add_parameter(v)  # type: ignore[attr-defined]
+                row_exprs.append(exp.var(param_name))
+        values_expr = exp.Values(expressions=[row_exprs])
+        self._expression.set("expression", values_expr)
+        return self
+
+    def add_values(self, values: Sequence[Any]) -> Self:
+        """Add a row of values to the INSERT statement (alternative signature).
+
+        Args:
+            values: Sequence of values for the row.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        return self.values(*values)