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

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
@@ -563,7 +567,8 @@ class CaseBuilder:
             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]
+        param_name = self._parent._generate_unique_parameter_name("case_when_value")
+        param_name = self._parent.add_parameter(value, name=param_name)[1]
         value_expr = exp.Placeholder(this=param_name)
 
         when_clause = exp.When(this=cond_expr, then=value_expr)
@@ -582,16 +587,17 @@ class CaseBuilder:
         Returns:
             CaseBuilder: The current builder instance for method chaining.
         """
-        param_name = self._parent.add_parameter(value)[1]
+        param_name = self._parent._generate_unique_parameter_name("case_else_value")
+        param_name = self._parent.add_parameter(value, name=param_name)[1]
         value_expr = exp.Placeholder(this=param_name)
         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)

sqlspec/builder/mixins/_update_operations.py

@@ -3,6 +3,7 @@
 from collections.abc import Mapping
 from typing import Any, Optional, Union
 
+from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
@@ -14,10 +15,14 @@ __all__ = ("UpdateFromClauseMixin", "UpdateSetClauseMixin", "UpdateTableClauseMi
 MIN_SET_ARGS = 2
 
 
+@trait
 class UpdateTableClauseMixin:
     """Mixin providing TABLE clause for UPDATE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
 
     def table(self, table_name: str, alias: Optional[str] = None) -> Self:
         """Set the table to update.
@@ -37,10 +42,24 @@ class UpdateTableClauseMixin:
         return self
 
 
+@trait
 class UpdateSetClauseMixin:
     """Mixin providing SET clause for UPDATE builders."""
 
-    _expression: Optional[exp.Expression] = None
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _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 set(self, *args: Any, **kwargs: Any) -> Self:
         """Set columns and values for the UPDATE statement.
@@ -79,9 +98,13 @@ class UpdateSetClauseMixin:
                 value_expr = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(self, "dialect", None)))
                 if has_query_builder_parameters(val):
                     for p_name, p_value in val.parameters.items():
-                        self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+                        self.add_parameter(p_value, name=p_name)
             else:
-                param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
+                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)[1]
                 value_expr = exp.Placeholder(this=param_name)
             assignments.append(exp.EQ(this=col_expr, expression=value_expr))
         elif (len(args) == 1 and isinstance(args[0], Mapping)) or kwargs:
@@ -95,9 +118,13 @@ class UpdateSetClauseMixin:
                     value_expr = exp.paren(exp.maybe_parse(sql_str, dialect=getattr(self, "dialect", None)))
                     if has_query_builder_parameters(val):
                         for p_name, p_value in val.parameters.items():
-                            self.add_parameter(p_value, name=p_name)  # type: ignore[attr-defined]
+                            self.add_parameter(p_value, name=p_name)
                 else:
-                    param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
+                    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)[1]
                     value_expr = exp.Placeholder(this=param_name)
                 assignments.append(exp.EQ(this=exp.column(col), expression=value_expr))
         else:
@@ -108,9 +135,15 @@ class UpdateSetClauseMixin:
         return self
 
 
+@trait
 class UpdateFromClauseMixin:
     """Mixin providing FROM clause for UPDATE builders (e.g., PostgreSQL style)."""
 
+    __slots__ = ()
+
+    # Type annotation for PyRight - this will be provided by the base class
+    _expression: Optional[exp.Expression]
+
     def from_(self, table: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
         """Add a FROM clause to the UPDATE statement.
 
@@ -124,7 +157,7 @@ class UpdateFromClauseMixin:
         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]
+        if self._expression is None or not isinstance(self._expression, exp.Update):
             msg = "Cannot add FROM clause to non-UPDATE expression. Set the main table first."
             raise SQLBuilderError(msg)
         table_expr: exp.Expression
@@ -142,9 +175,9 @@ class UpdateFromClauseMixin:
         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 self._expression.args.get("from") is None:
+            self._expression.set("from", exp.From(expressions=[]))
+        from_clause = self._expression.args["from"]
         if hasattr(from_clause, "append"):
             from_clause.append("expressions", table_expr)
         else: