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

sqlspec/builder/_insert.py

@@ -10,6 +10,7 @@ from sqlglot import exp
 from typing_extensions import Self
 
 from sqlspec.builder._base import QueryBuilder
+from sqlspec.builder._parsing_utils import extract_sql_object_expression
 from sqlspec.builder.mixins import InsertFromSelectMixin, InsertIntoClauseMixin, InsertValuesMixin, ReturningClauseMixin
 from sqlspec.core.result import SQLResult
 from sqlspec.exceptions import SQLBuilderError
@@ -46,7 +47,6 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
         """
         super().__init__(**kwargs)
 
-        # Initialize Insert-specific attributes
         self._table: Optional[str] = None
         self._columns: list[str] = []
         self._values_added_count: int = 0
@@ -90,6 +90,10 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
             raise SQLBuilderError(ERR_MSG_INTERNAL_EXPRESSION_TYPE)
         return self._expression
 
+    def get_insert_expression(self) -> exp.Insert:
+        """Get the insert expression (public API)."""
+        return self._get_insert_expression()
+
     def values(self, *values: Any, **kwargs: Any) -> "Self":
         """Adds a row of values to the INSERT statement.
 
@@ -126,10 +130,10 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
 
         if len(values) == 1:
             values_0 = values[0]
-            if hasattr(values_0, "items") and hasattr(values_0, "keys"):
+            if isinstance(values_0, dict):
                 return self.values_from_dict(values_0)
 
-        insert_expr = self._get_insert_expression()
+        insert_expr = self.get_insert_expression()
 
         if self._columns and len(values) != len(self._columns):
             msg = ERR_MSG_VALUES_COLUMNS_MISMATCH.format(values_len=len(values), columns_len=len(self._columns))
@@ -140,29 +144,15 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
             if isinstance(value, exp.Expression):
                 value_placeholders.append(value)
             elif has_expression_and_sql(value):
-                # Handle SQL objects (from sql.raw with parameters)
-                expression = getattr(value, "expression", None)
-                if expression is not None and isinstance(expression, exp.Expression):
-                    # Merge parameters from SQL object into builder
-                    self._merge_sql_object_parameters(value)
-                    value_placeholders.append(expression)
-                else:
-                    # If expression is None, fall back to parsing the raw SQL
-                    sql_text = getattr(value, "sql", "")
-                    # Merge parameters even when parsing raw SQL
-                    self._merge_sql_object_parameters(value)
-                    # Check if sql_text is callable (like Expression.sql method)
-                    if callable(sql_text):
-                        sql_text = str(value)
-                    value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
-                    value_placeholders.append(value_expr)
+                value_expr = extract_sql_object_expression(value, builder=self)
+                value_placeholders.append(value_expr)
             else:
                 if self._columns and i < len(self._columns):
                     column_str = str(self._columns[i])
                     column_name = column_str.rsplit(".", maxsplit=1)[-1] if "." in column_str else column_str
-                    param_name = self._generate_unique_parameter_name(column_name)
+                    param_name = self.generate_unique_parameter_name(column_name)
                 else:
-                    param_name = self._generate_unique_parameter_name(f"value_{i + 1}")
+                    param_name = self.generate_unique_parameter_name(f"value_{i + 1}")
                 _, param_name = self.add_parameter(value, name=param_name)
                 value_placeholders.append(exp.Placeholder(this=param_name))
 
@@ -254,17 +244,14 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
 
         Example:
             ```python
-            # ON CONFLICT (id) DO NOTHING
             sql.insert("users").values(id=1, name="John").on_conflict(
                 "id"
             ).do_nothing()
 
-            # ON CONFLICT (email, username) DO UPDATE SET updated_at = NOW()
             sql.insert("users").values(...).on_conflict(
                 "email", "username"
             ).do_update(updated_at=sql.raw("NOW()"))
 
-            # ON CONFLICT DO NOTHING (catches all conflicts)
             sql.insert("users").values(...).on_conflict().do_nothing()
             ```
         """
@@ -286,22 +273,41 @@ class Insert(QueryBuilder, ReturningClauseMixin, InsertValuesMixin, InsertFromSe
         return self.on_conflict(*columns).do_nothing()
 
     def on_duplicate_key_update(self, **kwargs: Any) -> "Insert":
-        """Adds conflict resolution using the ON CONFLICT syntax (cross-database compatible).
+        """Adds MySQL-style ON DUPLICATE KEY UPDATE clause.
 
         Args:
-            **kwargs: Column-value pairs to update on conflict.
+            **kwargs: Column-value pairs to update on duplicate key.
 
         Returns:
             The current builder instance for method chaining.
 
         Note:
-            This method uses PostgreSQL-style ON CONFLICT syntax but SQLGlot will
-            transpile it to the appropriate syntax for each database (MySQL's
-            ON DUPLICATE KEY UPDATE, etc.).
+            This method creates MySQL-specific ON DUPLICATE KEY UPDATE syntax.
+            For PostgreSQL, use on_conflict() instead.
         """
         if not kwargs:
             return self
-        return self.on_conflict().do_update(**kwargs)
+
+        insert_expr = self._get_insert_expression()
+
+        set_expressions = []
+        for col, val in kwargs.items():
+            if has_expression_and_sql(val):
+                value_expr = extract_sql_object_expression(val, builder=self)
+            elif isinstance(val, exp.Expression):
+                value_expr = val
+            else:
+                param_name = self.generate_unique_parameter_name(col)
+                _, param_name = self.add_parameter(val, name=param_name)
+                value_expr = exp.Placeholder(this=param_name)
+
+            set_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
+
+        on_conflict = exp.OnConflict(duplicate=True, action=exp.var("UPDATE"), expressions=set_expressions or None)
+
+        insert_expr.set("conflict", on_conflict)
+
+        return self
 
 
 class ConflictBuilder:
@@ -336,9 +342,8 @@ class ConflictBuilder:
             ).do_nothing()
             ```
         """
-        insert_expr = self._insert_builder._get_insert_expression()
+        insert_expr = self._insert_builder.get_insert_expression()
 
-        # Create ON CONFLICT with proper structure
         conflict_keys = [exp.to_identifier(col) for col in self._columns] if self._columns else None
         on_conflict = exp.OnConflict(conflict_keys=conflict_keys, action=exp.var("DO NOTHING"))
 
@@ -363,44 +368,21 @@ class ConflictBuilder:
             )
             ```
         """
-        insert_expr = self._insert_builder._get_insert_expression()
+        insert_expr = self._insert_builder.get_insert_expression()
 
-        # Create SET expressions for the UPDATE
         set_expressions = []
         for col, val in kwargs.items():
             if has_expression_and_sql(val):
-                # 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._insert_builder.add_parameter(param_value, name=param_name)
-                    value_expr = expression
-                else:
-                    # 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._insert_builder.add_parameter(param_value, name=param_name)
-                    # Check if sql_text is callable (like Expression.sql method)
-                    if callable(sql_text):
-                        sql_text = str(val)
-                    value_expr = exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
+                value_expr = extract_sql_object_expression(val, builder=self._insert_builder)
             elif isinstance(val, exp.Expression):
                 value_expr = val
             else:
-                # Create parameter for regular values
-                param_name = self._insert_builder._generate_unique_parameter_name(col)
+                param_name = self._insert_builder.generate_unique_parameter_name(col)
                 _, param_name = self._insert_builder.add_parameter(val, name=param_name)
                 value_expr = exp.Placeholder(this=param_name)
 
             set_expressions.append(exp.EQ(this=exp.column(col), expression=value_expr))
 
-        # Create ON CONFLICT with proper structure
         conflict_keys = [exp.to_identifier(col) for col in self._columns] if self._columns else None
         on_conflict = exp.OnConflict(
             conflict_keys=conflict_keys, action=exp.var("DO UPDATE"), expressions=set_expressions or None

sqlspec/builder/_parsing_utils.py

@@ -9,7 +9,7 @@ from typing import Any, Final, Optional, Union, cast
 
 from sqlglot import exp, maybe_parse, parse_one
 
-from sqlspec.core.parameters import ParameterStyle
+from sqlspec.core.parameters import ParameterStyle, ParameterValidator
 from sqlspec.utils.type_guards import (
     has_expression_and_parameters,
     has_expression_and_sql,
@@ -18,6 +18,27 @@ from sqlspec.utils.type_guards import (
 )
 
 
+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):
+        if "." in column:
+            return column.split(".")[-1]
+        return column
+    if isinstance(column, exp.Column):
+        try:
+            return str(column.this.this)
+        except AttributeError:
+            return str(column.this) if column.this else "column"
+    return "column"
+
+
 def parse_column_expression(
     column_input: Union[str, exp.Expression, Any], builder: Optional[Any] = None
 ) -> exp.Expression:
@@ -42,9 +63,7 @@ def parse_column_expression(
     if isinstance(column_input, exp.Expression):
         return column_input
 
-    # Handle SQL objects (from sql.raw with parameters)
     if has_expression_and_sql(column_input):
-        # This is likely a SQL object
         expression = getattr(column_input, "expression", None)
         if expression is not None and isinstance(expression, exp.Expression):
             # Merge parameters from SQL object into builder if available
@@ -53,9 +72,7 @@ def parse_column_expression(
                 for param_name, param_value in sql_parameters.items():
                     builder.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(column_input, "sql", "")
-        # Merge parameters even when parsing raw SQL
         if builder and has_expression_and_parameters(column_input) and hasattr(builder, "add_parameter"):
             sql_parameters = getattr(column_input, "parameters", {})
             for param_name, param_value in sql_parameters.items():
@@ -139,10 +156,8 @@ def parse_condition_expression(
         if value is None:
             return exp.Is(this=column_expr, expression=exp.null())
         if builder and has_parameter_builder(builder):
-            from sqlspec.builder.mixins._where_clause import _extract_column_name
-
-            column_name = _extract_column_name(column)
-            param_name = builder._generate_unique_parameter_name(column_name)
+            column_name = extract_column_name(column)
+            param_name = builder.generate_unique_parameter_name(column_name)
             _, param_name = builder.add_parameter(value, name=param_name)
             return exp.EQ(this=column_expr, expression=exp.Placeholder(this=param_name))
         if isinstance(value, str):
@@ -156,8 +171,6 @@ def parse_condition_expression(
 
     # Convert database-specific parameter styles to SQLGlot-compatible format
     # This ensures that placeholders like $1, %s, :1 are properly recognized as parameters
-    from sqlspec.core.parameters import ParameterValidator
-
     validator = ParameterValidator()
     param_info = validator.extract_parameters(condition_input)
 
@@ -186,4 +199,106 @@ def parse_condition_expression(
     return exp.condition(condition_input)
 
 
-__all__ = ("parse_column_expression", "parse_condition_expression", "parse_order_expression", "parse_table_expression")
+def extract_sql_object_expression(value: Any, builder: Optional[Any] = None) -> exp.Expression:
+    """Extract SQLGlot expression from SQL object value with parameter merging.
+
+    Handles the common pattern of:
+    1. Check if value has expression and SQL attributes
+    2. Try to get expression first, merge parameters if available
+    3. Fall back to parsing raw SQL text if expression is None
+    4. Merge parameters in both cases
+    5. Handle callable SQL text
+
+    This consolidates duplicated logic across builder files that process
+    SQL objects (like those from sql.raw() calls).
+
+    Args:
+        value: The SQL object value to process
+        builder: Optional builder instance for parameter merging (must have add_parameter method)
+
+    Returns:
+        SQLGlot Expression extracted from the SQL object
+
+    Raises:
+        ValueError: If the value doesn't appear to be a SQL object
+    """
+    if not has_expression_and_sql(value):
+        msg = f"Value does not have both expression and sql attributes: {type(value)}"
+        raise ValueError(msg)
+
+    # Try expression attribute first
+    expression = getattr(value, "expression", None)
+    if expression is not None and isinstance(expression, exp.Expression):
+        # Merge parameters if available and builder supports it
+        if builder and hasattr(value, "parameters") and hasattr(builder, "add_parameter"):
+            sql_parameters = getattr(value, "parameters", {})
+            for param_name, param_value in sql_parameters.items():
+                builder.add_parameter(param_value, name=param_name)
+        return cast("exp.Expression", expression)
+
+    # Fall back to parsing raw SQL text
+    sql_text = getattr(value, "sql", "")
+
+    # Merge parameters even when parsing raw SQL
+    if builder and hasattr(value, "parameters") and hasattr(builder, "add_parameter"):
+        sql_parameters = getattr(value, "parameters", {})
+        for param_name, param_value in sql_parameters.items():
+            builder.add_parameter(param_value, name=param_name)
+
+    # Handle callable SQL text
+    if callable(sql_text):
+        sql_text = str(value)
+
+    # Parse SQL text and return as expression
+    return exp.maybe_parse(sql_text) or exp.convert(str(sql_text))
+
+
+def extract_expression(value: Any) -> exp.Expression:
+    """Extract SQLGlot expression from value, handling wrapper types.
+
+    Args:
+        value: String, SQLGlot expression, or wrapper type.
+
+    Returns:
+        Raw SQLGlot expression.
+    """
+    from sqlspec.builder._column import Column
+    from sqlspec.builder._expression_wrappers import ExpressionWrapper
+    from sqlspec.builder.mixins._select_operations import Case
+
+    if isinstance(value, str):
+        return exp.column(value)
+    if isinstance(value, Column):
+        return value.sqlglot_expression
+    if isinstance(value, ExpressionWrapper):
+        return value.expression
+    if isinstance(value, Case):
+        return exp.Case(ifs=value.conditions, default=value.default)
+    if isinstance(value, exp.Expression):
+        return value
+    return exp.convert(value)
+
+
+def to_expression(value: Any) -> exp.Expression:
+    """Convert a Python value to a raw SQLGlot expression.
+
+    Args:
+        value: Python value or SQLGlot expression to convert.
+
+    Returns:
+        Raw SQLGlot expression.
+    """
+    if isinstance(value, exp.Expression):
+        return value
+    return exp.convert(value)
+
+
+__all__ = (
+    "extract_expression",
+    "extract_sql_object_expression",
+    "parse_column_expression",
+    "parse_condition_expression",
+    "parse_order_expression",
+    "parse_table_expression",
+    "to_expression",
+)

sqlspec/builder/_select.py

@@ -5,7 +5,7 @@ parameter binding and validation.
 """
 
 import re
-from typing import Any, Callable, Final, Optional, Union
+from typing import Any, Callable, Final, Optional, Union, cast
 
 from sqlglot import exp
 from typing_extensions import Self
@@ -24,6 +24,7 @@ from sqlspec.builder.mixins import (
     WhereClauseMixin,
 )
 from sqlspec.core.result import SQLResult
+from sqlspec.exceptions import SQLBuilderError
 
 __all__ = ("Select",)
 
@@ -73,7 +74,6 @@ class Select(
         """
         super().__init__(**kwargs)
 
-        # Initialize Select-specific attributes
         self._with_parts: dict[str, Union[exp.CTE, Select]] = {}
         self._hints: list[dict[str, object]] = []
 
@@ -169,3 +169,148 @@ class Select(
                 )
 
         return SafeQuery(sql=modified_sql, parameters=safe_query.parameters, dialect=safe_query.dialect)
+
+    def _validate_select_expression(self) -> None:
+        """Validate that current expression is a valid SELECT statement.
+
+        Raises:
+            SQLBuilderError: If expression is None or not a SELECT statement
+        """
+        if self._expression is None or not isinstance(self._expression, exp.Select):
+            msg = "Locking clauses can only be applied to SELECT statements"
+            raise SQLBuilderError(msg)
+
+    def _validate_lock_parameters(self, skip_locked: bool, nowait: bool) -> None:
+        """Validate locking parameters for conflicting options.
+
+        Args:
+            skip_locked: Whether SKIP LOCKED option is enabled
+            nowait: Whether NOWAIT option is enabled
+
+        Raises:
+            SQLBuilderError: If both skip_locked and nowait are True
+        """
+        if skip_locked and nowait:
+            msg = "Cannot use both skip_locked and nowait"
+            raise SQLBuilderError(msg)
+
+    def for_update(
+        self, *, skip_locked: bool = False, nowait: bool = False, of: "Optional[Union[str, list[str]]]" = None
+    ) -> "Self":
+        """Add FOR UPDATE clause to SELECT statement for row-level locking.
+
+        Args:
+            skip_locked: Skip rows that are already locked (SKIP LOCKED)
+            nowait: Return immediately if row is locked (NOWAIT)
+            of: Table names/aliases to lock (FOR UPDATE OF table)
+
+        Returns:
+            Self for method chaining
+        """
+        self._validate_select_expression()
+        self._validate_lock_parameters(skip_locked, nowait)
+
+        assert self._expression is not None
+        select_expr = cast("exp.Select", self._expression)
+
+        lock_args = {"update": True}
+
+        if skip_locked:
+            lock_args["wait"] = False
+        elif nowait:
+            lock_args["wait"] = True
+
+        if of:
+            tables = [of] if isinstance(of, str) else of
+            lock_args["expressions"] = [exp.table_(t) for t in tables]  # type: ignore[assignment]
+
+        lock = exp.Lock(**lock_args)
+
+        current_locks = select_expr.args.get("locks", [])
+        current_locks.append(lock)
+        select_expr.set("locks", current_locks)
+
+        return self
+
+    def for_share(
+        self, *, skip_locked: bool = False, nowait: bool = False, of: "Optional[Union[str, list[str]]]" = None
+    ) -> "Self":
+        """Add FOR SHARE clause for shared row-level locking.
+
+        Args:
+            skip_locked: Skip rows that are already locked (SKIP LOCKED)
+            nowait: Return immediately if row is locked (NOWAIT)
+            of: Table names/aliases to lock (FOR SHARE OF table)
+
+        Returns:
+            Self for method chaining
+        """
+        self._validate_select_expression()
+        self._validate_lock_parameters(skip_locked, nowait)
+
+        assert self._expression is not None
+        select_expr = cast("exp.Select", self._expression)
+
+        lock_args = {"update": False}
+
+        if skip_locked:
+            lock_args["wait"] = False
+        elif nowait:
+            lock_args["wait"] = True
+
+        if of:
+            tables = [of] if isinstance(of, str) else of
+            lock_args["expressions"] = [exp.table_(t) for t in tables]  # type: ignore[assignment]
+
+        lock = exp.Lock(**lock_args)
+
+        current_locks = select_expr.args.get("locks", [])
+        current_locks.append(lock)
+        select_expr.set("locks", current_locks)
+
+        return self
+
+    def for_key_share(self) -> "Self":
+        """Add FOR KEY SHARE clause (PostgreSQL-specific).
+
+        FOR KEY SHARE is like FOR SHARE, but the lock is weaker:
+        SELECT FOR UPDATE is blocked, but not SELECT FOR NO KEY UPDATE.
+
+        Returns:
+            Self for method chaining
+        """
+        self._validate_select_expression()
+
+        assert self._expression is not None
+        select_expr = cast("exp.Select", self._expression)
+
+        lock = exp.Lock(update=False, key=True)
+
+        current_locks = select_expr.args.get("locks", [])
+        current_locks.append(lock)
+        select_expr.set("locks", current_locks)
+
+        return self
+
+    def for_no_key_update(self) -> "Self":
+        """Add FOR NO KEY UPDATE clause (PostgreSQL-specific).
+
+        FOR NO KEY UPDATE is like FOR UPDATE, but the lock is weaker:
+        it does not block SELECT FOR KEY SHARE commands that attempt to
+        acquire a share lock on the same rows.
+
+        Returns:
+            Self for method chaining
+        """
+        self._validate_select_expression()
+
+        assert self._expression is not None
+        select_expr = cast("exp.Select", self._expression)
+
+        lock = exp.Lock(update=True, key=False)
+
+        current_locks = select_expr.args.get("locks", [])
+        current_locks.append(lock)
+        select_expr.set("locks", current_locks)
+
+        return self

sqlspec/builder/_update.py

@@ -131,7 +131,7 @@ class Update(
             subquery_exp = exp.paren(exp.maybe_parse(subquery.sql, dialect=self.dialect))
             table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
 
-            subquery_parameters = table._parameters
+            subquery_parameters = table.parameters
             if subquery_parameters:
                 for p_name, p_value in subquery_parameters.items():
                     self.add_parameter(p_value, name=p_name)