sqlspec 0.16.0__py3-none-any.whl → 0.16.1__py3-none-any.whl

sqlspec/builder/mixins/_insert_operations.py

@@ -1,20 +1,28 @@
 """Insert operation mixins for SQL builders."""
 
 from collections.abc import Sequence
-from typing import Any, Optional, Union
+from typing import Any, Optional, TypeVar, Union
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
 from sqlspec.exceptions import SQLBuilderError
+from sqlspec.protocols import SQLBuilderProtocol
+
+BuilderT = TypeVar("BuilderT", bound=SQLBuilderProtocol)
 
 __all__ = ("InsertFromSelectMixin", "InsertIntoClauseMixin", "InsertValuesMixin")
 
 
+@trait
 class InsertIntoClauseMixin:
     """Mixin providing INTO clause for INSERT builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def into(self, table: str) -> Self:
         """Set the target table for the INSERT statement.
@@ -39,10 +47,26 @@ class InsertIntoClauseMixin:
         return self
 
 
+@trait
 class InsertValuesMixin:
     """Mixin providing VALUES and columns methods for INSERT builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _columns: Any  # Provided by QueryBuilder
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate unique parameter name - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def columns(self, *columns: Union[str, exp.Expression]) -> Self:
         """Set the columns for the INSERT statement and synchronize the _columns attribute on the builder."""
@@ -54,7 +78,7 @@ class InsertValuesMixin:
         column_exprs = [exp.column(col) if isinstance(col, str) else col for col in columns]
         self._expression.set("columns", column_exprs)
         try:
-            cols = self._columns  # type: ignore[attr-defined]
+            cols = self._columns
             if not columns:
                 cols.clear()
             else:
@@ -63,37 +87,94 @@ class InsertValuesMixin:
             pass
         return self
 
-    def values(self, *values: Any) -> Self:
-        """Add a row of values to the INSERT statement, validating against _columns if set."""
+    def values(self, *values: Any, **kwargs: Any) -> Self:
+        """Add a row of values to the INSERT statement.
+
+        Supports:
+        - values(val1, val2, val3)
+        - values(col1=val1, col2=val2)
+        - values(mapping)
+
+        Args:
+            *values: Either positional values or a single mapping.
+            **kwargs: Column-value pairs.
+
+        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 add values to a non-INSERT expression."
             raise SQLBuilderError(msg)
-        try:
-            _columns = self._columns  # type: ignore[attr-defined]
-            if _columns and len(values) != len(_columns):
-                msg = f"Number of values ({len(values)}) does not match the number of specified columns ({len(_columns)})."
+
+        if kwargs:
+            if values:
+                msg = "Cannot mix positional values with keyword values."
                 raise SQLBuilderError(msg)
-        except AttributeError:
-            pass
-        row_exprs = []
-        for i, v in enumerate(values):
-            if isinstance(v, exp.Expression):
-                row_exprs.append(v)
-            else:
-                # Try to use column name if available, otherwise use position-based name
-                try:
-                    _columns = self._columns  # type: ignore[attr-defined]
-                    if _columns and i < len(_columns):
-                        column_name = str(_columns[i]).split(".")[-1] if "." in str(_columns[i]) else str(_columns[i])
-                        param_name = self._generate_unique_parameter_name(column_name)  # type: ignore[attr-defined]
-                    else:
-                        param_name = self._generate_unique_parameter_name(f"value_{i + 1}")  # type: ignore[attr-defined]
-                except AttributeError:
-                    param_name = self._generate_unique_parameter_name(f"value_{i + 1}")  # type: ignore[attr-defined]
-                _, param_name = self.add_parameter(v, name=param_name)  # type: ignore[attr-defined]
-                row_exprs.append(exp.var(param_name))
+            try:
+                _columns = self._columns
+                if not _columns:
+                    self.columns(*kwargs.keys())
+            except AttributeError:
+                pass
+            row_exprs = []
+            for col, val in kwargs.items():
+                if isinstance(val, exp.Expression):
+                    row_exprs.append(val)
+                else:
+                    column_name = col if isinstance(col, str) else str(col)
+                    if "." in column_name:
+                        column_name = column_name.split(".")[-1]
+                    param_name = self._generate_unique_parameter_name(column_name)
+                    _, param_name = self.add_parameter(val, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+        elif len(values) == 1 and hasattr(values[0], "items"):
+            mapping = values[0]
+            try:
+                _columns = self._columns
+                if not _columns:
+                    self.columns(*mapping.keys())
+            except AttributeError:
+                pass
+            row_exprs = []
+            for col, val in mapping.items():
+                if isinstance(val, exp.Expression):
+                    row_exprs.append(val)
+                else:
+                    column_name = col if isinstance(col, str) else str(col)
+                    if "." in column_name:
+                        column_name = column_name.split(".")[-1]
+                    param_name = self._generate_unique_parameter_name(column_name)
+                    _, param_name = self.add_parameter(val, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+        else:
+            try:
+                _columns = self._columns
+                if _columns and len(values) != len(_columns):
+                    msg = f"Number of values ({len(values)}) does not match the number of specified columns ({len(_columns)})."
+                    raise SQLBuilderError(msg)
+            except AttributeError:
+                pass
+            row_exprs = []
+            for i, v in enumerate(values):
+                if isinstance(v, exp.Expression):
+                    row_exprs.append(v)
+                else:
+                    try:
+                        _columns = self._columns
+                        if _columns and i < len(_columns):
+                            column_name = (
+                                str(_columns[i]).split(".")[-1] if "." in str(_columns[i]) else str(_columns[i])
+                            )
+                            param_name = self._generate_unique_parameter_name(column_name)
+                        else:
+                            param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
+                    except AttributeError:
+                        param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
+                    _, param_name = self.add_parameter(v, name=param_name)
+                    row_exprs.append(exp.var(param_name))
+
         values_expr = exp.Values(expressions=[row_exprs])
         self._expression.set("expression", values_expr)
         return self
@@ -110,10 +191,21 @@ class InsertValuesMixin:
         return self.values(*values)
 
 
+@trait
 class InsertFromSelectMixin:
     """Mixin providing INSERT ... SELECT support for INSERT builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
+    _table: Any  # Provided by QueryBuilder
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def from_select(self, select_builder: Any) -> Self:
         """Sets the INSERT source to a SELECT statement.
@@ -128,7 +220,7 @@ class InsertFromSelectMixin:
             SQLBuilderError: If the table is not set or the select_builder is invalid.
         """
         try:
-            if not self._table:  # type: ignore[attr-defined]
+            if not self._table:
                 msg = "The target table must be set using .into() before adding values."
                 raise SQLBuilderError(msg)
         except AttributeError:
@@ -139,11 +231,11 @@ class InsertFromSelectMixin:
         if not isinstance(self._expression, exp.Insert):
             msg = "Cannot set INSERT source on a non-INSERT expression."
             raise SQLBuilderError(msg)
-        subquery_parameters = select_builder._parameters  # pyright: ignore[attr-defined]
+        subquery_parameters = select_builder._parameters
         if subquery_parameters:
             for p_name, p_value in subquery_parameters.items():
-                self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
-        select_expr = select_builder._expression  # pyright: ignore[attr-defined]
+                self.add_parameter(p_value, name=p_name)
+        select_expr = select_builder._expression
         if select_expr and isinstance(select_expr, exp.Select):
             self._expression.set("expression", select_expr.copy())
         else:

sqlspec/builder/mixins/_join_operations.py

@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -13,9 +14,15 @@ if TYPE_CHECKING:
 __all__ = ("JoinClauseMixin",)
 
 
+@trait
 class JoinClauseMixin:
     """Mixin providing JOIN clause methods for SELECT builders."""
 
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
     def join(
         self,
         table: Union[str, exp.Expression, Any],
@@ -36,12 +43,12 @@ class JoinClauseMixin:
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
                 table_expr_value = getattr(table, "_expression", None)
                 if table_expr_value is not None:
-                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                    subquery_exp = exp.paren(table_expr_value)
                 else:
                     subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                subquery = table.build()  # pyright: ignore
+                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)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
@@ -99,12 +106,12 @@ class JoinClauseMixin:
             if hasattr(table, "_expression") and getattr(table, "_expression", None) is not None:
                 table_expr_value = getattr(table, "_expression", None)
                 if table_expr_value is not None:
-                    subquery_exp = exp.paren(table_expr_value.copy())  # pyright: ignore
+                    subquery_exp = exp.paren(table_expr_value)
                 else:
                     subquery_exp = exp.paren(exp.Anonymous(this=""))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
-                subquery = table.build()  # pyright: ignore
+                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)))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp

sqlspec/builder/mixins/_merge_operations.py

@@ -2,6 +2,7 @@
 
 from typing import Any, Optional, Union
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -18,10 +19,12 @@ __all__ = (
 )
 
 
+@trait
 class MergeIntoClauseMixin:
     """Mixin providing INTO clause for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+    _expression: Optional[exp.Expression]
 
     def into(self, table: Union[str, exp.Expression], alias: Optional[str] = None) -> Self:
         """Set the target table for the MERGE operation (INTO clause).
@@ -35,17 +38,24 @@ class MergeIntoClauseMixin:
             The current builder instance for method chaining.
         """
         if self._expression is None:
-            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))  # pyright: ignore
-        if not isinstance(self._expression, exp.Merge):  # pyright: ignore
-            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))  # pyright: ignore
+            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
+        if not isinstance(self._expression, exp.Merge):
+            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
         self._expression.set("this", exp.to_table(table, alias=alias) if isinstance(table, str) else table)
         return self
 
 
+@trait
 class MergeUsingClauseMixin:
     """Mixin providing USING clause for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+    _expression: Optional[exp.Expression]
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def using(self, source: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
         """Set the source data for the MERGE operation (USING clause).
@@ -73,7 +83,7 @@ class MergeUsingClauseMixin:
             subquery_builder_parameters = source.parameters
             if subquery_builder_parameters:
                 for p_name, p_value in subquery_builder_parameters.items():
-                    self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+                    self.add_parameter(p_value, name=p_name)
 
             subquery_exp = exp.paren(getattr(source, "_expression", exp.select()))
             source_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
@@ -89,10 +99,12 @@ class MergeUsingClauseMixin:
         return self
 
 
+@trait
 class MergeOnClauseMixin:
     """Mixin providing ON clause for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+    _expression: Optional[exp.Expression]
 
     def on(self, condition: Union[str, exp.Expression]) -> Self:
         """Set the join condition for the MERGE operation (ON clause).
@@ -131,10 +143,22 @@ class MergeOnClauseMixin:
         return self
 
 
+@trait
 class MergeMatchedClauseMixin:
     """Mixin providing WHEN MATCHED THEN ... clauses for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+    _expression: Optional[exp.Expression]
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate unique parameter name - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def _add_when_clause(self, when_clause: exp.When) -> None:
         """Helper to add a WHEN clause to the MERGE statement.
@@ -143,9 +167,9 @@ class MergeMatchedClauseMixin:
             when_clause: The WHEN clause to add.
         """
         if self._expression is None:
-            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
+            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))  # type: ignore[misc]
         if not isinstance(self._expression, exp.Merge):
-            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))
+            self._expression = exp.Merge(this=None, using=None, on=None, whens=exp.Whens(expressions=[]))  # type: ignore[misc]
 
         whens = self._expression.args.get("whens")
         if not whens:
@@ -175,8 +199,8 @@ class MergeMatchedClauseMixin:
             column_name = col if isinstance(col, str) else str(col)
             if "." in column_name:
                 column_name = column_name.split(".")[-1]
-            param_name = self._generate_unique_parameter_name(column_name)  # type: ignore[attr-defined]
-            param_name = self.add_parameter(val, name=param_name)[1]  # type: ignore[attr-defined]
+            param_name = self._generate_unique_parameter_name(column_name)
+            param_name = self.add_parameter(val, name=param_name)[1]
             update_expressions.append(exp.EQ(this=exp.column(col), expression=exp.var(param_name)))
 
         when_args: dict[str, Any] = {"matched": True, "then": exp.Update(expressions=update_expressions)}
@@ -238,10 +262,28 @@ class MergeMatchedClauseMixin:
         return self
 
 
+@trait
 class MergeNotMatchedClauseMixin:
     """Mixin providing WHEN NOT MATCHED THEN ... clauses for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    _expression: Optional[exp.Expression]
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate unique parameter name - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _add_when_clause(self, when_clause: exp.When) -> None:
+        """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def when_not_matched_then_insert(
         self,
@@ -278,8 +320,8 @@ class MergeNotMatchedClauseMixin:
                 column_name = columns[i] if isinstance(columns[i], str) else str(columns[i])
                 if "." in column_name:
                     column_name = column_name.split(".")[-1]
-                param_name = self._generate_unique_parameter_name(column_name)  # type: ignore[attr-defined]
-                param_name = self.add_parameter(val, name=param_name)[1]  # type: ignore[attr-defined]
+                param_name = self._generate_unique_parameter_name(column_name)
+                param_name = self.add_parameter(val, name=param_name)[1]
                 parameterized_values.append(exp.var(param_name))
 
             insert_args["this"] = exp.Tuple(expressions=[exp.column(c) for c in columns])
@@ -316,14 +358,32 @@ class MergeNotMatchedClauseMixin:
             when_args["this"] = condition_expr
 
         when_clause = exp.When(**when_args)
-        self._add_when_clause(when_clause)  # type: ignore[attr-defined]
+        self._add_when_clause(when_clause)
         return self
 
 
+@trait
 class MergeNotMatchedBySourceClauseMixin:
     """Mixin providing WHEN NOT MATCHED BY SOURCE THEN ... clauses for MERGE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    _expression: Optional[exp.Expression]
+
+    def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
+        """Add parameter - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _generate_unique_parameter_name(self, base_name: str) -> str:
+        """Generate unique parameter name - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
+
+    def _add_when_clause(self, when_clause: exp.When) -> None:
+        """Helper to add a WHEN clause to the MERGE statement - provided by QueryBuilder."""
+        msg = "Method must be provided by QueryBuilder subclass"
+        raise NotImplementedError(msg)
 
     def when_not_matched_by_source_then_update(
         self, set_values: dict[str, Any], condition: Optional[Union[str, exp.Expression]] = None
@@ -347,8 +407,8 @@ class MergeNotMatchedBySourceClauseMixin:
             column_name = col if isinstance(col, str) else str(col)
             if "." in column_name:
                 column_name = column_name.split(".")[-1]
-            param_name = self._generate_unique_parameter_name(column_name)  # type: ignore[attr-defined]
-            param_name = self.add_parameter(val, name=param_name)[1]  # type: ignore[attr-defined]
+            param_name = self._generate_unique_parameter_name(column_name)
+            param_name = self.add_parameter(val, name=param_name)[1]
             update_expressions.append(exp.EQ(this=exp.column(col), expression=exp.var(param_name)))
 
         when_args: dict[str, Any] = {
@@ -375,7 +435,7 @@ class MergeNotMatchedBySourceClauseMixin:
             when_args["this"] = condition_expr
 
         when_clause = exp.When(**when_args)
-        self._add_when_clause(when_clause)  # type: ignore[attr-defined]
+        self._add_when_clause(when_clause)
         return self
 
     def when_not_matched_by_source_then_delete(self, condition: Optional[Union[str, exp.Expression]] = None) -> Self:
@@ -412,5 +472,5 @@ class MergeNotMatchedBySourceClauseMixin:
             when_args["this"] = condition_expr
 
         when_clause = exp.When(**when_args)
-        self._add_when_clause(when_clause)  # type: ignore[attr-defined]
+        self._add_when_clause(when_clause)
         return self

sqlspec/builder/mixins/_order_limit_operations.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Optional, Union, cast
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -14,10 +15,14 @@ if TYPE_CHECKING:
 __all__ = ("LimitOffsetClauseMixin", "OrderByClauseMixin", "ReturningClauseMixin")
 
 
+@trait
 class OrderByClauseMixin:
     """Mixin providing ORDER BY clause."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def order_by(self, *items: Union[str, exp.Ordered], desc: bool = False) -> Self:
         """Add ORDER BY clause.
@@ -50,10 +55,14 @@ class OrderByClauseMixin:
         return cast("Self", builder)
 
 
+@trait
 class LimitOffsetClauseMixin:
     """Mixin providing LIMIT and OFFSET clauses."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def limit(self, value: int) -> Self:
         """Add LIMIT clause.
@@ -94,10 +103,13 @@ class LimitOffsetClauseMixin:
         return cast("Self", builder)
 
 
+@trait
 class ReturningClauseMixin:
     """Mixin providing RETURNING clause."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def returning(self, *columns: Union[str, exp.Expression]) -> Self:
         """Add RETURNING clause to the statement.

sqlspec/builder/mixins/_pivot_operations.py

@@ -2,6 +2,7 @@
 
 from typing import TYPE_CHECKING, Optional, Union, cast
 
+from mypy_extensions import trait
 from sqlglot import exp
 
 if TYPE_CHECKING:
@@ -12,10 +13,14 @@ if TYPE_CHECKING:
 __all__ = ("PivotClauseMixin", "UnpivotClauseMixin")
 
 
+@trait
 class PivotClauseMixin:
     """Mixin class to add PIVOT functionality to a Select."""
 
-    _expression: "Optional[exp.Expression]" = None
+    __slots__ = ()
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
     dialect: "DialectType" = None
 
     def pivot(
@@ -79,10 +84,14 @@ class PivotClauseMixin:
         return cast("Select", self)
 
 
+@trait
 class UnpivotClauseMixin:
     """Mixin class to add UNPIVOT functionality to a Select."""
 
-    _expression: "Optional[exp.Expression]" = None
+    __slots__ = ()
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
     dialect: "DialectType" = None
 
     def unpivot(

sqlspec/builder/mixins/_select_operations.py

@@ -3,6 +3,7 @@
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Optional, Union, cast
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -11,17 +12,20 @@ from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.type_guards import has_query_builder_parameters, is_expression
 
 if TYPE_CHECKING:
-    from sqlspec.builder._base import QueryBuilder
     from sqlspec.builder._column import Column, FunctionColumn
     from sqlspec.protocols import SelectBuilderProtocol, SQLBuilderProtocol
 
 __all__ = ("CaseBuilder", "SelectClauseMixin")
 
 
+@trait
 class SelectClauseMixin:
     """Consolidated mixin providing all SELECT-related clauses and functionality."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def select(self, *columns: Union[str, exp.Expression, "Column", "FunctionColumn"]) -> Self:
         """Add columns to SELECT clause.
@@ -529,7 +533,7 @@ class SelectClauseMixin:
         Returns:
             CaseBuilder: A CaseBuilder instance for building the CASE expression.
         """
-        builder = cast("QueryBuilder", self)  # pyright: ignore
+        builder = cast("SelectBuilderProtocol", self)
         return CaseBuilder(builder, alias)
 
 
@@ -537,15 +541,15 @@ class SelectClauseMixin:
 class CaseBuilder:
     """Builder for CASE expressions."""
 
-    _parent: "QueryBuilder"  # pyright: ignore
+    _parent: "SelectBuilderProtocol"
     _alias: Optional[str]
     _case_expr: exp.Case
 
-    def __init__(self, parent: "QueryBuilder", alias: "Optional[str]" = None) -> None:
+    def __init__(self, parent: "SelectBuilderProtocol", alias: "Optional[str]" = None) -> None:
         """Initialize CaseBuilder.
 
         Args:
-            parent: The parent builder.
+            parent: The parent builder with select capabilities.
             alias: Optional alias for the CASE expression.
         """
         self._parent = parent
@@ -589,11 +593,11 @@ class CaseBuilder:
         self._case_expr.set("default", value_expr)
         return self
 
-    def end(self) -> "QueryBuilder":
+    def end(self) -> "SelectBuilderProtocol":
         """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", self._parent.select(select_expr))  # type: ignore[attr-defined]
+        return self._parent.select(select_expr)