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

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,
@@ -63,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
@@ -74,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():
@@ -175,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)
 
@@ -205,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/mixins/_cte_and_set_ops.py

@@ -117,8 +117,7 @@ class CommonTableExpressionMixin:
             if recursive:
                 existing_with.set("recursive", recursive)
         else:
-            # Only SELECT, INSERT, UPDATE support WITH clauses
-            if hasattr(expression, "with_") and isinstance(expression, (exp.Select, exp.Insert, exp.Update)):
+            if isinstance(expression, (exp.Select, exp.Insert, exp.Update)):
                 updated_expression = expression.with_(cte_alias_expr, as_=name, copy=False)
                 builder.set_expression(updated_expression)
                 if recursive:

sqlspec/builder/mixins/_join_operations.py

@@ -82,16 +82,12 @@ class JoinClauseMixin:
         self, table: Any, alias: Optional[str], builder: "SQLBuilderProtocol"
     ) -> exp.Expression:
         """Handle table parameters that are query builders."""
-        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)
-            else:
-                subquery_exp = exp.paren(exp.Anonymous(this=""))
+        if hasattr(table, "_expression") and table._expression is not None:
+            subquery_exp = exp.paren(table._expression)
             return exp.alias_(subquery_exp, alias) if alias else subquery_exp
         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)))
+        subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=builder.dialect))
         return exp.alias_(subquery_exp, alias) if alias else subquery_exp
 
     def _parse_on_condition(
@@ -107,28 +103,20 @@ class JoinClauseMixin:
             return self._handle_sql_object_condition(on, builder)
         if isinstance(on, exp.Expression):
             return on
-        # Last resort - convert to string and parse
         return exp.condition(str(on))
 
     def _handle_sql_object_condition(self, on: Any, builder: "SQLBuilderProtocol") -> exp.Expression:
         """Handle SQL object conditions with parameter binding."""
-        expression = getattr(on, "expression", None)
-        if expression is not None and isinstance(expression, exp.Expression):
-            # Merge parameters from SQL object into builder
-            if hasattr(on, "parameters") and hasattr(builder, "add_parameter"):
-                sql_parameters = getattr(on, "parameters", {})
-                for param_name, param_value in sql_parameters.items():
+        if hasattr(on, "expression") and on.expression is not None:
+            if hasattr(on, "parameters"):
+                for param_name, param_value in on.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(on, "sql", "")
-        # Merge parameters even when parsing raw SQL
-        if hasattr(on, "parameters") and hasattr(builder, "add_parameter"):
-            sql_parameters = getattr(on, "parameters", {})
-            for param_name, param_value in sql_parameters.items():
+            return cast("exp.Expression", on.expression)
+        if hasattr(on, "parameters"):
+            for param_name, param_value in on.parameters.items():
                 builder.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.condition(str(sql_text))
+        parsed_expr = exp.maybe_parse(on.sql)
+        return parsed_expr if parsed_expr is not None else exp.condition(str(on.sql))
 
     def _create_join_expression(
         self, table_expr: exp.Expression, on_expr: Optional[exp.Expression], join_type: str
@@ -195,17 +183,13 @@ class JoinClauseMixin:
         if isinstance(table, str):
             table_expr = parse_table_expression(table, alias)
         elif has_query_builder_parameters(table):
-            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)
-                else:
-                    subquery_exp = exp.paren(exp.Anonymous(this=""))
+            if hasattr(table, "_expression") and table._expression is not None:
+                subquery_exp = exp.paren(table._expression)
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
             else:
                 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)))
+                subquery_exp = exp.paren(exp.maybe_parse(sql_str, dialect=builder.dialect))
                 table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp
         else:
             table_expr = table