feat: add FOR UPDATE locking support to query builder (#88)

Builder improvements related to row locking:
- Add row-level locking support with `FOR UPDATE`, `FOR SHARE`, `SKIP LOCKED`, and `NOWAIT` clauses
- Implement support for table-specific locking with `OF` clause

Author: cofin

sqlspec/_sql.py

@@ -40,6 +40,7 @@
     MathExpression,
     StringExpression,
 )
+from sqlspec.builder._parsing_utils import extract_expression, to_expression
 from sqlspec.builder.mixins._join_operations import JoinBuilder
 from sqlspec.builder.mixins._select_operations import Case, SubqueryBuilder, WindowFunctionBuilder
 from sqlspec.core.statement import SQL
@@ -746,7 +747,7 @@ def count(
         if isinstance(column, str) and column == "*":
             expr = exp.Count(this=exp.Star(), distinct=distinct)
         else:
-            col_expr = self._extract_expression(column)
+            col_expr = extract_expression(column)
             expr = exp.Count(this=col_expr, distinct=distinct)
         return AggregateExpression(expr)
 
@@ -774,7 +775,7 @@ def sum(
         Returns:
             SUM expression.
         """
-        col_expr = SQLFactory._extract_expression(column)
+        col_expr = extract_expression(column)
         return AggregateExpression(exp.Sum(this=col_expr, distinct=distinct))
 
     @staticmethod
@@ -787,7 +788,7 @@ def avg(column: Union[str, exp.Expression, "ExpressionWrapper", "Case"]) -> Aggr
         Returns:
             AVG expression.
         """
-        col_expr = SQLFactory._extract_expression(column)
+        col_expr = extract_expression(column)
         return AggregateExpression(exp.Avg(this=col_expr))
 
     @staticmethod
@@ -800,7 +801,7 @@ def max(column: Union[str, exp.Expression, "ExpressionWrapper", "Case"]) -> Aggr
         Returns:
             MAX expression.
         """
-        col_expr = SQLFactory._extract_expression(column)
+        col_expr = extract_expression(column)
         return AggregateExpression(exp.Max(this=col_expr))
 
     @staticmethod
@@ -813,7 +814,7 @@ def min(column: Union[str, exp.Expression, "ExpressionWrapper", "Case"]) -> Aggr
         Returns:
             MIN expression.
         """
-        col_expr = SQLFactory._extract_expression(column)
+        col_expr = extract_expression(column)
         return AggregateExpression(exp.Min(this=col_expr))
 
     @staticmethod
@@ -1034,45 +1035,6 @@ def to_literal(value: Any) -> FunctionExpression:
             return FunctionExpression(value)
         return FunctionExpression(exp.convert(value))
 
-    @staticmethod
-    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)
-
-    @staticmethod
-    def _extract_expression(value: Any) -> exp.Expression:
-        """Extract SQLGlot expression from value, handling our wrapper types.
-
-        Args:
-            value: String, SQLGlot expression, or our wrapper type.
-
-        Returns:
-            Raw SQLGlot expression.
-        """
-        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)
-
     @staticmethod
     def decode(column: Union[str, exp.Expression], *args: Union[str, exp.Expression, Any]) -> FunctionExpression:
         """Create a DECODE expression (Oracle-style conditional logic).
@@ -1109,14 +1071,14 @@ def decode(column: Union[str, exp.Expression], *args: Union[str, exp.Expression,
 
         for i in range(0, len(args) - 1, 2):
             if i + 1 >= len(args):
-                default = SQLFactory._to_expression(args[i])
+                default = to_expression(args[i])
                 break
 
             search_val = args[i]
             result_val = args[i + 1]
 
-            search_expr = SQLFactory._to_expression(search_val)
-            result_expr = SQLFactory._to_expression(result_val)
+            search_expr = to_expression(search_val)
+            result_expr = to_expression(result_val)
 
             condition = exp.EQ(this=col_expr, expression=search_expr)
             conditions.append(exp.If(this=condition, true=result_expr))
@@ -1164,7 +1126,7 @@ def nvl(
             COALESCE expression equivalent to NVL.
         """
         col_expr = exp.column(column) if isinstance(column, str) else column
-        sub_expr = SQLFactory._to_expression(substitute_value)
+        sub_expr = to_expression(substitute_value)
         return ConversionExpression(exp.Coalesce(expressions=[col_expr, sub_expr]))
 
     @staticmethod
@@ -1192,8 +1154,8 @@ def nvl2(
             ```
         """
         col_expr = exp.column(column) if isinstance(column, str) else column
-        not_null_expr = SQLFactory._to_expression(value_if_not_null)
-        null_expr = SQLFactory._to_expression(value_if_null)
+        not_null_expr = to_expression(value_if_not_null)
+        null_expr = to_expression(value_if_null)
 
         is_null = exp.Is(this=col_expr, expression=exp.Null())
         condition = exp.Not(this=is_null)

sqlspec/builder/_base.py

@@ -3,6 +3,8 @@
 Provides abstract base classes and core functionality for SQL query builders.
 """
 
+import hashlib
+import uuid
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, NoReturn, Optional, Union, cast
 
@@ -19,13 +21,15 @@
 from sqlspec.core.statement import SQL, StatementConfig
 from sqlspec.exceptions import SQLBuilderError
 from sqlspec.utils.logging import get_logger
-from sqlspec.utils.type_guards import has_expression_and_parameters, has_sql_method, has_with_method
+from sqlspec.utils.type_guards import has_expression_and_parameters, has_sql_method, has_with_method, is_expression
 
 if TYPE_CHECKING:
     from sqlspec.core.result import SQLResult
 
 __all__ = ("QueryBuilder", "SafeQuery")
 
+MAX_PARAMETER_COLLISION_ATTEMPTS = 1000
+
 logger = get_logger(__name__)
 
 
@@ -104,13 +108,9 @@ def set_expression(self, expression: exp.Expression) -> None:
 
         Args:
             expression: SQLGlot expression to set
-
-        Raises:
-            TypeError: If expression is not a SQLGlot Expression
         """
-        if not isinstance(expression, exp.Expression):
-            msg = f"Expected Expression, got {type(expression)}"
-            raise TypeError(msg)
+        if not is_expression(expression):
+            self._raise_invalid_expression_type(expression)
         self._expression = expression
 
     def has_expression(self) -> bool:
@@ -151,6 +151,46 @@ def _raise_sql_builder_error(message: str, cause: Optional[BaseException] = None
         """
         raise SQLBuilderError(message) from cause
 
+    @staticmethod
+    def _raise_invalid_expression_type(expression: Any) -> NoReturn:
+        """Raise error for invalid expression type.
+
+        Args:
+            expression: The invalid expression object
+
+        Raises:
+            TypeError: Always raised for type mismatch
+        """
+        msg = f"Expected Expression, got {type(expression)}"
+        raise TypeError(msg)
+
+    @staticmethod
+    def _raise_cte_query_error(alias: str, message: str) -> NoReturn:
+        """Raise error for CTE query issues.
+
+        Args:
+            alias: CTE alias name
+            message: Specific error message
+
+        Raises:
+            SQLBuilderError: Always raised for CTE errors
+        """
+        msg = f"CTE '{alias}': {message}"
+        raise SQLBuilderError(msg)
+
+    @staticmethod
+    def _raise_cte_parse_error(cause: BaseException) -> NoReturn:
+        """Raise error for CTE parsing failures.
+
+        Args:
+            cause: The original parsing exception
+
+        Raises:
+            SQLBuilderError: Always raised with chained cause
+        """
+        msg = f"Failed to parse CTE query: {cause!s}"
+        raise SQLBuilderError(msg) from cause
+
     def _add_parameter(self, value: Any, context: Optional[str] = None) -> str:
         """Adds a parameter to the query and returns its placeholder name.
 
@@ -229,13 +269,11 @@ def _generate_unique_parameter_name(self, base_name: str) -> str:
         if base_name not in self._parameters:
             return base_name
 
-        for i in range(1, 1000):
+        for i in range(1, MAX_PARAMETER_COLLISION_ATTEMPTS):
             name = f"{base_name}_{i}"
             if name not in self._parameters:
                 return name
 
-        import uuid
-
         return f"{base_name}_{uuid.uuid4().hex[:8]}"
 
     def _merge_cte_parameters(self, cte_name: str, parameters: dict[str, Any]) -> dict[str, str]:
@@ -284,8 +322,6 @@ def _generate_builder_cache_key(self, config: "Optional[StatementConfig]" = None
         Returns:
             A unique cache key representing the builder state and configuration
         """
-        import hashlib
-
         dialect_name: str = self.dialect_name or "default"
 
         if self._expression is None:
@@ -339,35 +375,29 @@ def with_cte(self: Self, alias: str, query: "Union[QueryBuilder, exp.Select, str
         if isinstance(query, QueryBuilder):
             query_expr = query.get_expression()
             if query_expr is None:
-                self._raise_sql_builder_error("CTE query builder has no expression.")
+                self._raise_cte_query_error(alias, "query builder has no expression")
             if not isinstance(query_expr, exp.Select):
-                msg = f"CTE query builder expression must be a Select, got {type(query_expr).__name__}."
-                self._raise_sql_builder_error(msg)
+                self._raise_cte_query_error(alias, f"expression must be a Select, got {type(query_expr).__name__}")
             cte_select_expression = query_expr
             param_mapping = self._merge_cte_parameters(alias, query.parameters)
-            updated_expression = self._update_placeholders_in_expression(cte_select_expression, param_mapping)
-            if not isinstance(updated_expression, exp.Select):
-                msg = f"Updated CTE expression must be a Select, got {type(updated_expression).__name__}."
-                self._raise_sql_builder_error(msg)
-            cte_select_expression = updated_expression
+            cte_select_expression = cast(
+                "exp.Select", self._update_placeholders_in_expression(cte_select_expression, param_mapping)
+            )
 
         elif isinstance(query, str):
             try:
                 parsed_expression = sqlglot.parse_one(query, read=self.dialect_name)
                 if not isinstance(parsed_expression, exp.Select):
-                    msg = f"CTE query string must parse to a SELECT statement, got {type(parsed_expression).__name__}."
-                    self._raise_sql_builder_error(msg)
+                    self._raise_cte_query_error(
+                        alias, f"query string must parse to SELECT, got {type(parsed_expression).__name__}"
+                    )
                 cte_select_expression = parsed_expression
             except SQLGlotParseError as e:
-                self._raise_sql_builder_error(f"Failed to parse CTE query string: {e!s}", e)
-            except Exception as e:
-                msg = f"An unexpected error occurred while parsing CTE query string: {e!s}"
-                self._raise_sql_builder_error(msg, e)
+                self._raise_cte_parse_error(e)
         elif isinstance(query, exp.Select):
             cte_select_expression = query
         else:
-            msg = f"Invalid query type for CTE: {type(query).__name__}"
-            self._raise_sql_builder_error(msg)
+            self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}")
 
         self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=exp.to_table(alias))
         return self
@@ -438,10 +468,9 @@ def _optimize_expression(self, expression: exp.Expression) -> exp.Expression:
             optimized = optimize(
                 expression, schema=self.schema, dialect=self.dialect_name, optimizer_settings=optimizer_settings
             )
-
             cache.put("optimized", cache_key, optimized)
-
         except Exception:
+            logger.debug("Expression optimization failed, using original expression")
             return expression
         else:
             return optimized
@@ -482,18 +511,7 @@ def _to_statement(self, config: "Optional[StatementConfig]" = None) -> "SQL":
         """
         safe_query = self.build()
 
-        if isinstance(safe_query.parameters, dict):
-            kwargs = safe_query.parameters
-            parameters: Optional[tuple[Any, ...]] = None
-        else:
-            kwargs = None
-            parameters = (
-                safe_query.parameters
-                if isinstance(safe_query.parameters, tuple)
-                else tuple(safe_query.parameters)
-                if safe_query.parameters
-                else None
-            )
+        kwargs, parameters = self._extract_statement_parameters(safe_query.parameters)
 
         if config is None:
             config = StatementConfig(
@@ -521,6 +539,28 @@ def _to_statement(self, config: "Optional[StatementConfig]" = None) -> "SQL":
             return SQL(sql_string, *parameters, statement_config=config)
         return SQL(sql_string, statement_config=config)
 
+    def _extract_statement_parameters(
+        self, raw_parameters: Any
+    ) -> "tuple[Optional[dict[str, Any]], Optional[tuple[Any, ...]]]":
+        """Extract parameters for SQL statement creation.
+
+        Args:
+            raw_parameters: Raw parameter data from SafeQuery
+
+        Returns:
+            Tuple of (kwargs, parameters) for SQL statement construction
+        """
+        if isinstance(raw_parameters, dict):
+            return raw_parameters, None
+
+        if isinstance(raw_parameters, tuple):
+            return None, raw_parameters
+
+        if raw_parameters:
+            return None, tuple(raw_parameters)
+
+        return None, None
+
     def __str__(self) -> str:
         """Return the SQL string representation of the query.