From b26028caa1f0a891d51bead671f54fc2415f4706 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:00:49 +0000 Subject: [PATCH 01/17] feat(builder): allow builder sources in UPDATE FROM and render CTEs on UPDATE and DELETE (#773) --- sqlspec/builder/_base.py | 97 +++++++++++---- sqlspec/builder/_dml.py | 57 +++++++-- sqlspec/builder/_update.py | 20 ++++ sqlspec/data_dictionary/_types.py | 1 + .../dialects/bigquery/config.py | 1 + .../dialects/cockroachdb/config.py | 1 + .../data_dictionary/dialects/duckdb/config.py | 1 + .../data_dictionary/dialects/mssql/config.py | 1 + .../data_dictionary/dialects/mysql/config.py | 2 + .../data_dictionary/dialects/oracle/config.py | 1 + .../dialects/postgres/config.py | 1 + .../dialects/spanner/config.py | 1 + .../data_dictionary/dialects/sqlite/config.py | 1 + tests/unit/builder/test_dml_cte.py | 103 ++++++++++++++++ tests/unit/builder/test_update_from.py | 112 ++++++++++++++++++ 15 files changed, 364 insertions(+), 36 deletions(-) create mode 100644 tests/unit/builder/test_dml_cte.py create mode 100644 tests/unit/builder/test_update_from.py diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 7b3c10f1c..f000664cb 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,13 +318,13 @@ 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) return final_expression @@ -337,34 +339,41 @@ 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.""" - if isinstance(query, QueryBuilder): - query_expr = query.get_expression() + def _resolve_cte_query( + self, alias: str, query: "QueryBuilder | exp.Select | exp.Values | str | Any" + ) -> exp.Expr: + """Resolve a CTE query into a Select or Values expression with merged parameters.""" + if isinstance(query, QueryBuilder) or hasattr(query, "get_expression"): + query_expr = ( + query._build_final_expression(copy=True) + if hasattr(query, "_build_final_expression") + else query.get_expression() + ) 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__}") + if not isinstance(query_expr, (exp.Select, exp.Values)): + self._raise_cte_query_error( + alias, f"expression must be a Select or Values, got {type(query_expr).__name__}" + ) cte_select_expression = query_expr.copy() - 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 hasattr(query, "parameters"): + param_mapping = self._merge_cte_parameters(alias, query.parameters) + if param_mapping: + cte_select_expression = self._update_placeholders(cte_select_expression, param_mapping) + return cte_select_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.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): + if isinstance(query, (exp.Select, exp.Values)): return query self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}") @@ -579,13 +588,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.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 +611,36 @@ 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)) + if columns: + alias_node: exp.Expr = exp.TableAlias( + this=exp.to_identifier(alias), + columns=[exp.to_identifier(c) for c in columns], + ) + else: + alias_node = exp.to_table(alias) + self._with_ctes[alias] = exp.CTE(this=cte_select_expression, alias=alias_node) return self + def with_( + self: Self, + alias: str, + query: "QueryBuilder | exp.Select | exp.Values | str | Any", + recursive: bool = False, + columns: "list[str] | None" = None, + ) -> Self: + """Alias for with_cte for parity across builders. + + Args: + alias: The alias 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(alias, query, recursive=recursive, columns=columns) + def build(self, dialect: DialectType = None) -> "BuiltQuery": """Builds the SQL query string and parameters. @@ -812,6 +856,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) diff --git a/sqlspec/builder/_dml.py b/sqlspec/builder/_dml.py index 92545136b..1cd505970 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -352,27 +352,62 @@ 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 + 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) + subquery_params = getattr(table, "parameters", {}) + if subquery_params and isinstance(subquery_params, dict): + param_mapping = base_builder._merge_cte_parameters(alias or "subquery", subquery_params) + if param_mapping: + subquery_copy = base_builder._update_placeholders(subquery_copy, param_mapping) + + if isinstance(subquery_copy, exp.Values): + table_expr = ( + exp.alias_(subquery_copy, alias) + if alias and not subquery_copy.args.get("alias") + else subquery_copy + ) + else: + table_expr = exp.Subquery(this=subquery_copy, alias=alias) else: msg = f"Unsupported table type for FROM clause: {type(table)}" raise SQLBuilderError(msg) diff --git a/sqlspec/builder/_update.py b/sqlspec/builder/_update.py index 6f76d860f..504b516d1 100644 --- a/sqlspec/builder/_update.py +++ b/sqlspec/builder/_update.py @@ -13,8 +13,10 @@ from sqlspec.builder._dml import UpdateFromClauseMixin, UpdateSetClauseMixin, UpdateTableClauseMixin from sqlspec.builder._explain import ExplainMixin from sqlspec.builder._join import build_join_clause +from sqlspec.builder._parsing_utils import _resolve_dialect from sqlspec.builder._select import ReturningClauseMixin, WhereClauseMixin from sqlspec.core import SQLResult +from sqlspec.data_dictionary import get_dialect_config from sqlspec.exceptions import SQLBuilderError if TYPE_CHECKING: @@ -132,4 +134,22 @@ def build(self, dialect: "DialectType" = None) -> "BuiltQuery": msg = "At least one SET clause must be specified for UPDATE statement." raise SQLBuilderError(msg) + if self._expression.args.get("from_") is not None: + target_dialect = _resolve_dialect(dialect, self.dialect) + dialect_name = ( + getattr(target_dialect, "name", str(target_dialect)) + if target_dialect + else (self.dialect_name or "default") + ) + try: + config = get_dialect_config(dialect_name) + 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) + except ValueError: + pass + return super().build(dialect=dialect) 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..c8faa035f 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": False, } 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..8b8b14002 100644 --- a/sqlspec/data_dictionary/dialects/oracle/config.py +++ b/sqlspec/data_dictionary/dialects/oracle/config.py @@ -54,6 +54,7 @@ "supports_for_update": True, "supports_skip_locked": True, "supports_on_conflict": False, + "supports_update_from": False, } ORACLE_TYPE_MAPPINGS: dict[str, str] = { 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/unit/builder/test_dml_cte.py b/tests/unit/builder/test_dml_cte.py new file mode 100644 index 000000000..abc9c620a --- /dev/null +++ b/tests/unit/builder/test_dml_cte.py @@ -0,0 +1,103 @@ +"""Unit tests for CTE rendering on UPDATE and DELETE statements.""" + +import pytest + +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) diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py new file mode 100644 index 000000000..99e5f89e3 --- /dev/null +++ b/tests/unit/builder/test_update_from.py @@ -0,0 +1,112 @@ +"""Unit tests for UPDATE ... FROM clause support and dialect validation.""" + +import pytest +from sqlglot import exp + +from sqlspec import sql +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", "bigquery"]) +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 From 3efb95146e30be6f939b4861123b4835dd63fbf8 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:08:27 +0000 Subject: [PATCH 02/17] feat(builder): add parameterized sql.values factory and Values builder (#773) --- sqlspec/builder/__init__.py | 2 + sqlspec/builder/_base.py | 62 ++++++--- sqlspec/builder/_dml.py | 24 +++- sqlspec/builder/_factory.py | 24 ++++ sqlspec/builder/_values.py | 221 ++++++++++++++++++++++++++++++ tests/unit/builder/test_values.py | 127 +++++++++++++++++ 6 files changed, 437 insertions(+), 23 deletions(-) create mode 100644 sqlspec/builder/_values.py create mode 100644 tests/unit/builder/test_values.py 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 f000664cb..897b4f9bd 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -343,25 +343,46 @@ def _resolve_cte_query( self, alias: str, query: "QueryBuilder | exp.Select | exp.Values | str | Any" ) -> exp.Expr: """Resolve a CTE query into a Select or Values expression with merged parameters.""" - if isinstance(query, QueryBuilder) or hasattr(query, "get_expression"): - query_expr = ( - query._build_final_expression(copy=True) - if hasattr(query, "_build_final_expression") - else query.get_expression() - ) + if isinstance(query, QueryBuilder): + 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, exp.Values)): self._raise_cte_query_error( alias, f"expression must be a Select or Values, got {type(query_expr).__name__}" ) - cte_select_expression = query_expr.copy() - if hasattr(query, "parameters"): - param_mapping = self._merge_cte_parameters(alias, query.parameters) - if param_mapping: - cte_select_expression = self._update_placeholders(cte_select_expression, param_mapping) + 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) + 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.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) @@ -611,10 +632,19 @@ def with_cte( self._raise_builder_error(f"CTE with alias '{alias}' already exists.") cte_select_expression = self._resolve_cte_query(alias, query) - if columns: + cte_columns = 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 columns], + columns=[exp.to_identifier(c) for c in cte_columns], ) else: alias_node = exp.to_table(alias) @@ -623,7 +653,7 @@ def with_cte( def with_( self: Self, - alias: str, + name: str, query: "QueryBuilder | exp.Select | exp.Values | str | Any", recursive: bool = False, columns: "list[str] | None" = None, @@ -631,7 +661,7 @@ def with_( """Alias for with_cte for parity across builders. Args: - alias: The alias for the CTE. + 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. @@ -639,7 +669,7 @@ def with_( Returns: Self: The current builder instance for method chaining. """ - return self.with_cte(alias, query, recursive=recursive, columns=columns) + return self.with_cte(name, query, recursive=recursive, columns=columns) def build(self, dialect: DialectType = None) -> "BuiltQuery": """Builds the SQL query string and parameters. diff --git a/sqlspec/builder/_dml.py b/sqlspec/builder/_dml.py index 1cd505970..573099469 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -394,20 +394,30 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: 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 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(alias or "subquery", subquery_params) + 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): - table_expr = ( - exp.alias_(subquery_copy, alias) - if alias and not subquery_copy.args.get("alias") - else subquery_copy - ) + 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 else: - table_expr = exp.Subquery(this=subquery_copy, alias=alias) + table_expr = exp.Subquery(this=subquery_copy, alias=alias or builder_alias) 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/_values.py b/sqlspec/builder/_values.py new file mode 100644 index 000000000..216297ab7 --- /dev/null +++ b/sqlspec/builder/_values.py @@ -0,0 +1,221 @@ +"""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) + + first_row = rows[0] + if isinstance(first_row, Mapping): + if self._columns is None: + self._columns = list(first_row.keys()) + expected_keys = set(self._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 self._columns]) + else: + if not isinstance(first_row, (list, tuple)): + 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, (list, tuple)): + 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 self._columns is not None and len(self._columns) != expected_len: + msg = f"Column count ({len(self._columns)}) does not match row width ({expected_len})." + raise SQLBuilderError(msg) + + self._rows.extend(normalized_rows) + self._rebuild_expression() + return self + + def _rebuild_expression(self) -> None: + """Rebuild the underlying sqlglot expression and parameter bindings.""" + self._parameters.clear() + 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) + return super()._build_final_expression(copy=copy) diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py new file mode 100644 index 000000000..46ef7376e --- /dev/null +++ b/tests/unit/builder/test_values.py @@ -0,0 +1,127 @@ +import pytest + +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 From da62f38b8386e37af1a4967a4e3fc3cf9fd98057 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:14:55 +0000 Subject: [PATCH 03/17] docs(builder): document UPDATE FROM, CTEs, and sql.values with integration tests (#773) --- docs/changelog.rst | 9 +++ docs/reference/builder/queries.rst | 7 ++ docs/usage/query_builder.rst | 62 +++++++++++++++ .../asyncpg/test_builder_claim_statement.py | 75 +++++++++++++++++++ .../sqlite/test_builder_claim_statement.py | 66 ++++++++++++++++ 5 files changed, 219 insertions(+) create mode 100644 tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py create mode 100644 tests/integration/adapters/sqlite/test_builder_claim_statement.py diff --git a/docs/changelog.rst b/docs/changelog.rst index b12fef729..403dffc23 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. + (`#773 `_) + **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..ce4b6333a 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 + - Raises ``SQLBuilderError``; use MERGE statement + * - Spanner + - No + - Raises ``SQLBuilderError`` + * - BigQuery + - No + - Raises ``SQLBuilderError``; use MERGE statement + Upserts (ON CONFLICT) --------------------- 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..aa9b0b871 --- /dev/null +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -0,0 +1,75 @@ +"""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") + .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")] 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..30395d87c --- /dev/null +++ b/tests/integration/adapters/sqlite/test_builder_claim_statement.py @@ -0,0 +1,66 @@ +"""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").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")] From 72e526f0dbf6210416afc113bb04afa3241c6a7d Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:31:08 +0000 Subject: [PATCH 04/17] test(builder): add edge-case coverage and type improvements for VALUES and DML CTEs (#773) --- sqlspec/builder/_base.py | 7 +- .../asyncpg/test_builder_claim_statement.py | 12 ++-- .../sqlite/test_builder_claim_statement.py | 6 +- tests/unit/builder/test_dml_cte.py | 31 ++------ tests/unit/builder/test_update_from.py | 44 ++---------- tests/unit/builder/test_values.py | 71 +++++++++++++++---- 6 files changed, 82 insertions(+), 89 deletions(-) diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 897b4f9bd..0187bedb9 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -339,9 +339,7 @@ def _spawn_like_self(self: Self) -> Self: simplify_expressions=self.simplify_expressions, ) - def _resolve_cte_query( - self, alias: str, query: "QueryBuilder | exp.Select | exp.Values | str | Any" - ) -> exp.Expr: + def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | 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._build_final_expression(copy=True) @@ -643,8 +641,7 @@ def with_cte( if cte_columns: alias_node: exp.Expr = exp.TableAlias( - this=exp.to_identifier(alias), - columns=[exp.to_identifier(c) for c in cte_columns], + this=exp.to_identifier(alias), columns=[exp.to_identifier(c) for c in cte_columns] ) else: alias_node = exp.to_table(alias) diff --git a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py index aa9b0b871..c5c917822 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -36,14 +36,11 @@ async def asyncpg_tasks_session(asyncpg_async_driver: AsyncpgDriver) -> AsyncGen 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") - .limit(1) - .for_update(skip_locked=True) + sql.select("id").from_("test_builder_tasks").where_eq("status", "pending").limit(1).for_update(skip_locked=True) ) claim_query = ( - sql.update("test_builder_tasks") + sql + .update("test_builder_tasks") .set(status="processing") .from_(subquery, alias="sub") .where("test_builder_tasks.id = sub.id") @@ -63,7 +60,8 @@ async def test_values_cte_bulk_update(asyncpg_tasks_session: AsyncpgDriver) -> N """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") + sql + .update("test_builder_tasks") .with_cte("v", val_cte) .set(status=exp.column("status", table="v")) .from_("v") diff --git a/tests/integration/adapters/sqlite/test_builder_claim_statement.py b/tests/integration/adapters/sqlite/test_builder_claim_statement.py index 30395d87c..01fbfd841 100644 --- a/tests/integration/adapters/sqlite/test_builder_claim_statement.py +++ b/tests/integration/adapters/sqlite/test_builder_claim_statement.py @@ -33,7 +33,8 @@ 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").limit(1) claim_query = ( - sql.update("tasks") + sql + .update("tasks") .set(status="processing") .from_(subquery, alias="sub") .where("tasks.id = sub.id") @@ -54,7 +55,8 @@ 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") + sql + .update("tasks") .with_cte("v", val_cte) .set(status=exp.column("status", table="v")) .from_("v") diff --git a/tests/unit/builder/test_dml_cte.py b/tests/unit/builder/test_dml_cte.py index abc9c620a..56078531e 100644 --- a/tests/unit/builder/test_dml_cte.py +++ b/tests/unit/builder/test_dml_cte.py @@ -9,12 +9,7 @@ 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") - ) + query = sql.update("t").with_cte("c", cte).set(a=1).returning("id") stmt = query.build(dialect="postgres") assert stmt.sql.startswith("WITH") @@ -27,12 +22,7 @@ def test_update_cte_renders_with_returning() -> None: 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") - ) + query = sql.update("t").with_("c", cte).set(a=1).returning("id") stmt = query.build(dialect="postgres") assert stmt.sql.startswith("WITH") @@ -43,12 +33,7 @@ def test_update_with_alias_renders_cte() -> None: 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") - ) + 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") @@ -60,12 +45,7 @@ def test_delete_cte_renders() -> None: 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") - ) + 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") @@ -79,7 +59,8 @@ def test_cte_parameter_merge_and_collision() -> None: cte2 = sql.select("id").from_("y").where_eq("status", "archived") query = ( - sql.update("t") + sql + .update("t") .with_cte("c1", cte1) .with_cte("c2", cte2) .set(status="active") diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py index 99e5f89e3..b5d3c9a32 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -10,13 +10,7 @@ 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") - ) + 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 @@ -44,12 +38,7 @@ def test_update_from_select_builder_matrix() -> None: 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") - ) + 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) @@ -57,12 +46,7 @@ def test_update_from_raises_oracle_mysql(dialect: str) -> None: 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") - ) + 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 @@ -71,12 +55,7 @@ def test_update_from_string_table() -> None: 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") - ) + 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 @@ -84,12 +63,7 @@ def test_update_from_expression_table() -> None: 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") - ) + 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()) @@ -101,12 +75,6 @@ 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") - ) + 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 diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 46ef7376e..cb37da13f 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -47,12 +47,7 @@ def test_values_as_cte_on_select() -> None: 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") - ) + 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 @@ -68,12 +63,7 @@ def test_values_as_cte_on_update() -> None: 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") - ) + 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 @@ -125,3 +115,60 @@ def test_column_named_values_still_works() -> None: 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([(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([1, 2]) From e8b005e07443421b1d64838ad3d82a4f70f1f375 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:32:23 +0000 Subject: [PATCH 05/17] style(test): format test_values.py (#773) --- tests/unit/builder/test_values.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index cb37da13f..766cb6abf 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -1,3 +1,5 @@ +from typing import Any, cast + import pytest from sqlspec import sql @@ -130,6 +132,7 @@ def test_values_as_and_set_columns() -> None: 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 @@ -151,6 +154,7 @@ def test_values_build_empty_raises() -> None: 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 @@ -165,10 +169,10 @@ def test_values_add_rows_validation_errors() -> None: Values([{"a": 1}, {"b": 2}]) with pytest.raises(SQLBuilderError, match=r"(?i)must be a sequence"): - Values([(1, 2), 3]) + 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([1, 2]) + Values(cast(Any, [1, 2])) From b80afc626f3f72cd0a2d00009dfa6292f32e8953 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:34:41 +0000 Subject: [PATCH 06/17] test(builder): register cte expression alias in KNOWN_SET_SITES (#773) --- tests/unit/builder/test_sqlglot_arg_contracts.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/unit/builder/test_sqlglot_arg_contracts.py b/tests/unit/builder/test_sqlglot_arg_contracts.py index e8c395d55..f4b882632 100644 --- a/tests/unit/builder/test_sqlglot_arg_contracts.py +++ b/tests/unit/builder/test_sqlglot_arg_contracts.py @@ -29,8 +29,9 @@ 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", "cte_duck_expression", "alias"), + ("sqlspec/builder/_base.py", "cte_select_expression", "alias"), ("sqlspec/builder/_base.py", "final_expression", "with_"), - ("sqlspec/builder/_base.py", "expression", "conflict"), ("sqlspec/builder/_base.py", "lock", "sqlspec_share_mode"), ("sqlspec/builder/_base.py", "node", "quoted"), ("sqlspec/builder/_base.py", "optimized", "conflict"), From 24dd188b4e5ca5dc1a51e3cf47b84cd90bc154d3 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:43:57 +0000 Subject: [PATCH 07/17] fix: preserve CTE semantics and validate DML sources --- sqlspec/builder/_base.py | 36 +++++++++++++++++++++----- sqlspec/builder/_dml.py | 18 +++++++++++-- sqlspec/builder/_update.py | 20 -------------- sqlspec/builder/_values.py | 6 +++++ tests/unit/builder/test_dml_cte.py | 9 +++++++ tests/unit/builder/test_update_from.py | 17 ++++++++++++ tests/unit/builder/test_values.py | 12 +++++++++ 7 files changed, 89 insertions(+), 29 deletions(-) diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 0187bedb9..0d1bd2374 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -326,6 +326,8 @@ def _build_final_expression(self, *, copy: bool = False) -> exp.Expr: 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: @@ -339,13 +341,13 @@ def _spawn_like_self(self: Self) -> Self: simplify_expressions=self.simplify_expressions, ) - def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | exp.Values | str | Any") -> exp.Expr: + 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._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, exp.Values)): + 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__}" ) @@ -367,7 +369,7 @@ def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | exp ) 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.Values)): + 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__}" ) @@ -386,13 +388,13 @@ def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | exp 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, exp.Values)): + if not isinstance(parsed_expression, (exp.Select, exp.SetOperation, exp.Values)): self._raise_cte_query_error( alias, f"query string must parse to SELECT or VALUES, got {type(parsed_expression).__name__}" ) return parsed_expression - if isinstance(query, (exp.Select, exp.Values)): + if isinstance(query, (exp.Select, exp.SetOperation, exp.Values)): return query self._raise_cte_query_error(alias, f"invalid query type: {type(query).__name__}") @@ -610,7 +612,7 @@ def _cache_key(self, config: "StatementConfig | None" = None) -> str: def with_cte( self: Self, alias: str, - query: "QueryBuilder | exp.Select | exp.Values | str | Any", + query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any", recursive: bool = False, columns: "list[str] | None" = None, ) -> Self: @@ -646,12 +648,13 @@ def with_cte( 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.Values | str | Any", + query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any", recursive: bool = False, columns: "list[str] | None" = None, ) -> Self: @@ -679,6 +682,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) @@ -988,10 +992,28 @@ 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(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 573099469..b539dd346 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -373,7 +373,12 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: if isinstance(table, str): table_expr = exp.to_table(table, alias=alias) elif isinstance(table, exp.Expr): - table_expr = exp.alias_(table, alias) if alias else table + if isinstance(table, exp.Select): + 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") @@ -395,6 +400,8 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: 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" @@ -416,8 +423,15 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: 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): + 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: - table_expr = exp.Subquery(this=subquery_copy, alias=alias or builder_alias) + 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/_update.py b/sqlspec/builder/_update.py index 504b516d1..6f76d860f 100644 --- a/sqlspec/builder/_update.py +++ b/sqlspec/builder/_update.py @@ -13,10 +13,8 @@ from sqlspec.builder._dml import UpdateFromClauseMixin, UpdateSetClauseMixin, UpdateTableClauseMixin from sqlspec.builder._explain import ExplainMixin from sqlspec.builder._join import build_join_clause -from sqlspec.builder._parsing_utils import _resolve_dialect from sqlspec.builder._select import ReturningClauseMixin, WhereClauseMixin from sqlspec.core import SQLResult -from sqlspec.data_dictionary import get_dialect_config from sqlspec.exceptions import SQLBuilderError if TYPE_CHECKING: @@ -134,22 +132,4 @@ def build(self, dialect: "DialectType" = None) -> "BuiltQuery": msg = "At least one SET clause must be specified for UPDATE statement." raise SQLBuilderError(msg) - if self._expression.args.get("from_") is not None: - target_dialect = _resolve_dialect(dialect, self.dialect) - dialect_name = ( - getattr(target_dialect, "name", str(target_dialect)) - if target_dialect - else (self.dialect_name or "default") - ) - try: - config = get_dialect_config(dialect_name) - 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) - except ValueError: - pass - return super().build(dialect=dialect) diff --git a/sqlspec/builder/_values.py b/sqlspec/builder/_values.py index 216297ab7..1df8ae9c8 100644 --- a/sqlspec/builder/_values.py +++ b/sqlspec/builder/_values.py @@ -134,6 +134,9 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: if isinstance(first_row, Mapping): if self._columns is None: self._columns = list(first_row.keys()) + if not self._columns: + msg = "VALUES clause rows must contain at least one column." + raise SQLBuilderError(msg) expected_keys = set(self._columns) normalized_rows: list[list[Any]] = [] for idx, r in enumerate(rows): @@ -166,6 +169,9 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: msg = f"Column count ({len(self._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._rows.extend(normalized_rows) self._rebuild_expression() return self diff --git a/tests/unit/builder/test_dml_cte.py b/tests/unit/builder/test_dml_cte.py index 56078531e..181d28422 100644 --- a/tests/unit/builder/test_dml_cte.py +++ b/tests/unit/builder/test_dml_cte.py @@ -82,3 +82,12 @@ def test_duplicate_cte_alias_raises_error() -> None: 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") diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py index b5d3c9a32..7dc0bf530 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -5,6 +5,7 @@ from sqlspec import sql from sqlspec.exceptions import SQLBuilderError +from sqlspec.core import StatementConfig def test_update_from_select_builder_matrix() -> None: @@ -78,3 +79,19 @@ def test_update_from_multiple_sources() -> None: 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"]) +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 diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 766cb6abf..1b23766f4 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -176,3 +176,15 @@ def test_values_add_rows_validation_errors() -> None: 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([{}]) From 8a148c418be08df1eafaa1f7e64b5146304e7a49 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:45:49 +0000 Subject: [PATCH 08/17] fix: align update capabilities and preserve values CTE bindings --- docs/usage/query_builder.rst | 6 ++--- sqlspec/builder/_values.py | 7 ++++++ .../dialects/bigquery/config.py | 2 +- .../data_dictionary/dialects/oracle/config.py | 3 +++ tests/unit/builder/test_update_from.py | 22 ++++++++++++++++++- tests/unit/builder/test_values.py | 6 +++++ 6 files changed, 41 insertions(+), 5 deletions(-) diff --git a/docs/usage/query_builder.rst b/docs/usage/query_builder.rst index ce4b6333a..2120ce8f3 100644 --- a/docs/usage/query_builder.rst +++ b/docs/usage/query_builder.rst @@ -92,13 +92,13 @@ Dialect Support Matrix - Raises ``SQLBuilderError``; use multi-table join update or MERGE * - Oracle - No - - Raises ``SQLBuilderError``; use MERGE statement + - Builder conservatively raises ``SQLBuilderError``; use MERGE or raw SQL on Oracle 23+ * - Spanner - No - Raises ``SQLBuilderError`` * - BigQuery - - No - - Raises ``SQLBuilderError``; use MERGE statement + - Yes + - Native ``UPDATE ... FROM``; a ``WHERE`` condition is required Upserts (ON CONFLICT) --------------------- diff --git a/sqlspec/builder/_values.py b/sqlspec/builder/_values.py index 1df8ae9c8..6c6007dac 100644 --- a/sqlspec/builder/_values.py +++ b/sqlspec/builder/_values.py @@ -178,7 +178,14 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> 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 diff --git a/sqlspec/data_dictionary/dialects/bigquery/config.py b/sqlspec/data_dictionary/dialects/bigquery/config.py index c8faa035f..2292ae3dc 100644 --- a/sqlspec/data_dictionary/dialects/bigquery/config.py +++ b/sqlspec/data_dictionary/dialects/bigquery/config.py @@ -27,7 +27,7 @@ "supports_for_update": False, "supports_skip_locked": False, "supports_on_conflict": False, - "supports_update_from": False, + "supports_update_from": True, } BIGQUERY_TYPE_MAPPINGS: dict[str, str] = { diff --git a/sqlspec/data_dictionary/dialects/oracle/config.py b/sqlspec/data_dictionary/dialects/oracle/config.py index 8b8b14002..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" @@ -202,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/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py index 7dc0bf530..57ce96994 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -35,7 +35,7 @@ def test_update_from_select_builder_matrix() -> None: assert stmt_tsql.parameters["a"] == 1 -@pytest.mark.parametrize("dialect", ["oracle", "mysql", "mariadb", "spanner", "bigquery"]) +@pytest.mark.parametrize("dialect", ["oracle", "mysql", "mariadb", "spanner"]) 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) @@ -95,3 +95,23 @@ def test_update_from_select_expression_is_parenthesized() -> None: 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 diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 1b23766f4..26255db11 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -188,3 +188,9 @@ def test_values_rejects_row_width_changes_between_calls() -> None: 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] From 3069b123a6f62265b4c21fa5d498d0803dcae4a9 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:51:43 +0000 Subject: [PATCH 09/17] fix: retain values CTEs in composed selects --- sqlspec/builder/_base.py | 17 ++++++++-- sqlspec/builder/_dml.py | 4 +-- sqlspec/builder/_select.py | 6 +++- sqlspec/builder/_values.py | 29 ++++++++++------ .../asyncpg/test_builder_claim_statement.py | 8 ++++- .../sqlite/test_builder_claim_statement.py | 2 +- .../builder/test_sqlglot_arg_contracts.py | 4 ++- tests/unit/builder/test_update_from.py | 15 ++++---- tests/unit/builder/test_values.py | 34 +++++++++++++++++++ 9 files changed, 92 insertions(+), 27 deletions(-) diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 0d1bd2374..467aecf3e 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -341,7 +341,9 @@ def _spawn_like_self(self: Self) -> Self: simplify_expressions=self.simplify_expressions, ) - def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | exp.SetOperation | exp.Values | str | Any") -> exp.Expr: + 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._build_final_expression(copy=True) @@ -395,7 +397,10 @@ def _resolve_cte_query(self, alias: str, query: "QueryBuilder | exp.Select | exp return parsed_expression if isinstance(query, (exp.Select, exp.SetOperation, exp.Values)): - return query + 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" @@ -633,6 +638,10 @@ def with_cte( cte_select_expression = self._resolve_cte_query(alias, query) 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) @@ -997,7 +1006,9 @@ def _validate_update_from(self, expression: exp.Expr, dialect: DialectType) -> N 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() + dialect_name = ( + dialect.lower() if isinstance(dialect, str) else type(Dialect.get_or_raise(dialect)).__name__.lower() + ) try: config = get_dialect_config(dialect_name) except ValueError: diff --git a/sqlspec/builder/_dml.py b/sqlspec/builder/_dml.py index b539dd346..52ff5b28b 100644 --- a/sqlspec/builder/_dml.py +++ b/sqlspec/builder/_dml.py @@ -373,7 +373,7 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: if isinstance(table, str): table_expr = exp.to_table(table, alias=alias) elif isinstance(table, exp.Expr): - if isinstance(table, exp.Select): + 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) @@ -425,7 +425,7 @@ def from_(self, table: str | exp.Expr | Any, alias: str | None = None) -> Self: 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): + 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) diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 0094babbf..48f4f0377 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) diff --git a/sqlspec/builder/_values.py b/sqlspec/builder/_values.py index 6c6007dac..6b23aa45e 100644 --- a/sqlspec/builder/_values.py +++ b/sqlspec/builder/_values.py @@ -130,14 +130,15 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: msg = "VALUES clause requires at least one row." raise SQLBuilderError(msg) + columns = self._columns first_row = rows[0] if isinstance(first_row, Mapping): - if self._columns is None: - self._columns = list(first_row.keys()) - if not self._columns: + 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(self._columns) + expected_keys = set(columns) normalized_rows: list[list[Any]] = [] for idx, r in enumerate(rows): if not isinstance(r, Mapping): @@ -146,9 +147,9 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: 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 self._columns]) + normalized_rows.append([r[k] for k in columns]) else: - if not isinstance(first_row, (list, tuple)): + 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) @@ -157,7 +158,7 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: raise SQLBuilderError(msg) normalized_rows = [] for idx, r in enumerate(rows): - if not isinstance(r, (list, tuple)): + 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: @@ -165,13 +166,14 @@ def add_rows(self, rows: Sequence[Sequence[Any] | Mapping[str, Any]]) -> Self: raise SQLBuilderError(msg) normalized_rows.append(list(r)) - if self._columns is not None and len(self._columns) != expected_len: - msg = f"Column count ({len(self._columns)}) does not match row width ({expected_len})." + 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 @@ -231,4 +233,11 @@ def _build_final_expression(self, *, copy: bool = False) -> exp.Expr: if not self._rows: msg = "VALUES clause requires at least one row." raise SQLBuilderError(msg) - return super()._build_final_expression(copy=copy) + 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/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py index c5c917822..ab84532a0 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -36,7 +36,13 @@ async def asyncpg_tasks_session(asyncpg_async_driver: AsyncpgDriver) -> AsyncGen 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").limit(1).for_update(skip_locked=True) + sql + .select("id") + .from_("test_builder_tasks") + .where_eq("status", "pending") + .order_by("id") + .limit(1) + .for_update(skip_locked=True) ) claim_query = ( sql diff --git a/tests/integration/adapters/sqlite/test_builder_claim_statement.py b/tests/integration/adapters/sqlite/test_builder_claim_statement.py index 01fbfd841..30de7ac40 100644 --- a/tests/integration/adapters/sqlite/test_builder_claim_statement.py +++ b/tests/integration/adapters/sqlite/test_builder_claim_statement.py @@ -31,7 +31,7 @@ def sqlite_tasks_session() -> Generator[SqliteDriver, None, None]: 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").limit(1) + subquery = sql.select("id").from_("tasks").where_eq("status", "pending").order_by("id").limit(1) claim_query = ( sql .update("tasks") diff --git a/tests/unit/builder/test_sqlglot_arg_contracts.py b/tests/unit/builder/test_sqlglot_arg_contracts.py index f4b882632..6f28161b8 100644 --- a/tests/unit/builder/test_sqlglot_arg_contracts.py +++ b/tests/unit/builder/test_sqlglot_arg_contracts.py @@ -30,8 +30,9 @@ ("sqlspec/adapters/bigquery/core.py", "statement_values", "expressions"), ("sqlspec/adapters/duckdb/core.py", "part", "quoted"), ("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", "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"), @@ -61,6 +62,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 index 57ce96994..375d8dd22 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -4,8 +4,8 @@ from sqlglot import exp from sqlspec import sql -from sqlspec.exceptions import SQLBuilderError from sqlspec.core import StatementConfig +from sqlspec.exceptions import SQLBuilderError def test_update_from_select_builder_matrix() -> None: @@ -108,10 +108,9 @@ def test_oracle_update_from_runtime_capability(major: int, expected: bool) -> No 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 + assert ( + resolve_oracle_feature_flag( + ORACLE_CONFIG, VersionInfo(major, 0, 0), "supports_update_from", compatible_major=major, is_autonomous=False + ) + is expected + ) diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 26255db11..869bf1e3b 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -1,5 +1,7 @@ from typing import Any, cast +from sqlglot import exp + import pytest from sqlspec import sql @@ -194,3 +196,35 @@ 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 From c9268e1fce8a7b3411eede7f7c346645407b36d1 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:53:36 +0000 Subject: [PATCH 10/17] fix: link DML changelog to its pull request --- docs/changelog.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/changelog.rst b/docs/changelog.rst index 403dffc23..68cfc01be 100644 --- a/docs/changelog.rst +++ b/docs/changelog.rst @@ -213,7 +213,7 @@ v0.63.0 - Transactions, table fixtures, SQL fragments, storage, and kwargs param (`#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. - (`#773 `_) + (`#779 `_) **Fixed:** From 9e4a3587d7d944fbb706257e6527c2229e581891 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:56:15 +0000 Subject: [PATCH 11/17] fix: validate update sources for Spanner PostgreSQL dialect --- sqlspec/builder/_base.py | 2 +- tests/unit/builder/test_update_from.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index 467aecf3e..dab246708 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -1010,7 +1010,7 @@ def _validate_update_from(self, expression: exp.Expr, dialect: DialectType) -> N dialect.lower() if isinstance(dialect, str) else type(Dialect.get_or_raise(dialect)).__name__.lower() ) try: - config = get_dialect_config(dialect_name) + 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): diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py index 375d8dd22..e8833e43f 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -35,7 +35,7 @@ def test_update_from_select_builder_matrix() -> None: assert stmt_tsql.parameters["a"] == 1 -@pytest.mark.parametrize("dialect", ["oracle", "mysql", "mariadb", "spanner"]) +@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) @@ -81,7 +81,7 @@ def test_update_from_multiple_sources() -> None: assert "AS s1" in stmt.sql and "AS s2" in stmt.sql -@pytest.mark.parametrize("dialect", ["mysql", "oracle", "mariadb"]) +@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"): From c0dc80ab00d6e3d424d093c0e483c3c1b6a7bd56 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 16:57:59 +0000 Subject: [PATCH 12/17] style: normalize values regression imports --- tests/unit/builder/test_values.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index 869bf1e3b..e28f8765e 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -1,8 +1,7 @@ from typing import Any, cast -from sqlglot import exp - import pytest +from sqlglot import exp from sqlspec import sql from sqlspec.builder import Values From 6ffc658daeea1f6c415c46079e3d4e1788d09733 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 18:08:10 +0000 Subject: [PATCH 13/17] fix(builder): assert quoted aliases in update from postgres test --- tests/unit/builder/test_update_from.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/unit/builder/test_update_from.py b/tests/unit/builder/test_update_from.py index e8833e43f..75a982304 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -78,7 +78,7 @@ def test_update_from_multiple_sources() -> None: 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 + assert 'AS "s1"' in stmt.sql and 'AS "s2"' in stmt.sql @pytest.mark.parametrize("dialect", ["mysql", "oracle", "mariadb", "spangres"]) From df00ca0f4cf1eb5bb3122c396aff2f376ec5a8cf Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 18:38:09 +0000 Subject: [PATCH 14/17] fix(builder): resolve CommonTableExpressionMixin mypyc vtable conflict --- sqlspec/builder/_select.py | 61 ++++---------------------------------- 1 file changed, 5 insertions(+), 56 deletions(-) diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index 48f4f0377..fec4fbd0f 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -1128,64 +1128,13 @@ def unpivot( @trait class CommonTableExpressionMixin: - __slots__ = () - - 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) + """Mixin for Common Table Expression (CTE) support. - 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 From 056c14ebffa1fd229a34d2e6a48314bebfdb0dcd Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 19:22:23 +0000 Subject: [PATCH 15/17] fix(builder): preserve VALUES columns when replacing source aliases --- sqlspec/builder/_select.py | 10 ++++++++-- .../asyncpg/test_builder_claim_statement.py | 10 ++++++++++ tests/unit/builder/test_values.py | 14 ++++++++++++++ 3 files changed, 32 insertions(+), 2 deletions(-) diff --git a/sqlspec/builder/_select.py b/sqlspec/builder/_select.py index fec4fbd0f..8437a9984 100644 --- a/sqlspec/builder/_select.py +++ b/sqlspec/builder/_select.py @@ -349,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 diff --git a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py index ab84532a0..1a816d75a 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -77,3 +77,13 @@ async def test_values_cte_bulk_update(asyncpg_tasks_session: AsyncpgDriver) -> N 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")] diff --git a/tests/unit/builder/test_values.py b/tests/unit/builder/test_values.py index e28f8765e..4cdcc23e4 100644 --- a/tests/unit/builder/test_values.py +++ b/tests/unit/builder/test_values.py @@ -227,3 +227,17 @@ def test_select_from_values_preserves_attached_cte() -> None: 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"} From d503550b08ae02593f0f997d334afff4ca7cd154 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 19:31:50 +0000 Subject: [PATCH 16/17] fix(builder): preserve named VALUES CTEs during optimization --- sqlspec/builder/_base.py | 4 ++ .../asyncpg/test_builder_claim_statement.py | 7 +++ tests/unit/builder/test_dml_cte.py | 34 ++++++++++++++ tests/unit/builder/test_update_from.py | 47 +++++++++++++++++++ 4 files changed, 92 insertions(+) diff --git a/sqlspec/builder/_base.py b/sqlspec/builder/_base.py index dab246708..80eb74eb6 100644 --- a/sqlspec/builder/_base.py +++ b/sqlspec/builder/_base.py @@ -889,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) diff --git a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py index 1a816d75a..9a0ac54ec 100644 --- a/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py +++ b/tests/integration/adapters/postgres/asyncpg/test_builder_claim_statement.py @@ -87,3 +87,10 @@ async def test_select_from_values_with_alias_override( 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/unit/builder/test_dml_cte.py b/tests/unit/builder/test_dml_cte.py index 181d28422..517075e7d 100644 --- a/tests/unit/builder/test_dml_cte.py +++ b/tests/unit/builder/test_dml_cte.py @@ -1,6 +1,9 @@ """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 @@ -91,3 +94,34 @@ def test_recursive_cte_flag_is_rendered(operation: str) -> None: 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_update_from.py b/tests/unit/builder/test_update_from.py index 75a982304..4ff1cd7ee 100644 --- a/tests/unit/builder/test_update_from.py +++ b/tests/unit/builder/test_update_from.py @@ -1,5 +1,7 @@ """Unit tests for UPDATE ... FROM clause support and dialect validation.""" +from types import SimpleNamespace + import pytest from sqlglot import exp @@ -114,3 +116,48 @@ def test_oracle_update_from_runtime_capability(major: int, expected: bool) -> No ) 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" From 17a6f2e6ef013c6d32b2e702c000696619c7a083 Mon Sep 17 00:00:00 2001 From: Cody Fincher Date: Sun, 13 Sep 2026 20:13:38 +0000 Subject: [PATCH 17/17] fix(test): restore known set sites in arg contracts suite --- tests/unit/builder/test_sqlglot_arg_contracts.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/unit/builder/test_sqlglot_arg_contracts.py b/tests/unit/builder/test_sqlglot_arg_contracts.py index 6f28161b8..1365c0c5e 100644 --- a/tests/unit/builder/test_sqlglot_arg_contracts.py +++ b/tests/unit/builder/test_sqlglot_arg_contracts.py @@ -32,6 +32,8 @@ ("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"),