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

sqlspec/builder/mixins/_update_operations.py

@@ -1,8 +1,9 @@
 """Update operation mixins for SQL builders."""
 
 from collections.abc import Mapping
-from typing import Any, Optional, Union
+from typing import Any, Optional, Union, cast
 
+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,70 @@ 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 _process_update_value(self, val: Any, col: Any) -> exp.Expression:
+        """Process a value for UPDATE assignment, handling SQL objects and parameters.
+
+        Args:
+            val: The value to process
+            col: The column name for parameter naming
+
+        Returns:
+            The processed expression for the value
+        """
+        if isinstance(val, exp.Expression):
+            return val
+        if has_query_builder_parameters(val):
+            subquery = val.build()
+            sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
+            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)
+            return value_expr
+        if hasattr(val, "expression") and hasattr(val, "sql"):
+            # Handle SQL objects (from sql.raw with parameters)
+            expression = getattr(val, "expression", None)
+            if expression is not None and isinstance(expression, exp.Expression):
+                # Merge parameters from SQL object into builder
+                if hasattr(val, "parameters"):
+                    sql_parameters = getattr(val, "parameters", {})
+                    for param_name, param_value in sql_parameters.items():
+                        self.add_parameter(param_value, name=param_name)
+                return cast("exp.Expression", expression)
+            # If expression is None, fall back to parsing the raw SQL
+            sql_text = getattr(val, "sql", "")
+            # Merge parameters even when parsing raw SQL
+            if hasattr(val, "parameters"):
+                sql_parameters = getattr(val, "parameters", {})
+                for param_name, param_value in sql_parameters.items():
+                    self.add_parameter(param_value, name=param_name)
+            parsed_expr = exp.maybe_parse(sql_text)
+            return parsed_expr if parsed_expr is not None else exp.convert(str(sql_text))
+        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]
+        return exp.Placeholder(this=param_name)
 
     def set(self, *args: Any, **kwargs: Any) -> Self:
         """Set columns and values for the UPDATE statement.
@@ -61,7 +126,6 @@ class UpdateSetClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-
         if self._expression is None:
             self._expression = exp.Update()
         if not isinstance(self._expression, exp.Update):
@@ -71,34 +135,12 @@ class UpdateSetClauseMixin:
         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 isinstance(val, exp.Expression):
-                value_expr = val
-            elif has_query_builder_parameters(val):
-                subquery = val.build()
-                sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
-                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]
-            else:
-                param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
-                value_expr = exp.Placeholder(this=param_name)
+            value_expr = self._process_update_value(val, col)
             assignments.append(exp.EQ(this=col_expr, expression=value_expr))
         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 isinstance(val, exp.Expression):
-                    value_expr = val
-                elif has_query_builder_parameters(val):
-                    subquery = val.build()
-                    sql_str = subquery.sql if hasattr(subquery, "sql") and not callable(subquery.sql) else str(subquery)
-                    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]
-                else:
-                    param_name = self.add_parameter(val)[1]  # type: ignore[attr-defined]
-                    value_expr = exp.Placeholder(this=param_name)
+                value_expr = self._process_update_value(val, col)
                 assignments.append(exp.EQ(this=exp.column(col), expression=value_expr))
         else:
             msg = "Invalid arguments for set(): use (column, value), mapping, or kwargs."
@@ -108,9 +150,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 +172,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 +190,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: