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

sqlspec/statement/builder/insert.py

@@ -25,7 +25,7 @@ if TYPE_CHECKING:
     from collections.abc import Mapping, Sequence
 
 
-__all__ = ("InsertBuilder",)
+__all__ = ("Insert",)
 
 ERR_MSG_TABLE_NOT_SET = "The target table must be set using .into() before adding values."
 ERR_MSG_VALUES_COLUMNS_MISMATCH = (
@@ -36,9 +36,7 @@ ERR_MSG_EXPRESSION_NOT_INITIALIZED = "Internal error: base expression not initia
 
 
 @dataclass(unsafe_hash=True)
-class InsertBuilder(
-    QueryBuilder[RowT], ReturningClauseMixin, InsertValuesMixin, InsertFromSelectMixin, InsertIntoClauseMixin
-):
+class Insert(QueryBuilder[RowT], ReturningClauseMixin, InsertValuesMixin, InsertFromSelectMixin, InsertIntoClauseMixin):
     """Builder for INSERT statements.
 
     This builder facilitates the construction of SQL INSERT queries
@@ -48,15 +46,20 @@ class InsertBuilder(
         ```python
         # Basic INSERT with values
         insert_query = (
-            InsertBuilder()
+            Insert()
             .into("users")
             .columns("name", "email", "age")
             .values("John Doe", "john@example.com", 30)
         )
 
+        # Even more concise with constructor
+        insert_query = Insert("users").values(
+            {"name": "John", "age": 30}
+        )
+
         # Multi-row INSERT
         insert_query = (
-            InsertBuilder()
+            Insert()
             .into("users")
             .columns("name", "email")
             .values("John", "john@example.com")
@@ -65,7 +68,7 @@ class InsertBuilder(
 
         # INSERT from dictionary
         insert_query = (
-            InsertBuilder()
+            Insert()
             .into("users")
             .values_from_dict(
                 {"name": "John", "email": "john@example.com"}
@@ -74,10 +77,10 @@ class InsertBuilder(
 
         # INSERT from SELECT
         insert_query = (
-            InsertBuilder()
+            Insert()
             .into("users_backup")
             .from_select(
-                SelectBuilder()
+                Select()
                 .select("name", "email")
                 .from_("users")
                 .where("active = true")
@@ -90,6 +93,23 @@ class InsertBuilder(
     _columns: list[str] = field(default_factory=list, init=False)
     _values_added_count: int = field(default=0, init=False)
 
+    def __init__(self, table: Optional[str] = None, **kwargs: Any) -> None:
+        """Initialize INSERT with optional table.
+
+        Args:
+            table: Target table name
+            **kwargs: Additional QueryBuilder arguments
+        """
+        super().__init__(**kwargs)
+
+        # Initialize fields from dataclass
+        self._table = None
+        self._columns = []
+        self._values_added_count = 0
+
+        if table:
+            self.into(table)
+
     def _create_base_expression(self) -> exp.Insert:
         """Create a base INSERT expression.
 
@@ -165,7 +185,6 @@ class InsertBuilder(
         else:
             # This case should ideally not be reached if logic is correct:
             # means _values_added_count > 0 but expression is not exp.Values.
-            # Fallback to creating a new Values node, though this might indicate an issue.
             new_values_node = exp.Values(expressions=[exp.Tuple(expressions=list(value_placeholders))])
             insert_expr.set("expression", new_values_node)
 
@@ -191,14 +210,12 @@ class InsertBuilder(
             raise SQLBuilderError(ERR_MSG_TABLE_NOT_SET)
 
         if not self._columns:
-            # Set columns from dictionary keys if not already set
             self.columns(*data.keys())
         elif set(self._columns) != set(data.keys()):
             # Verify that dictionary keys match existing columns
             msg = f"Dictionary keys {set(data.keys())} do not match existing columns {set(self._columns)}."
             raise SQLBuilderError(msg)
 
-        # Add values in the same order as columns
         return self.values(*[data[col] for col in self._columns])
 
     def values_from_dicts(self, data: "Sequence[Mapping[str, Any]]") -> "Self":
@@ -219,12 +236,10 @@ class InsertBuilder(
         if not data:
             return self
 
-        # Use the first dictionary to establish columns
         first_dict = data[0]
         if not self._columns:
             self.columns(*first_dict.keys())
 
-        # Validate that all dictionaries have the same keys
         expected_keys = set(self._columns)
         for i, row_dict in enumerate(data):
             if set(row_dict.keys()) != expected_keys:
@@ -234,7 +249,6 @@ class InsertBuilder(
                 )
                 raise SQLBuilderError(msg)
 
-        # Add each row
         for row_dict in data:
             self.values(*[row_dict[col] for col in self._columns])
 

sqlspec/statement/builder/merge.py

@@ -20,11 +20,11 @@ from sqlspec.statement.builder.mixins import (
 from sqlspec.statement.result import SQLResult
 from sqlspec.typing import RowT
 
-__all__ = ("MergeBuilder",)
+__all__ = ("Merge",)
 
 
 @dataclass(unsafe_hash=True)
-class MergeBuilder(
+class Merge(
     QueryBuilder[RowT],
     MergeUsingClauseMixin,
     MergeOnClauseMixin,
@@ -42,7 +42,7 @@ class MergeBuilder(
         ```python
         # Basic MERGE statement
         merge_query = (
-            MergeBuilder()
+            Merge()
             .into("target_table")
             .using("source_table", "src")
             .on("target_table.id = src.id")
@@ -64,7 +64,7 @@ class MergeBuilder(
         )
 
         merge_query = (
-            MergeBuilder()
+            Merge()
             .into("users")
             .using(source_query, "src")
             .on("users.email = src.email")

sqlspec/statement/builder/mixins/_aggregate_functions.py

@@ -4,7 +4,7 @@ from sqlglot import exp
 from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from sqlspec.statement.builder.protocols import SelectBuilderProtocol
+    from sqlspec.protocols import SelectBuilderProtocol
 
 __all__ = ("AggregateFunctionsMixin",)
 
@@ -112,40 +112,139 @@ class AggregateFunctionsMixin:
         select_expr = exp.alias_(array_agg_expr, alias) if alias else array_agg_expr
         return cast("Self", builder.select(select_expr))
 
-    def bool_and(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
-        """Add BOOL_AND aggregate function to SELECT clause (PostgreSQL, DuckDB, etc).
+    def count_distinct(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add COUNT(DISTINCT column) to SELECT clause.
 
         Args:
-            column: The boolean column to aggregate.
+            column: The column to count distinct values of.
+            alias: Optional alias for the count.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        count_expr = exp.Count(this=exp.Distinct(expressions=[col_expr]))
+        select_expr = exp.alias_(count_expr, alias) if alias else count_expr
+        return cast("Self", builder.select(select_expr))
+
+    def stddev(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add STDDEV aggregate function to SELECT clause.
+
+        Args:
+            column: The column to calculate standard deviation of.
             alias: Optional alias for the result.
 
         Returns:
             The current builder instance for method chaining.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        stddev_expr = exp.Stddev(this=col_expr)
+        select_expr = exp.alias_(stddev_expr, alias) if alias else stddev_expr
+        return cast("Self", builder.select(select_expr))
 
-        Note:
-            Uses exp.Anonymous for BOOL_AND. Not all dialects support this function.
+    def stddev_pop(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add STDDEV_POP aggregate function to SELECT clause.
+
+        Args:
+            column: The column to calculate population standard deviation of.
+            alias: Optional alias for the result.
+
+        Returns:
+            The current builder instance for method chaining.
         """
         builder = cast("SelectBuilderProtocol", self)
         col_expr = exp.column(column) if isinstance(column, str) else column
-        bool_and_expr = exp.Anonymous(this="BOOL_AND", expressions=[col_expr])
-        select_expr = exp.alias_(bool_and_expr, alias) if alias else bool_and_expr
+        stddev_pop_expr = exp.StddevPop(this=col_expr)
+        select_expr = exp.alias_(stddev_pop_expr, alias) if alias else stddev_pop_expr
         return cast("Self", builder.select(select_expr))
 
-    def bool_or(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
-        """Add BOOL_OR aggregate function to SELECT clause (PostgreSQL, DuckDB, etc).
+    def stddev_samp(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add STDDEV_SAMP aggregate function to SELECT clause.
 
         Args:
-            column: The boolean column to aggregate.
+            column: The column to calculate sample standard deviation of.
+            alias: Optional alias for the result.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        stddev_samp_expr = exp.StddevSamp(this=col_expr)
+        select_expr = exp.alias_(stddev_samp_expr, alias) if alias else stddev_samp_expr
+        return cast("Self", builder.select(select_expr))
+
+    def variance(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add VARIANCE aggregate function to SELECT clause.
+
+        Args:
+            column: The column to calculate variance of.
+            alias: Optional alias for the result.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        variance_expr = exp.Variance(this=col_expr)
+        select_expr = exp.alias_(variance_expr, alias) if alias else variance_expr
+        return cast("Self", builder.select(select_expr))
+
+    def var_pop(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add VAR_POP aggregate function to SELECT clause.
+
+        Args:
+            column: The column to calculate population variance of.
+            alias: Optional alias for the result.
+
+        Returns:
+            The current builder instance for method chaining.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        var_pop_expr = exp.VariancePop(this=col_expr)
+        select_expr = exp.alias_(var_pop_expr, alias) if alias else var_pop_expr
+        return cast("Self", builder.select(select_expr))
+
+    def string_agg(self, column: Union[str, exp.Expression], separator: str = ",", alias: Optional[str] = None) -> Self:
+        """Add STRING_AGG aggregate function to SELECT clause.
+
+        Args:
+            column: The column to aggregate into a string.
+            separator: The separator between values (default is comma).
             alias: Optional alias for the result.
 
         Returns:
             The current builder instance for method chaining.
 
         Note:
-            Uses exp.Anonymous for BOOL_OR. Not all dialects support this function.
+            Different databases have different names for this function:
+            - PostgreSQL: STRING_AGG
+            - MySQL: GROUP_CONCAT
+            - SQLite: GROUP_CONCAT
+            SQLGlot will handle the translation.
+        """
+        builder = cast("SelectBuilderProtocol", self)
+        col_expr = exp.column(column) if isinstance(column, str) else column
+        # Use GroupConcat which SQLGlot can translate to STRING_AGG for Postgres
+        string_agg_expr = exp.GroupConcat(this=col_expr, separator=exp.Literal.string(separator))
+        select_expr = exp.alias_(string_agg_expr, alias) if alias else string_agg_expr
+        return cast("Self", builder.select(select_expr))
+
+    def json_agg(self, column: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
+        """Add JSON_AGG aggregate function to SELECT clause.
+
+        Args:
+            column: The column to aggregate into a JSON array.
+            alias: Optional alias for the result.
+
+        Returns:
+            The current builder instance for method chaining.
         """
         builder = cast("SelectBuilderProtocol", self)
         col_expr = exp.column(column) if isinstance(column, str) else column
-        bool_or_expr = exp.Anonymous(this="BOOL_OR", expressions=[col_expr])
-        select_expr = exp.alias_(bool_or_expr, alias) if alias else bool_or_expr
+        json_agg_expr = exp.JSONArrayAgg(this=col_expr)
+        select_expr = exp.alias_(json_agg_expr, alias) if alias else json_agg_expr
         return cast("Self", builder.select(select_expr))

sqlspec/statement/builder/mixins/_common_table_expr.py

@@ -61,7 +61,6 @@ class CommonTableExpressionMixin:
             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])

sqlspec/statement/builder/mixins/_delete_from.py

@@ -26,7 +26,7 @@ class DeleteFromClauseMixin:
             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."
+            msg = f"Base expression for Delete is {current_expr_type}, expected Delete."
             raise SQLBuilderError(msg)
 
         setattr(self, "_table", table)

sqlspec/statement/builder/mixins/_from.py

@@ -5,10 +5,10 @@ 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
+from sqlspec.utils.type_guards import has_query_builder_parameters, is_expression
 
 if TYPE_CHECKING:
-    from sqlspec.statement.builder.protocols import BuilderProtocol
+    from sqlspec.protocols import SQLBuilderProtocol
 
 __all__ = ("FromClauseMixin",)
 
@@ -29,7 +29,7 @@ class FromClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if builder._expression is None:
             builder._expression = exp.Select()
         if not isinstance(builder._expression, exp.Select):
@@ -41,16 +41,18 @@ class FromClauseMixin:
         elif is_expression(table):
             # Direct sqlglot expression - use as is
             from_expr = exp.alias_(table, alias) if alias else table
-        elif hasattr(table, "build"):
+        elif has_query_builder_parameters(table):
             # Query builder with build() method
-            subquery = table.build()  # pyright: ignore
-            subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(builder, "dialect", None)))
+            subquery = table.build()
+            sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
+            subquery_exp = exp.paren(exp.maybe_parse(sql_str, 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:
+            if merged_params and hasattr(subquery, "parameters"):
+                subquery_params = getattr(subquery, "parameters", {})
                 merged_params = merged_params.merge_parameters(
-                    parameters=subquery.parameters,
+                    parameters=subquery_params,
                     args=current_params if isinstance(current_params, list) else None,
                     kwargs=current_params if isinstance(current_params, dict) else {},
                 )

sqlspec/statement/builder/mixins/_group_by.py

@@ -106,7 +106,6 @@ class GroupByClauseMixin:
         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]

sqlspec/statement/builder/mixins/_insert_from_select.py

@@ -38,7 +38,6 @@ class InsertFromSelectMixin:
         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())

sqlspec/statement/builder/mixins/_insert_values.py

@@ -39,7 +39,6 @@ class InsertValuesMixin:
         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
         ):
@@ -50,7 +49,6 @@ class InsertValuesMixin:
             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])

sqlspec/statement/builder/mixins/_join.py

@@ -5,9 +5,10 @@ from typing_extensions import Self
 
 from sqlspec.exceptions import SQLBuilderError
 from sqlspec.statement.builder._parsing_utils import parse_table_expression
+from sqlspec.utils.type_guards import has_query_builder_parameters
 
 if TYPE_CHECKING:
-    from sqlspec.statement.builder.protocols import BuilderProtocol
+    from sqlspec.protocols import SQLBuilderProtocol
 
 __all__ = ("JoinClauseMixin",)
 
@@ -22,7 +23,7 @@ class JoinClauseMixin:
         alias: Optional[str] = None,
         join_type: str = "INNER",
     ) -> Self:
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if builder._expression is None:
             builder._expression = exp.Select()
         if not isinstance(builder._expression, exp.Select):
@@ -31,16 +32,19 @@ class JoinClauseMixin:
         table_expr: exp.Expression
         if isinstance(table, str):
             table_expr = parse_table_expression(table, alias)
-        elif hasattr(table, "build"):
-            # Handle builder objects with build() method
+        elif has_query_builder_parameters(table):
             # Work directly with AST when possible to avoid string parsing
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
-                subquery_exp = exp.paren(table._expression.copy())  # pyright: ignore
+                table_expr_value = getattr(table, "_expression", None)
+                if table_expr_value is not None:
+                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                else:
+                    subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                # Fallback to string parsing
                 subquery = table.build()  # pyright: ignore
-                subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(builder, "dialect", None)))
+                sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
+                subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(builder, "dialect", None)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             # Parameter merging logic can be added here if needed
         else:
@@ -84,7 +88,7 @@ class JoinClauseMixin:
         return self.join(table, on, alias, "FULL")
 
     def cross_join(self, table: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if builder._expression is None:
             builder._expression = exp.Select()
         if not isinstance(builder._expression, exp.Select):
@@ -93,15 +97,18 @@ class JoinClauseMixin:
         table_expr: exp.Expression
         if isinstance(table, str):
             table_expr = parse_table_expression(table, alias)
-        elif hasattr(table, "build"):
-            # Handle builder objects with build() method
+        elif has_query_builder_parameters(table):
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
-                subquery_exp = exp.paren(table._expression.copy())  # pyright: ignore
+                table_expr_value = getattr(table, "_expression", None)
+                if table_expr_value is not None:
+                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                else:
+                    subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                # Fallback to string parsing
                 subquery = table.build()  # pyright: ignore
-                subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=getattr(builder, "dialect", None)))
+                sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
+                subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(builder, "dialect", None)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
         else:
             table_expr = table

sqlspec/statement/builder/mixins/_limit_offset.py

@@ -4,7 +4,7 @@ from sqlglot import exp
 from typing_extensions import Self
 
 if TYPE_CHECKING:
-    from sqlspec.statement.builder.protocols import BuilderProtocol
+    from sqlspec.protocols import SQLBuilderProtocol
 
 from sqlspec.exceptions import SQLBuilderError
 
@@ -26,7 +26,7 @@ class LimitOffsetClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if not isinstance(builder._expression, exp.Select):
             msg = "LIMIT is only supported for SELECT statements."
             raise SQLBuilderError(msg)
@@ -45,7 +45,7 @@ class LimitOffsetClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if not isinstance(builder._expression, exp.Select):
             msg = "OFFSET is only supported for SELECT statements."
             raise SQLBuilderError(msg)

sqlspec/statement/builder/mixins/_merge_clauses.py

@@ -4,6 +4,7 @@ from sqlglot import exp
 from typing_extensions import Self
 
 from sqlspec.exceptions import SQLBuilderError
+from sqlspec.utils.type_guards import has_query_builder_parameters
 
 __all__ = (
     "MergeIntoClauseMixin",
@@ -66,9 +67,9 @@ class MergeUsingClauseMixin:
         source_expr: exp.Expression
         if isinstance(source, str):
             source_expr = exp.to_table(source, alias=alias)
-        elif hasattr(source, "_parameters") and hasattr(source, "_expression"):
+        elif has_query_builder_parameters(source) and hasattr(source, "_expression"):
             # Merge parameters from the SELECT builder or other builder
-            subquery_builder_params = getattr(source, "_parameters", {})
+            subquery_builder_params = source.parameters
             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]
@@ -145,13 +146,11 @@ class MergeMatchedClauseMixin:
         if not isinstance(self._expression, exp.Merge):
             self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
 
-        # Get or create the whens object
         whens = self._expression.args.get("whens")
         if not whens:
             whens = exp.Whens(expressions=[])
             self._expression.set("whens", whens)
 
-        # Add the when clause to the whens expressions using SQLGlot's append method
         whens.append("expressions", when_clause)
 
     def when_matched_then_update(

sqlspec/statement/builder/mixins/_order_by.py

@@ -7,7 +7,7 @@ from sqlspec.exceptions import SQLBuilderError
 from sqlspec.statement.builder._parsing_utils import parse_order_expression
 
 if TYPE_CHECKING:
-    from sqlspec.statement.builder.protocols import BuilderProtocol
+    from sqlspec.protocols import SQLBuilderProtocol
 
 __all__ = ("OrderByClauseMixin",)
 
@@ -28,7 +28,7 @@ class OrderByClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        builder = cast("BuilderProtocol", self)
+        builder = cast("SQLBuilderProtocol", self)
         if not isinstance(builder._expression, exp.Select):
             msg = "ORDER BY is only supported for SELECT statements."
             raise SQLBuilderError(msg)

sqlspec/statement/builder/mixins/_pivot.py

@@ -5,13 +5,13 @@ from sqlglot import exp
 if TYPE_CHECKING:
     from sqlglot.dialects.dialect import DialectType
 
-    from sqlspec.statement.builder.select import SelectBuilder
+    from sqlspec.statement.builder.select import Select
 
 __all__ = ("PivotClauseMixin",)
 
 
 class PivotClauseMixin:
-    """Mixin class to add PIVOT functionality to a SelectBuilder."""
+    """Mixin class to add PIVOT functionality to a Select."""
 
     _expression: "Optional[exp.Expression]" = None
     dialect: "DialectType" = None
@@ -23,7 +23,7 @@ class PivotClauseMixin:
         pivot_column: Union[str, exp.Expression],
         pivot_values: list[Union[str, int, float, exp.Expression]],
         alias: Optional[str] = None,
-    ) -> "SelectBuilder":
+    ) -> "Select":
         """Adds a PIVOT clause to the SELECT statement.
 
         Example:
@@ -61,7 +61,6 @@ class PivotClauseMixin:
             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)
@@ -69,14 +68,12 @@ class PivotClauseMixin:
         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)
+        return cast("Select", self)