diff --git a/docs/changelog.rst b/docs/changelog.rst index b12fef729..68cfc01be 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -212,6 +212,9 @@ v0.63.0 - Transactions, table fixtures, SQL fragments, storage, and kwargs param consumption without intermediate tuple relays. (`#771 `_) +* ``sql.values`` creates a :class:`~sqlspec.builder.Values` builder for parameterized bulk row lists rather than resolving as a column named ``values``. Use ``sql.column("values")`` to construct column expressions referencing that identifier. + (`#779 `_) + **Fixed:** * Preserve JSON objects and arrays as individual query parameters after placeholder conversion, @@ -233,6 +236,12 @@ v0.63.0 - Transactions, table fixtures, SQL fragments, storage, and kwargs param deduplicate concurrent event IDs with a single key-range-locked statement. (`#781 `_) +* ``Update.from_()`` accepts query builders and subqueries with parameter merging instead of raising a runtime type error, and dialect checks raise a descriptive ``SQLBuilderError`` on dialects without native ``FROM`` clause support. + (`#779 `_) + +* CTEs registered via ``with_cte()`` or ``with_()`` render on ``Update`` and ``Delete`` statements and merge bound parameters without inlining or dropping explicit common table expressions. + (`#779 `_) + * Query builder keeps ``ON CONFLICT ... DO UPDATE`` and ``ON DUPLICATE KEY UPDATE`` assignments in written order. Assignments such as ``do_update(name=exp.column("name", table="excluded"))`` no longer render reversed. Conflict targets and update columns are quoted, allowing reserved words (e.g. ``order``, ``group``) diff --git a/docs/reference/builder/queries.rst b/docs/reference/builder/queries.rst index 635b55143..c9080ba6a 100644 --- a/docs/reference/builder/queries.rst +++ b/docs/reference/builder/queries.rst @@ -158,3 +158,10 @@ Joins .. autoclass:: JoinBuilder :members: :show-inheritance: + +Values +====== + +.. autoclass:: Values + :members: + :show-inheritance: diff --git a/docs/usage/query_builder.rst b/docs/usage/query_builder.rst index d83ec0368..2120ce8f3 100644 --- a/docs/usage/query_builder.rst +++ b/docs/usage/query_builder.rst @@ -38,6 +38,68 @@ Inserts and Updates :dedent: 4 :no-upgrade: +UPDATE ... FROM and CTEs +------------------------ + +Builder queries support ``UPDATE ... FROM`` with subqueries or Common Table Expressions (CTEs), +enabling queue claim statements and batch updates across supported dialects. + +.. code-block:: python + + from sqlspec import sql + + claim_query = ( + sql.update("tasks") + .set(status="processing") + .from_( + sql.select("id") + .from_("tasks") + .where_eq("status", "pending") + .limit(1) + .for_update(skip_locked=True), + alias="sub", + ) + .where("tasks.id = sub.id") + .returning(sql.column("id", table="tasks")) + ) + +Dialect Support Matrix +~~~~~~~~~~~~~~~~~~~~~~ + +.. list-table:: + :header-rows: 1 + + * - Dialect + - UPDATE ... FROM Support + - Notes + * - PostgreSQL + - Yes + - Native ``UPDATE ... FROM`` with ``RETURNING`` and row locking + * - CockroachDB + - Yes + - Native ``UPDATE ... FROM`` + * - SQLite + - Yes + - Native ``UPDATE ... FROM`` (SQLite 3.33.0+) + * - DuckDB + - Yes + - Native ``UPDATE ... FROM`` + * - SQL Server (MSSQL) + - Yes + - Native ``UPDATE ... FROM`` + * - MySQL / MariaDB + - No + - Raises ``SQLBuilderError``; use multi-table join update or MERGE + * - Oracle + - No + - Builder conservatively raises ``SQLBuilderError``; use MERGE or raw SQL on Oracle 23+ + * - Spanner + - No + - Raises ``SQLBuilderError`` + * - BigQuery + - Yes + - Native ``UPDATE ... FROM``; a ``WHERE`` condition is required + Upserts (ON CONFLICT) --------------------- diff --git a/sqlspec/builder/__init__.py b/sqlspec/builder/__init__.py index a919f42b9..eef646bbc 100644 --- a/sqlspec/builder/__init__.py +++ b/sqlspec/builder/__init__.py @@ -91,6 +91,7 @@ ) from sqlspec.builder._temporal import create_temporal_table, register_version_generators from sqlspec.builder._update import Update +from sqlspec.builder._values import Values from sqlspec.builder._vector_distance import VectorDistance from sqlspec.exceptions import SQLBuilderError @@ -151,6 +152,7 @@ "UpdateFromClauseMixin", "UpdateSetClauseMixin", "UpdateTableClauseMixin", + "Values", "VectorDistance", "WhereClauseMixin", "WindowFunctionBuilder", diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 7b3c10f1c..80eb74eb6 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -13,6 +13,8 @@ from sqlglot.dialects.dialect import DialectType from sqlglot.errors import ParseError as SQLGlotParseError from sqlglot.optimizer import RULES, optimize +from sqlglot.optimizer.eliminate_ctes import eliminate_ctes as _eliminate_ctes_rule +from sqlglot.optimizer.merge_subqueries import merge_subqueries as _merge_subqueries_rule from sqlglot.optimizer.normalize_identifiers import normalize_identifiers as _normalize_identifiers_rule from sqlglot.optimizer.optimize_joins import optimize_joins as _optimize_joins_rule from sqlglot.optimizer.pushdown_predicates import pushdown_predicates as _pushdown_predicates_rule @@ -38,7 +40,7 @@ from sqlspec.data_dictionary import get_dialect_config from sqlspec.exceptions import SQLBuilderError from sqlspec.utils.logging import get_logger -from sqlspec.utils.type_guards import has_expression_and_parameters, has_name, has_with_method, is_expression +from sqlspec.utils.type_guards import has_expression_and_parameters, has_name, is_expression from sqlspec.utils.uuids import uuid4 __all__ = ("BuiltQuery", "ExpressionBuilder", "QueryBuilder") @@ -316,14 +318,16 @@ def _build_final_expression(self, *, copy: bool = False) -> exp.Expr: return base_expression final_expression: exp.Expr = base_expression - if has_with_method(final_expression): - for alias, cte_node in self._with_ctes.items(): - final_expression = cast("Any", final_expression).with_(alias, as_=cte_node.args["this"], copy=False) - return cast("exp.Expr", final_expression) - - if "with_" in type(final_expression).arg_types: + existing_with = final_expression.args.get("with_") + if existing_with is None: final_expression.set("with_", exp.With(expressions=list(self._with_ctes.values()))) + else: + for cte_node in self._with_ctes.values(): + if cte_node not in existing_with.expressions: + existing_with.append("expressions", cte_node) + if any(cte.meta.get("recursive") for cte in self._with_ctes.values()): + final_expression.args["with_"].set("recursive", True) return final_expression def _spawn_like_self(self: Self) -> Self: @@ -337,35 +341,66 @@ def _spawn_like_self(self: Self) -> Self: simplify_expressions=self.simplify_expressions, ) - def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | str") -> exp.Select: - """Resolve a CTE query into a Select expression with merged parameters.""" + def _resolve_cte_query( + self, alias: str, query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any" + ) -> exp.Expr: + """Resolve a CTE query into a Select or Values expression with merged parameters.""" if isinstance(query, QueryBuilder): - query_expr = query.get_expression() + query_expr = query._build_final_expression(copy=True) if query_expr is None: self._raise_cte_query_error(alias, "query builder has no expression") - if not isinstance(query_expr, exp.Select): - self._raise_cte_query_error(alias, f"expression must be a Select, got {type(query_expr).__name__}") - cte_select_expression = query_expr.copy() + if not isinstance(query_expr, (exp.Select, exp.SetOperation, exp.Values)): + self._raise_cte_query_error( + alias, f"expression must be a Select or Values, got {type(query_expr).__name__}" + ) + cte_select_expression: exp.Expr = query_expr.copy() + if isinstance(cte_select_expression, exp.Values) and cte_select_expression.args.get("alias"): + cte_select_expression = cte_select_expression.copy() + cte_select_expression.set("alias", None) param_mapping = self._merge_cte_parameters(alias, query.parameters) - updated_expression = self._update_placeholders(cte_select_expression, param_mapping) - if not isinstance(updated_expression, exp.Select): # pragma: no cover - msg = "CTE placeholder update produced non-select expression" - raise SQLBuilderError(msg) - return updated_expression + if param_mapping: + cte_select_expression = self._update_placeholders(cte_select_expression, param_mapping) + return cte_select_expression + + if hasattr(query, "get_expression"): + raw_query = cast("Any", query) + raw_query_expr = ( + raw_query._build_final_expression(copy=True) + if hasattr(raw_query, "_build_final_expression") + else raw_query.get_expression() + ) + if raw_query_expr is None: + self._raise_cte_query_error(alias, "query builder has no expression") + if not isinstance(raw_query_expr, (exp.Select, exp.SetOperation, exp.Values)): + self._raise_cte_query_error( + alias, f"expression must be a Select or Values, got {type(raw_query_expr).__name__}" + ) + cte_duck_expression: exp.Expr = raw_query_expr.copy() + if isinstance(cte_duck_expression, exp.Values) and cte_duck_expression.args.get("alias"): + cte_duck_expression = cte_duck_expression.copy() + cte_duck_expression.set("alias", None) + if hasattr(raw_query, "parameters"): + param_mapping = self._merge_cte_parameters(alias, raw_query.parameters) + if param_mapping: + cte_duck_expression = self._update_placeholders(cte_duck_expression, param_mapping) + return cte_duck_expression if isinstance(query, str): try: parsed_expression = sqlglot.parse_one(query, read=self.dialect_name) except SQLGlotParseError as e: # pragma: no cover self._raise_cte_parse_error(e) - if not isinstance(parsed_expression, exp.Select): + if not isinstance(parsed_expression, (exp.Select, exp.SetOperation, exp.Values)): self._raise_cte_query_error( - alias, f"query string must parse to SELECT, got {type(parsed_expression).__name__}" + alias, f"query string must parse to SELECT or VALUES, got {type(parsed_expression).__name__}" ) return parsed_expression - if isinstance(query, exp.Select): - return query + if isinstance(query, (exp.Select, exp.SetOperation, exp.Values)): + cte_expression = query.copy() + if isinstance(cte_expression, exp.Values): + cte_expression.set("alias", None) + return cte_expression self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}") msg = "Unreachable" @@ -579,13 +614,21 @@ def _cache_key(self, config: "StatementConfig | None" = None) -> str: ) return f"builder:{fingerprint}" - def with_cte(self: Self, alias: str, query: "QueryBuilder | exp.Select | str") -> Self: + def with_cte( + self: Self, + alias: str, + query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any", + recursive: bool = False, + columns: "list[str] | None" = None, + ) -> Self: """Adds a Common Table Expression (CTE) to the query. Args: alias: The alias for the CTE. query: The CTE query, which can be another QueryBuilder instance, - a raw SQL string, or a sqlglot Select expression. + a raw SQL string, or a sqlglot Select or Values expression. + recursive: Whether the CTE is recursive. + columns: Optional list of column aliases for the CTE. Returns: Self: The current builder instance for method chaining. @@ -594,9 +637,49 @@ def with_cte(self: Self, alias: str, query: "QueryBuilder | exp.Select | str") - self._raise_builder_error(f"CTE with alias '{alias}' already exists.") cte_select_expression = self._resolve_cte_query(alias, query) - self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=exp.to_table(alias)) + cte_columns = columns + if cte_columns is None and isinstance(query, exp.Values): + values_alias = query.args.get("alias") + if isinstance(values_alias, exp.TableAlias): + cte_columns = [column.name for column in values_alias.columns] + if cte_columns is None: + query_cols = getattr(query, "columns", None) + query_private_cols = getattr(query, "_columns", None) + if isinstance(query_cols, (list, tuple)): + cte_columns = [str(c) for c in query_cols] + elif isinstance(query_private_cols, (list, tuple)): + cte_columns = [str(c) for c in query_private_cols] + + if cte_columns: + alias_node: exp.Expr = exp.TableAlias( + this=exp.to_identifier(alias), columns=[exp.to_identifier(c) for c in cte_columns] + ) + else: + alias_node = exp.to_table(alias) + self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=alias_node) + self._with_ctes[alias].meta["recursive"] = recursive return self + def with_( + self: Self, + name: str, + query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any", + recursive: bool = False, + columns: "list[str] | None" = None, + ) -> Self: + """Alias for with_cte for parity across builders. + + Args: + name: The alias/name for the CTE. + query: The CTE query expression or builder. + recursive: Whether the CTE is recursive. + columns: Optional list of column aliases for the CTE. + + Returns: + Self: The current builder instance for method chaining. + """ + return self.with_cte(name, query, recursive=recursive, columns=columns) + def build(self, dialect: DialectType = None) -> "BuiltQuery": """Builds the SQL query string and parameters. @@ -608,6 +691,7 @@ def build(self, dialect: DialectType = None) -> "BuiltQuery": BuiltQuery: A dataclass containing the SQL string and parameters. """ final_expression = self._build_final_expression() + self._validate_update_from(final_expression, _resolve_dialect(dialect, self.dialect)) if self.enable_optimization and isinstance(final_expression, exp.Expr): final_expression = self._optimize_expression(final_expression) @@ -805,6 +889,10 @@ def _optimize_expression(self, expression: exp.Expr, *, force: bool = False) -> if cached_optimized is not None: return cast("exp.Expr", cached_optimized).copy() + # Qualification drops VALUES CTE column aliases without projecting replacements. + if any(isinstance(cte.this, exp.Values) and cte.alias_column_names for cte in expression.find_all(exp.CTE)): + return expression + excluded_rules = set() if not self.optimize_joins: excluded_rules.add(_optimize_joins_rule) @@ -812,6 +900,9 @@ def _optimize_expression(self, expression: exp.Expr, *, force: bool = False) -> excluded_rules.add(_pushdown_predicates_rule) if not self.simplify_expressions: excluded_rules.add(_simplify_rule) + if expression.args.get("with_") is not None or self._with_ctes: + excluded_rules.add(_eliminate_ctes_rule) + excluded_rules.add(_merge_subqueries_rule) rules = RULES if not excluded_rules else tuple(rule for rule in RULES if rule not in excluded_rules) @@ -914,10 +1005,30 @@ def _to_statement(self, config: "StatementConfig | None" = None) -> "SQL": cache_entry = self._create_builder_cache_entry(config) return self._statement_from_cache_entry(cache_entry, config) + def _validate_update_from(self, expression: exp.Expr, dialect: DialectType) -> None: + if not dialect or not any(node.args.get("from_") is not None for node in expression.find_all(exp.Update)): + return + from sqlspec.data_dictionary import get_dialect_config + + dialect_name = ( + dialect.lower() if isinstance(dialect, str) else type(Dialect.get_or_raise(dialect)).__name__.lower() + ) + try: + config = get_dialect_config("spanner" if dialect_name == "spangres" else dialect_name) + except ValueError: + return + if not config.feature_flags.get("supports_update_from", True): + msg = ( + f"Dialect '{dialect_name}' does not support UPDATE ... FROM clauses. " + "Consider using MERGE or a JOIN-based UPDATE instead." + ) + raise SQLBuilderError(msg) + def _create_builder_cache_entry(self, config: "StatementConfig | None") -> "_BuilderCacheEntry": dialect_override = config.dialect if config is not None else None resolved_dialect = self._build_dialect(dialect_override) statement_expression = self._build_final_expression(copy=True) + self._validate_update_from(statement_expression, resolved_dialect) if self.enable_optimization and isinstance(statement_expression, exp.Expr): statement_expression = self._optimize_expression(statement_expression) diff --git a/sqlspec/builder/_dml.py b/sqlspec/builder/_dml.py index 92545136b..52ff5b28b 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -352,27 +352,86 @@ def get_expression(self) -> exp.Expr | None: ... def set_expression(self, expression: exp.Expr) -> None: ... def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: + """Add a table or subquery to the UPDATE statement's FROM clause. + + Args: + table: Target table name, expression, or builder instance. + alias: Optional alias for the source table or subquery. + + Returns: + The current builder instance for method chaining. + + Raises: + SQLBuilderError: If called on a non-UPDATE expression or with an unsupported table type. + """ 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) - assert current_expr is not None table_expr: exp.Expr if isinstance(table, str): table_expr = exp.to_table(table, alias=alias) - elif isinstance(table, SQLBuilderProtocol): - subquery_params = table.parameters - if subquery_params: - builder_with_params = cast("SQLBuilderProtocol", self) - for param_name, param_value in subquery_params.items(): - builder_with_params.add_parameter(param_value, name=param_name) - raw_expression = table.get_expression() - subquery_source = raw_expression if isinstance(raw_expression, exp.Expr) else exp.select() - subquery_exp = exp.paren(subquery_source) - table_expr = exp.alias_(subquery_exp, alias) if alias else subquery_exp elif isinstance(table, exp.Expr): - table_expr = exp.alias_(table, alias) if alias else table + if isinstance(table, (exp.Select, exp.SetOperation)): + table_expr = exp.Subquery(this=table.copy()) + if alias: + table_expr = exp.alias_(table_expr, alias, table=True) + else: + table_expr = exp.alias_(table.copy(), alias, table=True) if alias else table.copy() + elif ( + hasattr(table, "build") + or hasattr(table, "to_statement") + or hasattr(table, "get_expression") + or hasattr(table, "_expression") + ): + raw_expression = None + if hasattr(table, "_build_final_expression"): + raw_expression = table._build_final_expression(copy=True) + elif hasattr(table, "get_expression"): + raw_expression = table.get_expression() + elif hasattr(table, "_expression"): + raw_expression = table._expression + + if raw_expression is None: + msg = "Subquery builder has no expression to include in FROM clause." + raise SQLBuilderError(msg) + + subquery_copy = raw_expression.copy() if hasattr(raw_expression, "copy") else raw_expression + base_builder = cast("QueryBuilder", self) + builder_alias = getattr(table, "alias_name", None) or getattr(table, "alias", None) + if not isinstance(builder_alias, str): + builder_alias = None + if not builder_alias and hasattr(raw_expression, "alias_or_name"): + builder_alias = raw_expression.alias_or_name + effective_alias = alias or builder_alias or "subquery" + + subquery_params = getattr(table, "parameters", {}) + if subquery_params and isinstance(subquery_params, dict): + param_mapping = base_builder._merge_cte_parameters(effective_alias, subquery_params) + if param_mapping: + subquery_copy = base_builder._update_placeholders(subquery_copy, param_mapping) + + if isinstance(subquery_copy, exp.Values): + if alias: + cols: list[str] = [] + existing_alias = subquery_copy.args.get("alias") + if existing_alias and existing_alias.args.get("columns"): + cols = [c.name for c in existing_alias.args["columns"]] + elif hasattr(table, "columns") and isinstance(table.columns, (list, tuple)): + cols = [str(c) for c in table.columns] + table_expr = exp.alias_(subquery_copy, alias, table=cols or False) + else: + table_expr = subquery_copy + elif isinstance(subquery_copy, exp.Subquery): + table_expr = exp.alias_(subquery_copy, alias, table=True) if alias else subquery_copy + elif isinstance(subquery_copy, (exp.Select, exp.SetOperation)): + table_expr = exp.Subquery(this=subquery_copy) + if alias or builder_alias: + table_expr = exp.alias_(table_expr, alias or builder_alias, table=True) + else: + msg = "UPDATE FROM builder sources must be SELECT, VALUES, or subquery expressions." + raise SQLBuilderError(msg) else: msg = f"Unsupported table type for FROM clause: {type(table)}" raise SQLBuilderError(msg) diff --git a/sqlspec/builder/_factory.py b/sqlspec/builder/_factory.py index 5951ecbc3..e60149d6c 100644 --- a/sqlspec/builder/_factory.py +++ b/sqlspec/builder/_factory.py @@ -56,6 +56,7 @@ ) from sqlspec.builder._select import Case, Select, SubqueryBuilder, WindowFunctionBuilder from sqlspec.builder._update import Update +from sqlspec.builder._values import Values from sqlspec.core import SQL from sqlspec.core.explain import ExplainFormat, ExplainOptions from sqlspec.exceptions import SQLBuilderError @@ -93,6 +94,7 @@ "Select", "Truncate", "Update", + "Values", "WindowFunctionBuilder", "build_copy_from_statement", "build_copy_statement", @@ -402,6 +404,28 @@ def merge(self, table_or_sql: str | None = None, dialect: DialectType = None) -> return Merge(table_or_sql, dialect=builder_dialect) if table_or_sql else Merge(dialect=builder_dialect) + def values( + self, + rows: "Sequence[Sequence[Any] | Mapping[str, Any]]", + *, + alias: str | None = None, + columns: "Sequence[str] | None" = None, + dialect: DialectType = None, + ) -> Values: + """Create a VALUES builder for bulk rows. + + Args: + rows: Sequence of row tuples/lists or mappings. + alias: Optional table alias for the VALUES clause. + columns: Optional column names for the table alias. + dialect: Optional SQL dialect override. + + Returns: + Values builder instance. + """ + builder_dialect = _resolve_dialect(dialect, self.dialect) + return Values(rows, alias=alias, columns=columns, dialect=builder_dialect) + def explain( self, statement: "str | exp.Expr | SQL | SQLBuilderProtocol", diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 0094babbf..8437a9984 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -334,7 +334,11 @@ def from_( elif is_expression(table): from_expr = exp.alias_(table, alias) if alias else table elif has_parameter_builder(table): - subquery_expression = table.get_expression() + subquery_expression = ( + cast("QueryBuilder", table)._build_final_expression(copy=True) + if hasattr(table, "_build_final_expression") + else table.get_expression() + ) if subquery_expression is None: msg = "Subquery builder has no expression to include in FROM clause." raise SQLBuilderError(msg) @@ -345,8 +349,14 @@ def from_( if param_mapping: subquery_copy = base_builder._update_placeholders(subquery_copy, param_mapping) - wrapped_subquery = exp.paren(subquery_copy) - from_expr = exp.alias_(wrapped_subquery, alias) if alias else wrapped_subquery + if isinstance(subquery_copy, exp.Values): + if alias: + columns = getattr(table, "columns", None) + subquery_copy = exp.alias_(subquery_copy, alias, table=columns or True) + from_expr = subquery_copy + else: + wrapped_subquery = exp.paren(subquery_copy) + from_expr = exp.alias_(wrapped_subquery, alias) if alias else wrapped_subquery else: from_expr = table @@ -1124,64 +1134,13 @@ def unpivot( @trait class CommonTableExpressionMixin: - __slots__ = () + """Mixin for Common Table Expression (CTE) support. - def get_expression(self) -> exp.Expr | None: ... - def set_expression(self, expression: exp.Expr) -> None: ... - - _with_ctes: Any - dialect: Any - - def with_(self, name: str, query: Any | str, recursive: bool = False, columns: list[str] | None = None) -> Self: - """Add a CTE via the WITH clause. - - When ``query`` is another builder we reuse its expression, merge parameters with unique names, and let sqlglot handle the actual CTE wrapping to avoid duplicating ``_with_ctes`` state. - """ - builder = cast("QueryBuilder", self) - expression = builder.get_expression() - if expression is None: - msg = "Cannot add WITH clause: expression not initialized." - raise SQLBuilderError(msg) - - if not isinstance(expression, (exp.Select, exp.Insert, exp.Update, exp.Delete)): - msg = f"Cannot add WITH clause to {type(expression).__name__} expression." - raise SQLBuilderError(msg) - - cte_select: exp.Expr | None - if isinstance(query, str): - cte_select = exp.maybe_parse(query, dialect=self.dialect) - elif isinstance(query, exp.Expr): - cte_select = query - else: - cte_select = query.get_expression() - if cte_select is None: - msg = f"Could not get expression from builder: {query}" - raise SQLBuilderError(msg) - - built_query = query.to_statement() - parameters = built_query.parameters - if isinstance(parameters, dict): - param_mapping: dict[str, str] = {} - for param_name, param_value in parameters.items(): - unique_name = builder._next_parameter_name(f"{name}_{param_name}") - param_mapping[param_name] = unique_name - builder.add_parameter(param_value, name=unique_name) - cte_select = builder._update_placeholders(cte_select, param_mapping) - elif isinstance(parameters, (list, tuple)): - for param_value in parameters: - builder.add_parameter(param_value) - elif parameters is not None: - builder.add_parameter(parameters) - - if cte_select is None: - msg = f"Could not parse CTE query: {query}" - raise SQLBuilderError(msg) - - if isinstance(expression, (exp.Select, exp.Insert, exp.Update)): - updated = expression.with_(name, as_=cte_select.copy(), recursive=recursive, copy=True) - builder.set_expression(updated) + CTE functionality is provided directly by :class:`~sqlspec.builder._base.QueryBuilder`. + This class is retained for backward compatibility. + """ - return cast("Self", builder) + __slots__ = () @trait diff --git a/sqlspec/builder/_values.py b/sqlspec/builder/_values.py new file mode 100644 index 000000000..6b23aa45e --- /dev/null +++ b/sqlspec/builder/_values.py @@ -0,0 +1,243 @@ +"""VALUES expression builder. + +Provides a builder interface for constructing SQL VALUES clauses with +parameter binding and optional table aliasing. +""" + +from collections.abc import Mapping, Sequence +from typing import Any + +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.core import SQLResult +from sqlspec.exceptions import SQLBuilderError +from sqlspec.utils.type_guards import has_expression_and_sql + +__all__ = ("Values",) + + +class Values(QueryBuilder): + """Builder for SQL VALUES clauses. + + Constructs parameterized VALUES expressions that can be executed directly, + used as Common Table Expressions (CTEs), or embedded in FROM clauses. + """ + + __slots__ = ("_alias", "_columns", "_rows") + + def __init__( + self, + rows: Sequence[Sequence[Any] | Mapping[str, Any]] | None = None, + *, + alias: str | None = None, + columns: Sequence[str] | None = None, + **kwargs: Any, + ) -> None: + """Initialize a VALUES builder. + + Args: + rows: Sequence of row tuples/lists or mappings. + alias: Optional table alias for the VALUES clause. + columns: Optional column names for the table alias. + **kwargs: Additional QueryBuilder options. + """ + self._init_query_builder(kwargs) + self._alias: str | None = alias + self._columns: list[str] | None = list(columns) if columns is not None else None + self._rows: list[list[Any]] = [] + self._initialize_expression() + + if rows is not None: + self.add_rows(rows) + + def _create_base_expression(self) -> exp.Values: + """Create initial empty VALUES expression.""" + return exp.Values() + + @property + def _expected_result_type(self) -> type[SQLResult]: + """Return expected result type for VALUES queries.""" + return SQLResult + + @property + def alias_name(self) -> str | None: + """Get the table alias name if defined. + + Returns: + The alias name, or None if unaliased. + """ + return self._alias + + @property + def columns(self) -> list[str] | None: + """Get the column names if defined. + + Returns: + List of column names, or None if not defined. + """ + return list(self._columns) if self._columns is not None else None + + def as_(self, alias: str) -> Self: + """Set the table alias for this VALUES clause. + + Args: + alias: Table alias name. + + Returns: + Self for method chaining. + """ + self._alias = alias + if self._rows: + self._rebuild_expression() + return self + + def set_columns(self, *columns: str) -> Self: + """Set the column names for this VALUES clause. + + Args: + *columns: Column names to assign. + + Returns: + Self for method chaining. + + Raises: + SQLBuilderError: If column count does not match row width. + """ + if self._rows and len(columns) != len(self._rows[0]): + msg = f"Column count ({len(columns)}) does not match row width ({len(self._rows[0])})." + raise SQLBuilderError(msg) + self._columns = list(columns) + if self._rows: + self._rebuild_expression() + return self + + def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: + """Add rows to the VALUES clause. + + Args: + rows: Sequence of row tuples, lists, or mappings. + + Returns: + Self for method chaining. + + Raises: + SQLBuilderError: If rows is empty, non-uniform, or column count mismatches. + """ + if not rows: + msg = "VALUES clause requires at least one row." + raise SQLBuilderError(msg) + + columns = self._columns + first_row = rows[0] + if isinstance(first_row, Mapping): + if columns is None: + columns = list(first_row.keys()) + if not columns: + msg = "VALUES clause rows must contain at least one column." + raise SQLBuilderError(msg) + expected_keys = set(columns) + normalized_rows: list[list[Any]] = [] + for idx, r in enumerate(rows): + if not isinstance(r, Mapping): + msg = f"Row {idx} is not a mapping like the initial row." + raise SQLBuilderError(msg) + if set(r.keys()) != expected_keys: + msg = "All rows in VALUES clause must have the same keys as the initial row." + raise SQLBuilderError(msg) + normalized_rows.append([r[k] for k in columns]) + else: + if not isinstance(first_row, Sequence) or isinstance(first_row, (str, bytes, bytearray)): + msg = "VALUES rows must be sequences or mappings." + raise SQLBuilderError(msg) + expected_len = len(first_row) + if expected_len == 0: + msg = "VALUES clause rows must contain at least one column." + raise SQLBuilderError(msg) + normalized_rows = [] + for idx, r in enumerate(rows): + if not isinstance(r, Sequence) or isinstance(r, (str, bytes, bytearray)): + msg = f"Row {idx} must be a sequence." + raise SQLBuilderError(msg) + if len(r) != expected_len: + msg = "All rows in VALUES clause must have the same number of columns." + raise SQLBuilderError(msg) + normalized_rows.append(list(r)) + + if columns is not None and len(columns) != expected_len: + msg = f"Column count ({len(columns)}) does not match row width ({expected_len})." + raise SQLBuilderError(msg) + + if self._rows and len(normalized_rows[0]) != len(self._rows[0]): + msg = "All rows in VALUES clause must have the same number of columns." + raise SQLBuilderError(msg) + self._columns = columns + self._rows.extend(normalized_rows) + self._rebuild_expression() + return self + + def _rebuild_expression(self) -> None: + """Rebuild the underlying sqlglot expression and parameter bindings.""" + cte_parameters = { + str(placeholder.this): self._parameters[str(placeholder.this)] + for cte in self._with_ctes.values() + for placeholder in cte.find_all(exp.Placeholder) + if str(placeholder.this) in self._parameters + } + self._parameters.clear() + self._parameters.update(cte_parameters) + self._parameter_name_counters.clear() + self._parameter_counter = 0 + + tuple_expressions: list[exp.Tuple] = [] + for row in self._rows: + row_expressions: list[exp.Expr] = [] + for col_idx, val in enumerate(row): + if self._columns and col_idx < len(self._columns): + col_name = self._columns[col_idx] + else: + col_name = f"col_{col_idx + 1}" + + if isinstance(val, exp.Expr): + row_expressions.append(val) + elif has_expression_and_sql(val): + row_expressions.append(extract_sql_object_expression(val, builder=self)) + else: + placeholder, _ = self.create_placeholder(val, col_name) + row_expressions.append(placeholder) + tuple_expressions.append(exp.Tuple(expressions=row_expressions)) + + values_expr = exp.Values(expressions=tuple_expressions) + if self._alias: + if self._columns: + self._expression = exp.alias_(values_expr, alias=self._alias, table=self._columns) + else: + self._expression = exp.alias_(values_expr, alias=self._alias) + else: + self._expression = values_expr + + def _build_final_expression(self, *, copy: bool = False) -> exp.Expr: + """Construct the final expression for the VALUES clause. + + Args: + copy: Whether to copy the expression. + + Returns: + SQLGlot expression representing the VALUES clause. + + Raises: + SQLBuilderError: If no rows have been provided. + """ + if not self._rows: + msg = "VALUES clause requires at least one row." + raise SQLBuilderError(msg) + expression = super()._build_final_expression(copy=copy) + with_clause = expression.args.pop("with_", None) + if with_clause is not None: + if not expression.args.get("alias"): + expression = exp.alias_(expression, "_values", table=self._columns or True) + expression = exp.select("*").from_(expression) + expression.set("with_", with_clause) + return expression diff --git a/sqlspec/data_dictionary/_types.py b/sqlspec/data_dictionary/_types.py index d3d54ed5f..f7065b3ab 100644 --- a/sqlspec/data_dictionary/_types.py +++ b/sqlspec/data_dictionary/_types.py @@ -1434,6 +1434,7 @@ class FeatureFlags(TypedDict, total=False): supports_skip_locked: bool supports_structs: bool supports_transactions: bool + supports_update_from: bool supports_upsert: bool supports_uuid: bool supports_window_functions: bool diff --git a/sqlspec/data_dictionary/dialects/bigquery/config.py b/sqlspec/data_dictionary/dialects/bigquery/config.py index f315536a3..2292ae3dc 100644 --- a/sqlspec/data_dictionary/dialects/bigquery/config.py +++ b/sqlspec/data_dictionary/dialects/bigquery/config.py @@ -27,6 +27,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": False, + "supports_update_from": True, } BIGQUERY_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/cockroachdb/config.py b/sqlspec/data_dictionary/dialects/cockroachdb/config.py index cb94ecbc5..2c61546d5 100644 --- a/sqlspec/data_dictionary/dialects/cockroachdb/config.py +++ b/sqlspec/data_dictionary/dialects/cockroachdb/config.py @@ -25,6 +25,7 @@ "supports_skip_locked": True, "supports_crdb_internal_metadata": False, "supports_on_conflict": True, + "supports_update_from": True, } COCKROACHDB_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/duckdb/config.py b/sqlspec/data_dictionary/dialects/duckdb/config.py index a180479a9..ecb88e677 100644 --- a/sqlspec/data_dictionary/dialects/duckdb/config.py +++ b/sqlspec/data_dictionary/dialects/duckdb/config.py @@ -23,6 +23,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": True, + "supports_update_from": True, } DUCKDB_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/mssql/config.py b/sqlspec/data_dictionary/dialects/mssql/config.py index 35563fa89..622e3dea7 100644 --- a/sqlspec/data_dictionary/dialects/mssql/config.py +++ b/sqlspec/data_dictionary/dialects/mssql/config.py @@ -124,6 +124,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": False, + "supports_update_from": True, } MSSQL_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/mysql/config.py b/sqlspec/data_dictionary/dialects/mysql/config.py index 498fdfaa7..39e007fa6 100644 --- a/sqlspec/data_dictionary/dialects/mysql/config.py +++ b/sqlspec/data_dictionary/dialects/mysql/config.py @@ -62,6 +62,7 @@ "supports_sequences": False, "supports_system_versioned_tables": False, "supports_on_conflict": False, + "supports_update_from": False, } MYSQL_TYPE_MAPPINGS: dict[str, str] = { @@ -115,6 +116,7 @@ "supports_invisible_indexes": False, "supports_resource_groups": False, "supports_on_conflict": False, + "supports_update_from": False, } MARIADB_CONFIG = DialectConfig( diff --git a/sqlspec/data_dictionary/dialects/oracle/config.py b/sqlspec/data_dictionary/dialects/oracle/config.py index a87a8006c..5848b5ec9 100644 --- a/sqlspec/data_dictionary/dialects/oracle/config.py +++ b/sqlspec/data_dictionary/dialects/oracle/config.py @@ -31,6 +31,7 @@ ORACLE_MIN_JSON_NATIVE_COMPATIBLE: Final[int] = 20 ORACLE_MIN_JSON_BLOB_VERSION: Final[int] = 12 ORACLE_MIN_OSON_VERSION: Final[int] = 19 +ORACLE_MIN_UPDATE_FROM_VERSION: Final[int] = 23 ORACLE_JSON_STORAGE_NATIVE: Final[str] = "json" ORACLE_JSON_STORAGE_BLOB_JSON: Final[str] = "blob_json" @@ -54,6 +55,7 @@ "supports_for_update": True, "supports_skip_locked": True, "supports_on_conflict": False, + "supports_update_from": False, } ORACLE_TYPE_MAPPINGS: dict[str, str] = { @@ -201,6 +203,8 @@ def resolve_oracle_feature_flag( return bool(version_info and is_autonomous) if version_info is None: return False + if feature == "supports_update_from": + return version_info.major >= ORACLE_MIN_UPDATE_FROM_VERSION if feature == "supports_native_json": return oracle_supports_native_json(version_info.major, compatible_major) if feature == "supports_oson_blob": diff --git a/sqlspec/data_dictionary/dialects/postgres/config.py b/sqlspec/data_dictionary/dialects/postgres/config.py index f3b01942b..f088a51fd 100644 --- a/sqlspec/data_dictionary/dialects/postgres/config.py +++ b/sqlspec/data_dictionary/dialects/postgres/config.py @@ -26,6 +26,7 @@ "supports_schemas": True, "supports_for_update": True, "supports_on_conflict": True, + "supports_update_from": True, } POSTGRES_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/spanner/config.py b/sqlspec/data_dictionary/dialects/spanner/config.py index 8b63db7f6..b6950d7b7 100644 --- a/sqlspec/data_dictionary/dialects/spanner/config.py +++ b/sqlspec/data_dictionary/dialects/spanner/config.py @@ -14,6 +14,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": True, + "supports_update_from": False, } SPANNER_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/sqlite/config.py b/sqlspec/data_dictionary/dialects/sqlite/config.py index 466e764cd..b6acf3782 100644 --- a/sqlspec/data_dictionary/dialects/sqlite/config.py +++ b/sqlspec/data_dictionary/dialects/sqlite/config.py @@ -24,6 +24,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": True, + "supports_update_from": True, } SQLITE_TYPE_MAPPINGS: dict[str, str] = { diff --git a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py new file mode 100644 index 000000000..9a0ac54ec --- /dev/null +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -0,0 +1,96 @@ +"""Integration tests for builder claim statements and VALUES CTE updates on PostgreSQL asyncpg.""" + +from collections.abc import AsyncGenerator + +import pytest +from sqlglot import exp + +from sqlspec import sql +from sqlspec.adapters.asyncpg import AsyncpgDriver + +pytestmark = pytest.mark.xdist_group("postgres") + + +@pytest.fixture +async def asyncpg_tasks_session(asyncpg_async_driver: AsyncpgDriver) -> AsyncGenerator[AsyncpgDriver, None]: + """Create a tasks table with three pending rows for claim testing.""" + await asyncpg_async_driver.execute_script( + """ + CREATE TABLE IF NOT EXISTS test_builder_tasks ( + id TEXT PRIMARY KEY, + status TEXT NOT NULL + ); + TRUNCATE test_builder_tasks; + INSERT INTO test_builder_tasks (id, status) VALUES + ('task-1', 'pending'), + ('task-2', 'pending'), + ('task-3', 'pending'); + """ + ) + try: + yield asyncpg_async_driver + finally: + await asyncpg_async_driver.execute_script("DROP TABLE IF EXISTS test_builder_tasks") + + +async def test_claim_one_row(asyncpg_tasks_session: AsyncpgDriver) -> None: + """Test claiming exactly one row using UPDATE FROM with a subquery, FOR UPDATE SKIP LOCKED, and RETURNING.""" + subquery = ( + sql + .select("id") + .from_("test_builder_tasks") + .where_eq("status", "pending") + .order_by("id") + .limit(1) + .for_update(skip_locked=True) + ) + claim_query = ( + sql + .update("test_builder_tasks") + .set(status="processing") + .from_(subquery, alias="sub") + .where("test_builder_tasks.id = sub.id") + .returning(sql.column("id", table="test_builder_tasks")) + ) + + result = await asyncpg_tasks_session.execute(claim_query) + claimed_rows = result.data + assert len(claimed_rows) == 1 + assert claimed_rows[0][0] == "task-1" + + all_rows = (await asyncpg_tasks_session.execute("SELECT id, status FROM test_builder_tasks ORDER BY id")).data + assert all_rows == [("task-1", "processing"), ("task-2", "pending"), ("task-3", "pending")] + + +async def test_values_cte_bulk_update(asyncpg_tasks_session: AsyncpgDriver) -> None: + """Test bulk updating rows using a sql.values() Common Table Expression.""" + val_cte = sql.values([("task-1", "completed"), ("task-2", "failed")], alias="v", columns=["id", "status"]) + bulk_update = ( + sql + .update("test_builder_tasks") + .with_cte("v", val_cte) + .set(status=exp.column("status", table="v")) + .from_("v") + .where("test_builder_tasks.id = v.id") + ) + + await asyncpg_tasks_session.execute(bulk_update) + all_rows = (await asyncpg_tasks_session.execute("SELECT id, status FROM test_builder_tasks ORDER BY id")).data + assert all_rows == [("task-1", "completed"), ("task-2", "failed"), ("task-3", "pending")] + + +@pytest.mark.parametrize("source_alias", [None, "original"]) +async def test_select_from_values_with_alias_override( + asyncpg_async_driver: AsyncpgDriver, source_alias: str | None +) -> None: + source = sql.values([("first", "alice"), ("second", "bob")], alias=source_alias, columns=["id", "name"]) + query = sql.select("renamed.id", "renamed.name").from_(source, alias="renamed").order_by("renamed.id") + + assert (await asyncpg_async_driver.execute(query)).data == [("first", "alice"), ("second", "bob")] + + +async def test_select_from_values_cte_preserves_column_names(asyncpg_async_driver: AsyncpgDriver) -> None: + source = sql.values([("first", "alice"), ("second", "bob")], columns=["id", "name"]) + query = sql.select("id", "name").from_("v").with_cte("v", source).order_by("id") + + assert (await asyncpg_async_driver.execute(query)).data == [("first", "alice"), ("second", "bob")] diff --git a/tests/integration/adapters/sqlite/test_builder_claim_statement.py b/tests/integration/adapters/sqlite/test_builder_claim_statement.py new file mode 100644 index 000000000..30de7ac40 --- /dev/null +++ b/tests/integration/adapters/sqlite/test_builder_claim_statement.py @@ -0,0 +1,68 @@ +"""Integration tests for builder claim statements and VALUES CTE updates on SQLite.""" + +from collections.abc import Generator + +import pytest +from sqlglot import exp + +from sqlspec import sql +from sqlspec.adapters.sqlite import SqliteConfig, SqliteDriver + +pytestmark = pytest.mark.xdist_group("sqlite") + + +@pytest.fixture +def sqlite_tasks_session() -> Generator[SqliteDriver, None, None]: + """Provide a fresh SQLite in-memory session with a tasks table and three pending rows.""" + config = SqliteConfig(connection_config={"database": ":memory:"}) + try: + with config.provide_session() as driver: + driver.execute_script(""" + CREATE TABLE tasks ( + id INTEGER PRIMARY KEY, + status TEXT NOT NULL + ); + INSERT INTO tasks (id, status) VALUES (1, 'pending'), (2, 'pending'), (3, 'pending'); + """) + yield driver + finally: + config.close_pool() + + +def test_claim_one_row(sqlite_tasks_session: SqliteDriver) -> None: + """Test claiming exactly one row using UPDATE FROM with a subquery and RETURNING.""" + subquery = sql.select("id").from_("tasks").where_eq("status", "pending").order_by("id").limit(1) + claim_query = ( + sql + .update("tasks") + .set(status="processing") + .from_(subquery, alias="sub") + .where("tasks.id = sub.id") + .returning("id") + ) + + result = sqlite_tasks_session.execute(claim_query) + claimed_rows = result.data + assert len(claimed_rows) == 1 + claimed_id = claimed_rows[0][0] + assert claimed_id == 1 + + all_rows = sqlite_tasks_session.execute("SELECT id, status FROM tasks ORDER BY id").data + assert all_rows == [(1, "processing"), (2, "pending"), (3, "pending")] + + +def test_values_cte_bulk_update(sqlite_tasks_session: SqliteDriver) -> None: + """Test bulk updating rows using a sql.values() Common Table Expression.""" + val_cte = sql.values([(1, "completed"), (2, "failed")], alias="v", columns=["id", "status"]) + bulk_update = ( + sql + .update("tasks") + .with_cte("v", val_cte) + .set(status=exp.column("status", table="v")) + .from_("v") + .where("tasks.id = v.id") + ) + + sqlite_tasks_session.execute(bulk_update) + all_rows = sqlite_tasks_session.execute("SELECT id, status FROM tasks ORDER BY id").data + assert all_rows == [(1, "completed"), (2, "failed"), (3, "pending")] diff --git a/tests/unit/builder/test_dml_cte.py b/tests/unit/builder/test_dml_cte.py new file mode 100644 index 000000000..517075e7d --- /dev/null +++ b/tests/unit/builder/test_dml_cte.py @@ -0,0 +1,127 @@ +"""Unit tests for CTE rendering on UPDATE and DELETE statements.""" + +from types import SimpleNamespace + +import pytest +from sqlglot import exp + +from sqlspec import sql +from sqlspec.exceptions import SQLBuilderError + + +def test_update_cte_renders_with_returning() -> None: + """Test that CTEs attached to UPDATE statements render in SQL with RETURNING.""" + cte = sql.select("id").from_("source") + query = sql.update("t").with_cte("c", cte).set(a=1).returning("id") + + stmt = query.build(dialect="postgres") + assert stmt.sql.startswith("WITH") + assert ('"c" AS (' in stmt.sql or "c AS (" in stmt.sql) and "SELECT" in stmt.sql + assert "UPDATE" in stmt.sql + assert "RETURNING" in stmt.sql + assert stmt.parameters["a"] == 1 + + +def test_update_with_alias_renders_cte() -> None: + """Test that with_() method on UPDATE works as an alias for with_cte.""" + cte = sql.select("id").from_("source") + query = sql.update("t").with_("c", cte).set(a=1).returning("id") + + stmt = query.build(dialect="postgres") + assert stmt.sql.startswith("WITH") + assert "UPDATE" in stmt.sql + assert "RETURNING" in stmt.sql + + +def test_delete_cte_renders() -> None: + """Test that CTEs attached to DELETE statements render in SQL with RETURNING.""" + cte = sql.select("id").from_("source") + query = sql.delete("t").with_cte("c", cte).where("t.id in (select id from c)").returning("id") + + stmt = query.build(dialect="postgres") + assert stmt.sql.startswith("WITH") + assert ('"c" AS (' in stmt.sql or "c AS (" in stmt.sql) and "SELECT" in stmt.sql + assert "DELETE FROM" in stmt.sql + assert "RETURNING" in stmt.sql + + +def test_delete_with_alias_renders_cte() -> None: + """Test that with_() method on DELETE works as an alias for with_cte.""" + cte = sql.select("id").from_("source") + query = sql.delete("t").with_("c", cte).where("t.id in (select id from c)").returning("id") + + stmt = query.build(dialect="postgres") + assert stmt.sql.startswith("WITH") + assert "DELETE FROM" in stmt.sql + assert "RETURNING" in stmt.sql + + +def test_cte_parameter_merge_and_collision() -> None: + """Test that CTE parameters merge properly on UPDATE and DELETE with collision handling.""" + cte1 = sql.select("id").from_("x").where_eq("status", "pending") + cte2 = sql.select("id").from_("y").where_eq("status", "archived") + + query = ( + sql + .update("t") + .with_cte("c1", cte1) + .with_cte("c2", cte2) + .set(status="active") + .where("t.id in (select id from c1)") + ) + + stmt = query.build(dialect="postgres") + param_values = list(stmt.parameters.values()) + assert "pending" in param_values + assert "archived" in param_values + assert "active" in param_values + assert len(stmt.parameters) == 3 + + +def test_duplicate_cte_alias_raises_error() -> None: + """Test that registering duplicate CTE aliases on UPDATE raises SQLBuilderError.""" + cte = sql.select("id").from_("x") + query = sql.update("t").with_cte("c", cte) + + with pytest.raises(SQLBuilderError, match=r"CTE with alias 'c' already exists"): + query.with_cte("c", cte) + + +@pytest.mark.parametrize("operation", ["update", "delete", "select"]) +def test_recursive_cte_flag_is_rendered(operation: str) -> None: + query = getattr(sql, operation)("t") + if operation == "update": + query = query.set(a=1) + query = query.with_cte("c", sql.select("id").from_("source"), recursive=True) + assert query.build(dialect="postgres").sql.startswith("WITH RECURSIVE") + + +@pytest.mark.parametrize("final_expression", [False, True]) +def test_cte_accepts_expression_provider_with_bound_values(final_expression: bool) -> None: + expression = exp.values([(exp.Placeholder(this="id"),)], alias="original", columns=["id"]) + source = SimpleNamespace(get_expression=lambda: expression, parameters={"id": 7}, _columns=["id"]) + if final_expression: + source._build_final_expression = lambda **kwargs: expression.copy() + query = sql.select("id").from_("v").with_cte("v", source) + built = query.build(dialect="postgres") + + assert built.parameters == {"v_id": 7} + assert '"v"("id") AS' in built.sql + assert "original" not in built.sql + assert expression.alias == "original" + + +@pytest.mark.parametrize("expression", [None, exp.delete("t")]) +def test_cte_rejects_invalid_expression_provider(expression: exp.Expr | None) -> None: + source = SimpleNamespace(get_expression=lambda: expression) + with pytest.raises(SQLBuilderError, match="CTE 'invalid'"): + sql.select("*").with_cte("invalid", source) + + +def test_cte_preserves_existing_parsed_with_clause() -> None: + query = sql.select("WITH original AS (SELECT 1 AS id) SELECT id FROM original") + query.with_cte("additional", sql.select(exp.Literal.number(2).as_("id"))) + first = query.build(dialect="postgres").sql + assert first.count('"original" AS (') == 1 + assert first.count('"additional" AS (') == 1 + assert query.build(dialect="postgres").sql == first diff --git a/tests/unit/builder/test_sqlglot_arg_contracts.py b/tests/unit/builder/test_sqlglot_arg_contracts.py index e8c395d55..1365c0c5e 100644 --- a/tests/unit/builder/test_sqlglot_arg_contracts.py +++ b/tests/unit/builder/test_sqlglot_arg_contracts.py @@ -29,8 +29,12 @@ KNOWN_SET_SITES: frozenset[tuple[str, str, str]] = frozenset({ ("sqlspec/adapters/bigquery/core.py", "statement_values", "expressions"), ("sqlspec/adapters/duckdb/core.py", "part", "quoted"), - ("sqlspec/builder/_base.py", "final_expression", "with_"), + ("sqlspec/builder/_base.py", "cte_duck_expression", "alias"), + ("sqlspec/builder/_base.py", "cte_expression", "alias"), + ("sqlspec/builder/_base.py", "cte_select_expression", "alias"), ("sqlspec/builder/_base.py", "expression", "conflict"), + ("sqlspec/builder/_base.py", "final_expression", "with_"), + ("sqlspec/builder/_base.py", "final_expression.args['with_']", "recursive"), ("sqlspec/builder/_base.py", "lock", "sqlspec_share_mode"), ("sqlspec/builder/_base.py", "node", "quoted"), ("sqlspec/builder/_base.py", "optimized", "conflict"), @@ -60,6 +64,7 @@ ("sqlspec/builder/_select.py", "table", "pivots"), ("sqlspec/builder/_select.py", "where_clause", "this"), ("sqlspec/builder/_temporal.py", "table_expr", "version"), + ("sqlspec/builder/_values.py", "expression", "with_"), ("sqlspec/core/query_modifiers.py", "existing_where", "this"), ("sqlspec/core/query_modifiers.py", "expression", "expressions"), ("sqlspec/core/query_modifiers.py", "result", "with_"), diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py new file mode 100644 index 000000000..4ff1cd7ee --- /dev/null +++ b/tests/unit/builder/test_update_from.py @@ -0,0 +1,163 @@ +"""Unit tests for UPDATE ... FROM clause support and dialect validation.""" + +from types import SimpleNamespace + +import pytest +from sqlglot import exp + +from sqlspec import sql +from sqlspec.core import StatementConfig +from sqlspec.exceptions import SQLBuilderError + + +def test_update_from_select_builder_matrix() -> None: + """Test UPDATE FROM with Select builder across supported dialects.""" + subquery = sql.select("id").from_("t").limit(1) + query = sql.update("t").set(a=1).from_(subquery, alias="s").where("t.id = s.id").returning("id") + + stmt_pg = query.build(dialect="postgres") + assert "FROM (" in stmt_pg.sql and "SELECT" in stmt_pg.sql + assert "AS s" in stmt_pg.sql or 'AS "s"' in stmt_pg.sql + assert "RETURNING" in stmt_pg.sql + assert stmt_pg.parameters["a"] == 1 + + stmt_sqlite = query.build(dialect="sqlite") + assert "FROM (" in stmt_sqlite.sql and "SELECT" in stmt_sqlite.sql + assert "RETURNING" in stmt_sqlite.sql + assert stmt_sqlite.parameters["a"] == 1 + + stmt_duckdb = query.build(dialect="duckdb") + assert "FROM (" in stmt_duckdb.sql and "SELECT" in stmt_duckdb.sql + assert "RETURNING" in stmt_duckdb.sql + assert stmt_duckdb.parameters["a"] == 1 + + stmt_tsql = query.build(dialect="tsql") + assert "FROM (" in stmt_tsql.sql and "TOP 1" in stmt_tsql.sql + assert "OUTPUT" in stmt_tsql.sql + assert stmt_tsql.parameters["a"] == 1 + + +@pytest.mark.parametrize("dialect", ["oracle", "mysql", "mariadb", "spanner", "spangres"]) +def test_update_from_raises_oracle_mysql(dialect: str) -> None: + """Test UPDATE FROM raises SQLBuilderError on unsupported dialects.""" + subquery = sql.select("id").from_("t").limit(1) + query = sql.update("t").set(a=1).from_(subquery, alias="s").where("t.id = s.id") + + with pytest.raises(SQLBuilderError, match=r"(?i)MERGE|join"): + query.build(dialect=dialect) + + +def test_update_from_string_table() -> None: + """Test UPDATE FROM with string table name and alias.""" + query = sql.update("t").set(a=1).from_("source_table", alias="s").where("t.id = s.id") + stmt = query.build(dialect="postgres") + assert 'FROM "source_table" AS "s"' in stmt.sql or 'FROM "source_table" AS s' in stmt.sql + assert stmt.parameters["a"] == 1 + + +def test_update_from_expression_table() -> None: + """Test UPDATE FROM with sqlglot expression.""" + table_expr = exp.to_table("source_table") + query = sql.update("t").set(a=1).from_(table_expr, alias="s").where("t.id = s.id") + stmt = query.build(dialect="postgres") + assert 'FROM "source_table" AS "s"' in stmt.sql or 'FROM "source_table" AS s' in stmt.sql + + +def test_update_from_parameter_merge_and_collision() -> None: + """Test parameter collision resolution when subquery shares parameter names with main query.""" + subquery = sql.select("id").from_("source").where_eq("status", "pending") + query = sql.update("t").set(status="active").from_(subquery, alias="s").where("t.id = s.id") + stmt = query.build(dialect="postgres") + assert len(stmt.parameters) >= 2 + param_values = list(stmt.parameters.values()) + assert "pending" in param_values + assert "active" in param_values + + +def test_update_from_multiple_sources() -> None: + """Test adding multiple FROM sources creates join clauses.""" + s1 = sql.select("id").from_("src1") + s2 = sql.select("id").from_("src2") + query = sql.update("t").set(a=1).from_(s1, alias="s1").from_(s2, alias="s2").where("t.id = s1.id") + stmt = query.build(dialect="postgres") + assert 'AS "s1"' in stmt.sql and 'AS "s2"' in stmt.sql + + +@pytest.mark.parametrize("dialect", ["mysql", "oracle", "mariadb", "spangres"]) +def test_update_from_statement_config_rejects_unsupported_dialect(dialect: str) -> None: + query = sql.update("t").set(a=1).from_("source") + with pytest.raises(SQLBuilderError, match="MERGE"): + query.to_statement(StatementConfig(dialect=dialect)) + + +def test_update_from_select_expression_is_parenthesized() -> None: + source = exp.select("id").from_("source") + query = sql.update("t").set(a=1).from_(source, alias="s") + expression = query.get_expression() + assert expression is not None + assert isinstance(expression.args["from_"].this, exp.Subquery) + assert source.args.get("alias") is None + + +def test_bigquery_update_from_is_supported() -> None: + query = sql.update("t").set(a=1).from_("source", alias="s").where("t.id = s.id") + assert "FROM" in query.build(dialect="bigquery").sql + assert "FROM" in query.to_statement(StatementConfig(dialect="bigquery")).sql + + +@pytest.mark.parametrize(("major", "expected"), [(19, False), (23, True)]) +def test_oracle_update_from_runtime_capability(major: int, expected: bool) -> None: + from sqlspec.data_dictionary import VersionInfo + from sqlspec.data_dictionary.dialects.oracle.config import ORACLE_CONFIG, resolve_oracle_feature_flag + + assert ( + resolve_oracle_feature_flag( + ORACLE_CONFIG, VersionInfo(major, 0, 0), "supports_update_from", compatible_major=major, is_autonomous=False + ) + is expected + ) + + +@pytest.mark.parametrize("getter", [False, True]) +def test_update_from_expression_provider_preserves_parameters(getter: bool) -> None: + expression = exp.select("id").from_("source").where(exp.column("id").eq(exp.Placeholder(this="id"))) + source = SimpleNamespace(parameters={"id": 7}, alias="candidate") + if getter: + source.get_expression = lambda: expression + else: + source._expression = expression + query = sql.update("target").set(id=8).from_(source).where("target.id = candidate.id") + built = query.build(dialect="postgres") + + assert built.parameters == {"id": 8, "candidate_id": 7} + assert 'AS "candidate"' in built.sql + placeholder = expression.find(exp.Placeholder) + assert placeholder is not None + assert placeholder.name == "id" + + +@pytest.mark.parametrize("expression", [None, exp.delete("source")]) +def test_update_from_rejects_invalid_expression_provider(expression: exp.Expr | None) -> None: + source = SimpleNamespace(get_expression=lambda: expression) + with pytest.raises(SQLBuilderError, match=r"no expression|must be SELECT"): + sql.update("target").set(id=1).from_(source) + + +@pytest.mark.parametrize("source_alias", [None, "original"]) +def test_update_from_values_alias_preserves_columns(source_alias: str | None) -> None: + values = sql.values([(1, "updated")], alias=source_alias, columns=["id", "name"]) + query = sql.update("target").set(name=exp.column("name", table="v")).from_(values, alias="v") + built = query.build(dialect="postgres") + + assert 'AS "v"("id", "name")' in built.sql + assert built.parameters == {"v_id": 1, "v_name": "updated"} + assert values.alias_name == source_alias + + +def test_update_from_subquery_provider_replaces_alias() -> None: + expression = exp.select("id").from_("source").subquery("original") + source = SimpleNamespace(get_expression=lambda: expression) + query = sql.update("target").set(id=1).from_(source, alias="candidate") + + assert 'AS "candidate"' in query.build(dialect="postgres").sql + assert expression.alias == "original" diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py new file mode 100644 index 000000000..4cdcc23e4 --- /dev/null +++ b/tests/unit/builder/test_values.py @@ -0,0 +1,243 @@ +from typing import Any, cast + +import pytest +from sqlglot import exp + +from sqlspec import sql +from sqlspec.builder import Values +from sqlspec.exceptions import SQLBuilderError + + +def test_values_renders_with_parameters() -> None: + """Test that sql.values() produces a parameterized VALUES expression.""" + query = sql.values([(1, "a"), (2, "b")], alias="v", columns=["id", "name"]) + stmt = query.build() + + assert stmt.parameters == {"id": 1, "name": "a", "id_1": 2, "name_1": "b"} + assert "VALUES" in stmt.sql + assert "v" in stmt.sql + assert "id" in stmt.sql + assert "name" in stmt.sql + + +def test_values_renders_dialect_parameters() -> None: + """Test that sql.values() formats dialect-specific parameter placeholders.""" + query = sql.values([(1, "a"), (2, "b")], alias="v", columns=["id", "name"]) + stmt_pg = query.build(dialect="postgres") + + assert "%(id)s" in stmt_pg.sql + assert "%(name)s" in stmt_pg.sql + assert "%(id_1)s" in stmt_pg.sql + assert "%(name_1)s" in stmt_pg.sql + + +def test_values_as_cte_on_select() -> None: + """Test using sql.values() as a Common Table Expression on a SELECT query.""" + val = sql.values([(1, "alice"), (2, "bob")], columns=["id", "name"]) + query = sql.select("*").from_("v").with_cte("v", val) + stmt = query.build(dialect="postgres") + + assert "WITH" in stmt.sql + assert "v" in stmt.sql + assert "VALUES" in stmt.sql + assert "SELECT" in stmt.sql + assert stmt.parameters.get("v_id") == 1 + assert stmt.parameters.get("v_name") == "alice" + assert stmt.parameters.get("v_id_1") == 2 + assert stmt.parameters.get("v_name_1") == "bob" + + +def test_values_as_cte_on_update() -> None: + """Test using sql.values() as a Common Table Expression on an UPDATE query.""" + val = sql.values([(1, "active"), (2, "inactive")], alias="v", columns=["id", "status"]) + query = sql.update("users").with_cte("v", val).set(status="v.status").where("users.id = v.id") + stmt = query.build(dialect="postgres") + + assert "WITH" in stmt.sql + assert "v" in stmt.sql + assert "UPDATE" in stmt.sql + assert "SET" in stmt.sql + assert stmt.parameters.get("v_id") == 1 + assert stmt.parameters.get("v_status") == "active" + assert stmt.parameters.get("v_id_1") == 2 + assert stmt.parameters.get("v_status_1") == "inactive" + + +def test_values_as_update_from_source() -> None: + """Test using sql.values() as an UPDATE ... FROM source table expression.""" + val = sql.values([(1, "alice"), (2, "bob")], alias="v", columns=["id", "name"]) + query = sql.update("users").set(name="v.name").from_(val).where("users.id = v.id") + stmt = query.build(dialect="postgres") + + assert "UPDATE" in stmt.sql + assert "FROM (VALUES" in stmt.sql + assert 'AS "v"' in stmt.sql or "AS v" in stmt.sql + assert "WHERE" in stmt.sql + assert stmt.parameters.get("v_id") == 1 + assert stmt.parameters.get("v_name") == "alice" + assert stmt.parameters.get("v_id_1") == 2 + assert stmt.parameters.get("v_name_1") == "bob" + + +def test_values_from_dict_rows() -> None: + """Test creating sql.values() from mapping/dict rows.""" + query = sql.values([{"id": 1, "name": "a"}, {"id": 2, "name": "b"}], alias="v") + stmt = query.build() + + assert stmt.parameters == {"id": 1, "name": "a", "id_1": 2, "name_1": "b"} + assert "VALUES" in stmt.sql + assert "v" in stmt.sql + + +def test_values_empty_rows_raises() -> None: + """Test that empty row input raises SQLBuilderError.""" + with pytest.raises(SQLBuilderError, match=r"(?i)at least one row"): + sql.values([]) + + with pytest.raises(SQLBuilderError, match=r"(?i)at least one row"): + Values([]) + + +def test_values_ragged_rows_raises() -> None: + """Test that non-uniform row lengths raise SQLBuilderError.""" + with pytest.raises(SQLBuilderError, match=r"(?i)same number of columns"): + sql.values([(1, "a"), (2, "b", "extra")]) + + +def test_values_column_count_mismatch_raises() -> None: + """Test that column count mismatch with row width raises SQLBuilderError.""" + with pytest.raises(SQLBuilderError, match=r"(?i)does not match"): + sql.values([(1, "a")], columns=["id"]) + + +def test_column_named_values_still_works() -> None: + """Test that referencing a column named values via sql.column still works.""" + col = sql.column("values") + assert col.name == "values" + + query = sql.select(col).from_("events") + stmt = query.build() + assert "values" in stmt.sql + + +def test_values_as_and_set_columns() -> None: + """Test as_ and set_columns methods on Values builder.""" + val = Values([(1, "a")]).as_("my_alias").set_columns("col1", "col2") + assert val.alias_name == "my_alias" + assert val.columns == ["col1", "col2"] + + with pytest.raises(SQLBuilderError, match=r"(?i)does not match"): + val.set_columns("only_one") + + +def test_values_expected_result_type() -> None: + """Test expected result type property.""" + from sqlspec.core import SQLResult + + val = Values([(1, "a")]) + assert val._expected_result_type == SQLResult + + +def test_values_alias_without_columns() -> None: + """Test Values with alias but without column list.""" + val = Values([(1, "a")], alias="v") + stmt = val.build() + assert 'AS "v"' in stmt.sql or "AS v" in stmt.sql + + +def test_values_build_empty_raises() -> None: + """Test building an empty Values instance raises SQLBuilderError.""" + val = Values() + with pytest.raises(SQLBuilderError, match=r"(?i)at least one row"): + val.build() + + +def test_values_with_sqlglot_expressions() -> None: + """Test Values containing SQLGlot expressions.""" + from sqlglot import exp + + val = Values([(exp.convert(1), "text")]) + stmt = val.build() + assert "1" in stmt.sql + + +def test_values_add_rows_validation_errors() -> None: + """Test various validation failure branches in add_rows.""" + with pytest.raises(SQLBuilderError, match=r"(?i)not a mapping"): + Values([{"a": 1}, (2,)]) + + with pytest.raises(SQLBuilderError, match=r"(?i)same keys"): + Values([{"a": 1}, {"b": 2}]) + + with pytest.raises(SQLBuilderError, match=r"(?i)must be a sequence"): + Values(cast(Any, [(1, 2), 3])) + + with pytest.raises(SQLBuilderError, match=r"(?i)at least one column"): + Values([()]) + + with pytest.raises(SQLBuilderError, match=r"(?i)must be sequences or mappings"): + Values(cast(Any, [1, 2])) + + +def test_values_rejects_row_width_changes_between_calls() -> None: + values = sql.values([(1, 2)]) + with pytest.raises(SQLBuilderError, match="same number of columns"): + values.add_rows([(3,)]) + assert list(values.build().parameters.values()) == [1, 2] + + +def test_values_rejects_empty_mapping() -> None: + with pytest.raises(SQLBuilderError, match="at least one column"): + sql.values([{}]) + + +def test_values_rebuild_preserves_cte_parameters() -> None: + query = sql.values([(1,)], columns=["id"]).with_cte("c", sql.select("id").from_("t").where_eq("id", 2)) + query.add_rows([(3,)]) + assert sorted(query.build().parameters.values()) == [1, 2, 3] + + +def test_values_accepts_general_row_sequences() -> None: + assert list(sql.values([range(2)]).build().parameters.values()) == [0, 1] + + +def test_values_invalid_mapping_does_not_change_column_state() -> None: + query = Values() + with pytest.raises(SQLBuilderError, match="same keys"): + query.add_rows([{"id": 1}, {"other": 2}]) + query.add_rows([{"other": 3}]) + assert query.columns == ["other"] + assert list(query.build().parameters.values()) == [3] + + +def test_raw_values_cte_moves_alias_columns_without_mutation() -> None: + values = exp.values([(1,)], alias="old", columns=["id"]) + query = sql.select("id").from_("v").with_cte("v", values) + expression = query._build_final_expression(copy=True) + cte = expression.args["with_"].expressions[0] + assert cte.this.args.get("alias") is None + assert [column.name for column in cte.args["alias"].columns] == ["id"] + assert values.alias == "old" + + +def test_select_from_values_preserves_attached_cte() -> None: + source = sql.select(exp.Literal.number(1).as_("id")) + values = sql.values([(exp.Subquery(this=exp.select("id").from_("c")),)], columns=["id"]) + values.with_cte("c", source) + query = sql.select("*").from_(values, alias="v") + assert "WITH" in query.build(dialect="postgres").sql + assert "WITH" in values.build(dialect="postgres").sql + + +@pytest.mark.parametrize("source_alias", [None, "original"]) +def test_select_from_values_replaces_alias_and_preserves_columns(source_alias: str | None) -> None: + source = sql.values([(1, "alice")], alias=source_alias, columns=["id", "name"]) + query = sql.select("renamed.id", "renamed.name").from_(source, alias="renamed") + expression = query._build_final_expression(copy=True) + from_source = expression.args["from_"].this + + assert isinstance(from_source, exp.Values) + assert from_source.alias == "renamed" + assert [column.name for column in from_source.args["alias"].columns] == ["id", "name"] + assert source.alias_name == source_alias + assert query.build(dialect="postgres").parameters == {"renamed_id": 1, "renamed_name": "alice"}