sqlspec 0.24.1__py3-none-any.whl → 0.26.0__py3-none-any.whl

sqlspec/builder/mixins/_update_operations.py

@@ -1,3 +1,4 @@
+# pyright: reportPrivateUsage=false
 """UPDATE operation mixins.
 
 Provides mixins for UPDATE statement functionality including
@@ -5,12 +6,13 @@ table specification, SET clauses, and FROM clauses.
 """
 
 from collections.abc import Mapping
-from typing import Any, Optional, Union, cast
+from typing import Any, Optional, Union
 
 from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
+from sqlspec.builder._parsing_utils import extract_sql_object_expression
 from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.type_guards import has_query_builder_parameters
 
@@ -25,8 +27,8 @@ class UpdateTableClauseMixin:
 
     __slots__ = ()
 
-    # Type annotation for PyRight - this will be provided by the base class
-    _expression: Optional[exp.Expression]
+    def get_expression(self) -> Optional[exp.Expression]: ...
+    def set_expression(self, expression: exp.Expression) -> None: ...
 
     def table(self, table_name: str, alias: Optional[str] = None) -> Self:
         """Set the table to update.
@@ -38,10 +40,14 @@ class UpdateTableClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        if self._expression is None or not isinstance(self._expression, exp.Update):
-            self._expression = exp.Update(this=None, expressions=[], joins=[])
+        current_expr = self.get_expression()
+        if current_expr is None or not isinstance(current_expr, exp.Update):
+            self.set_expression(exp.Update(this=None, expressions=[], joins=[]))
+            current_expr = self.get_expression()
+
+        assert current_expr is not None
         table_expr: exp.Expression = exp.to_table(table_name, alias=alias)
-        self._expression.set("this", table_expr)
+        current_expr.set("this", table_expr)
         setattr(self, "_table", table_name)
         return self
 
@@ -52,8 +58,8 @@ class UpdateSetClauseMixin:
 
     __slots__ = ()
 
-    # Type annotation for PyRight - this will be provided by the base class
-    _expression: Optional[exp.Expression]
+    def get_expression(self) -> Optional[exp.Expression]: ...
+    def set_expression(self, expression: exp.Expression) -> None: ...
 
     def add_parameter(self, value: Any, name: Optional[str] = None) -> tuple[Any, str]:
         """Add parameter - provided by QueryBuilder."""
@@ -86,24 +92,7 @@ class UpdateSetClauseMixin:
                     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))
+            return extract_sql_object_expression(val, builder=self)
         column_name = col if isinstance(col, str) else str(col)
         if "." in column_name:
             column_name = column_name.split(".")[-1]
@@ -130,9 +119,12 @@ 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):
+        current_expr = self.get_expression()
+        if current_expr is None:
+            self.set_expression(exp.Update())
+            current_expr = self.get_expression()
+
+        if not isinstance(current_expr, exp.Update):
             msg = "Cannot add SET clause to non-UPDATE expression."
             raise SQLBuilderError(msg)
         assignments = []
@@ -149,8 +141,8 @@ class UpdateSetClauseMixin:
         else:
             msg = "Invalid arguments for set(): use (column, value), mapping, or kwargs."
             raise SQLBuilderError(msg)
-        existing = self._expression.args.get("expressions", [])
-        self._expression.set("expressions", existing + assignments)
+        existing = current_expr.args.get("expressions", [])
+        current_expr.set("expressions", existing + assignments)
         return self
 
 
@@ -160,8 +152,8 @@ class UpdateFromClauseMixin:
 
     __slots__ = ()
 
-    # Type annotation for PyRight - this will be provided by the base class
-    _expression: Optional[exp.Expression]
+    def get_expression(self) -> Optional[exp.Expression]: ...
+    def set_expression(self, expression: exp.Expression) -> None: ...
 
     def from_(self, table: Union[str, exp.Expression, Any], alias: Optional[str] = None) -> Self:
         """Add a FROM clause to the UPDATE statement.
@@ -176,7 +168,8 @@ 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):
+        current_expr = self.get_expression()
+        if current_expr is None or not isinstance(current_expr, exp.Update):
             msg = "Cannot add FROM clause to non-UPDATE expression. Set the main table first."
             raise SQLBuilderError(msg)
         table_expr: exp.Expression
@@ -194,9 +187,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:
-            self._expression.set("from", exp.From(expressions=[]))
-        from_clause = self._expression.args["from"]
+        if current_expr.args.get("from") is None:
+            current_expr.set("from", exp.From(expressions=[]))
+        from_clause = current_expr.args["from"]
         if hasattr(from_clause, "append"):
             from_clause.append("expressions", table_expr)
         else:

sqlspec/builder/mixins/_where_clause.py

@@ -1,4 +1,5 @@
 # ruff: noqa: PLR2004
+# pyright: reportPrivateUsage=false, reportPrivateImportUsage=false
 """WHERE and HAVING clause mixins.
 
 Provides mixins for WHERE and HAVING clause functionality with
@@ -14,7 +15,9 @@ from mypy_extensions import trait
 from sqlglot import exp
 from typing_extensions import Self
 
-from sqlspec.builder._parsing_utils import parse_column_expression, parse_condition_expression
+from sqlspec.builder._parsing_utils import extract_column_name, parse_column_expression, parse_condition_expression
+from sqlspec.core.parameters import ParameterStyle, ParameterValidator
+from sqlspec.core.statement import SQL
 from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.type_guards import (
     has_expression_and_parameters,
@@ -24,31 +27,6 @@ from sqlspec.utils.type_guards import (
     is_iterable_parameters,
 )
 
-
-def _extract_column_name(column: Union[str, exp.Column]) -> str:
-    """Extract column name from column expression for parameter naming.
-
-    Args:
-        column: Column expression (string or SQLGlot Column)
-
-    Returns:
-        Column name as string for use as parameter name
-    """
-    if isinstance(column, str):
-        # Handle simple column names and table.column references
-        if "." in column:
-            return column.split(".")[-1]  # Return just the column part
-        return column
-    if isinstance(column, exp.Column):
-        # Extract the column name from SQLGlot Column expression
-        try:
-            return str(column.this.this)
-        except AttributeError:
-            return str(column.this) if column.this else "column"
-    # Fallback for any unexpected types (defensive programming)
-    return "column"
-
-
 if TYPE_CHECKING:
     from sqlspec.builder._column import ColumnExpression
     from sqlspec.protocols import SQLBuilderProtocol
@@ -62,8 +40,9 @@ class WhereClauseMixin:
 
     __slots__ = ()
 
-    # Type annotation for PyRight - this will be provided by the base class
-    _expression: Optional[exp.Expression]
+    # Type annotations for PyRight - these will be provided by the base class
+    def get_expression(self) -> Optional[exp.Expression]: ...
+    def set_expression(self, expression: exp.Expression) -> None: ...
 
     def _create_parameterized_condition(
         self,
@@ -82,7 +61,7 @@ class WhereClauseMixin:
             The created condition expression
         """
         builder = cast("SQLBuilderProtocol", self)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         param_name = builder._generate_unique_parameter_name(column_name)
         _, param_name = builder.add_parameter(value, name=param_name)
         col_expr = parse_column_expression(column) if not isinstance(column, exp.Column) else column
@@ -118,18 +97,20 @@ class WhereClauseMixin:
             Self with OR condition applied
         """
         # Create a temporary clone to capture the condition
-        original_expr = self._expression
+        original_expr = self.get_expression()
 
         # Apply the where method to get the condition
         where_method(*args, **kwargs)
 
         # Get the last condition added by extracting it from the modified expression
-        if isinstance(self._expression, (exp.Select, exp.Update, exp.Delete)) and original_expr != self._expression:
-            last_where = self._expression.find(exp.Where)
+        current_expr = self.get_expression()
+        if isinstance(current_expr, (exp.Select, exp.Update, exp.Delete)) and original_expr != current_expr:
+            last_where = current_expr.find(exp.Where)
             if last_where and last_where.this:
                 condition = last_where.this
                 # Restore original expression
-                self._expression = original_expr
+                if original_expr is not None:
+                    self.set_expression(original_expr)
                 # Apply as OR
                 return self.or_where(condition)
 
@@ -236,7 +217,7 @@ class WhereClauseMixin:
         column_name_raw, operator, value = condition
         operator = str(operator).upper()
         column_exp = parse_column_expression(column_name_raw)
-        column_name = _extract_column_name(column_name_raw)
+        column_name = extract_column_name(column_name_raw)
 
         # Simple operators that use direct parameterization
         simple_operators = {
@@ -299,16 +280,17 @@ class WhereClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        if self.__class__.__name__ == "Update" and not isinstance(self._expression, exp.Update):
+        current_expr = self.get_expression()
+        if self.__class__.__name__ == "Update" and not isinstance(current_expr, exp.Update):
             msg = "Cannot add WHERE clause to non-UPDATE expression"
             raise SQLBuilderError(msg)
 
         builder = cast("SQLBuilderProtocol", self)
-        if builder._expression is None:
+        if current_expr is None:
             msg = "Cannot add WHERE clause: expression is not initialized."
             raise SQLBuilderError(msg)
 
-        if isinstance(builder._expression, exp.Delete) and not builder._expression.args.get("this"):
+        if isinstance(current_expr, exp.Delete) and not current_expr.args.get("this"):
             msg = "WHERE clause requires a table to be set. Use from() to set the table first."
             raise SQLBuilderError(msg)
 
@@ -319,15 +301,11 @@ class WhereClauseMixin:
                 raise SQLBuilderError(msg)
 
             # Check if condition contains parameter placeholders
-            from sqlspec.core.parameters import ParameterStyle, ParameterValidator
-
             validator = ParameterValidator()
             param_info = validator.extract_parameters(condition)
 
             if param_info:
                 # String condition with placeholders - create SQL object with parameters
-                from sqlspec import sql as sql_factory
-
                 # Create parameter mapping based on the detected parameter info
                 param_dict = dict(kwargs)  # Start with named parameters
 
@@ -347,7 +325,7 @@ class WhereClauseMixin:
                     param_dict[f"param_{i}"] = value
 
                 # Create SQL object with parameters that will be processed correctly
-                condition = sql_factory.raw(condition, **param_dict)
+                condition = SQL(condition, param_dict)
                 # Fall through to existing SQL object handling logic
 
             elif len(values) == 1 and not kwargs:
@@ -357,10 +335,11 @@ class WhereClauseMixin:
                 else:
                     where_expr = self._process_tuple_condition((condition, values[0]))
                 # Process this condition and skip the rest
-                if isinstance(builder._expression, (exp.Select, exp.Update, exp.Delete)):
-                    builder._expression = builder._expression.where(where_expr, copy=False)
+                if isinstance(current_expr, (exp.Select, exp.Update, exp.Delete)):
+                    updated_expr = current_expr.where(where_expr, copy=False)
+                    self.set_expression(updated_expr)
                 else:
-                    msg = f"WHERE clause not supported for {type(builder._expression).__name__}"
+                    msg = f"WHERE clause not supported for {type(current_expr).__name__}"
                     raise SQLBuilderError(msg)
                 return self
             else:
@@ -400,10 +379,11 @@ class WhereClauseMixin:
             msg = f"Unsupported condition type: {type(condition).__name__}"
             raise SQLBuilderError(msg)
 
-        if isinstance(builder._expression, (exp.Select, exp.Update, exp.Delete)):
-            builder._expression = builder._expression.where(where_expr, copy=False)
+        if isinstance(current_expr, (exp.Select, exp.Update, exp.Delete)):
+            updated_expr = current_expr.where(where_expr, copy=False)
+            self.set_expression(updated_expr)
         else:
-            msg = f"WHERE clause not supported for {type(builder._expression).__name__}"
+            msg = f"WHERE clause not supported for {type(current_expr).__name__}"
             raise SQLBuilderError(msg)
         return self
 
@@ -448,7 +428,7 @@ class WhereClauseMixin:
     def where_between(self, column: Union[str, exp.Column], low: Any, high: Any) -> Self:
         """Add WHERE column BETWEEN low AND high clause."""
         builder = cast("SQLBuilderProtocol", self)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         low_param = builder._generate_unique_parameter_name(f"{column_name}_low")
         high_param = builder._generate_unique_parameter_name(f"{column_name}_high")
         _, low_param = builder.add_parameter(low, name=low_param)
@@ -460,7 +440,7 @@ class WhereClauseMixin:
     def where_like(self, column: Union[str, exp.Column], pattern: str, escape: Optional[str] = None) -> Self:
         """Add WHERE column LIKE pattern clause."""
         builder = cast("SQLBuilderProtocol", self)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         param_name = builder._generate_unique_parameter_name(column_name)
         _, param_name = builder.add_parameter(pattern, name=param_name)
         col_expr = parse_column_expression(column) if not isinstance(column, exp.Column) else column
@@ -519,7 +499,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, (str, bytes)):
             msg = "Unsupported type for 'values' in WHERE IN"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -548,7 +528,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, (str, bytes)):
             msg = "Values for where_not_in must be a non-string iterable or subquery."
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -638,7 +618,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, bytes):
             msg = "Unsupported type for 'values' in WHERE ANY"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -678,7 +658,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, bytes):
             msg = "Unsupported type for 'values' in WHERE NOT ANY"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -837,15 +817,11 @@ class WhereClauseMixin:
                 raise SQLBuilderError(msg)
 
             # Check if condition contains parameter placeholders
-            from sqlspec.core.parameters import ParameterStyle, ParameterValidator
-
             validator = ParameterValidator()
             param_info = validator.extract_parameters(condition)
 
             if param_info:
                 # String condition with placeholders - create SQL object with parameters
-                from sqlspec import sql as sql_factory
-
                 # Create parameter mapping based on the detected parameter info
                 param_dict = dict(kwargs)  # Start with named parameters
 
@@ -865,7 +841,7 @@ class WhereClauseMixin:
                     param_dict[f"param_{i}"] = value
 
                 # Create SQL object with parameters that will be processed correctly
-                condition = sql_factory.raw(condition, **param_dict)
+                condition = SQL(condition, param_dict)
                 # Fall through to existing SQL object handling logic
 
             elif len(values) == 1 and not kwargs:
@@ -975,7 +951,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, (str, bytes)):
             msg = "Unsupported type for 'values' in OR WHERE IN"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -990,7 +966,7 @@ class WhereClauseMixin:
     def or_where_like(self, column: Union[str, exp.Column], pattern: str, escape: Optional[str] = None) -> Self:
         """Add OR column LIKE pattern clause."""
         builder = cast("SQLBuilderProtocol", self)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         param_name = builder._generate_unique_parameter_name(column_name)
         _, param_name = builder.add_parameter(pattern, name=param_name)
         col_expr = parse_column_expression(column) if not isinstance(column, exp.Column) else column
@@ -1044,7 +1020,7 @@ class WhereClauseMixin:
     def or_where_between(self, column: Union[str, exp.Column], low: Any, high: Any) -> Self:
         """Add OR column BETWEEN low AND high clause."""
         builder = cast("SQLBuilderProtocol", self)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         low_param = builder._generate_unique_parameter_name(f"{column_name}_low")
         high_param = builder._generate_unique_parameter_name(f"{column_name}_high")
         _, low_param = builder.add_parameter(low, name=low_param)
@@ -1092,7 +1068,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, (str, bytes)):
             msg = "Values for or_where_not_in must be a non-string iterable or subquery."
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -1216,7 +1192,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, bytes):
             msg = "Unsupported type for 'values' in OR WHERE ANY"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -1272,7 +1248,7 @@ class WhereClauseMixin:
         if not is_iterable_parameters(values) or isinstance(values, bytes):
             msg = "Unsupported type for 'values' in OR WHERE NOT ANY"
             raise SQLBuilderError(msg)
-        column_name = _extract_column_name(column)
+        column_name = extract_column_name(column)
         parameters = []
         for i, v in enumerate(values):
             if len(values) == 1:
@@ -1292,7 +1268,9 @@ class HavingClauseMixin:
 
     __slots__ = ()
 
-    _expression: Optional[exp.Expression]
+    # Type annotations for PyRight - these will be provided by the base class
+    def get_expression(self) -> Optional[exp.Expression]: ...
+    def set_expression(self, expression: exp.Expression) -> None: ...
 
     def having(self, condition: Union[str, exp.Expression]) -> Self:
         """Add HAVING clause.
@@ -1306,11 +1284,15 @@ class HavingClauseMixin:
         Returns:
             The current builder instance for method chaining.
         """
-        if self._expression is None:
-            self._expression = exp.Select()
-        if not isinstance(self._expression, exp.Select):
+        current_expr = self.get_expression()
+        if current_expr is None:
+            self.set_expression(exp.Select())
+            current_expr = self.get_expression()
+
+        if not isinstance(current_expr, 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)
+        updated_expr = current_expr.having(having_expr, copy=False)
+        self.set_expression(updated_expr)
         return self