Add support for DML statements (INSERT/UPDATE/DELETE) as CTE bodies

.github/patches/extensions/sqlsmith/fix.patch

@@ -26,6 +26,15 @@ index 72e8f5b..8eafec0 100644
  	}
 
  	return delete_statement;
+@@ -302,7 +303,7 @@ void StatementGenerator::GenerateCTEs(QueryNode &node) {
+ 	}
+ 	while (RandomPercentage(20)) {
+ 		auto cte = make_uniq<CommonTableExpressionInfo>();
+-		cte->query = unique_ptr_cast<SQLStatement, SelectStatement>(GenerateSelect());
++		cte->query_node = std::move(unique_ptr_cast<SQLStatement, SelectStatement>(GenerateSelect())->node);
+ 		for (idx_t i = 0; i < 1 + RandomValue(10); i++) {
+ 			cte->aliases.push_back(GenerateIdentifier());
+ 		}
 @@ -957,9 +958,12 @@ unique_ptr<ParsedExpression> StatementGenerator::GenerateWindowFunction(optional
  		case ExpressionType::WINDOW_NTILE:
  		case ExpressionType::WINDOW_FIRST_VALUE:
@@ -80,6 +89,15 @@ index 4602928..3e56aca 100644
  #include "duckdb/parser/statement/select_statement.hpp"
  #endif
 
+@@ -206,7 +209,7 @@ void StatementSimplifier::Simplify(CommonTableExpressionMap &cte) {
+ 	SimplifyMap(cte.map);
+ 	for (auto &cte_child : cte.map) {
+ 		// simplify individual ctes
+-		Simplify(cte_child.second->query->node);
++		Simplify(cte_child.second->query_node);
+ 	}
+ }
+
 @@ -357,8 +360,6 @@ void StatementSimplifier::SimplifyExpression(duckdb::unique_ptr<ParsedExpression
  		SimplifyChildExpression(expr, window.filter_expr);
  		SimplifyChildExpression(expr, window.start_expr);

extension/autocomplete/transformer/transform_create_view.cpp

@@ -55,9 +55,7 @@ void PEGTransformerFactory::WrapRecursiveView(unique_ptr<CreateViewInfo> &info,
 	auto cte_info = make_uniq<CommonTableExpressionInfo>();
 	cte_info->aliases = info->aliases;
 
-	auto cte_select = make_uniq<SelectStatement>();
-	cte_select->node = std::move(inner_node);
-	cte_info->query = std::move(cte_select);
+	cte_info->query_node = std::move(inner_node);
 
 	outer_select->cte_map.map.insert(info->view_name, std::move(cte_info));
 

extension/autocomplete/transformer/transform_select.cpp

@@ -14,6 +14,12 @@
 #include "duckdb/parser/tableref/at_clause.hpp"
 #include "duckdb/parser/query_node/set_operation_node.hpp"
 #include "duckdb/parser/tableref/pivotref.hpp"
+#include "duckdb/parser/statement/insert_statement.hpp"
+#include "duckdb/parser/statement/update_statement.hpp"
+#include "duckdb/parser/statement/delete_statement.hpp"
+#include "duckdb/parser/query_node/insert_query_node.hpp"
+#include "duckdb/parser/query_node/update_query_node.hpp"
+#include "duckdb/parser/query_node/delete_query_node.hpp"
 
 namespace duckdb {
 
@@ -1597,7 +1603,15 @@ CommonTableExpressionMap PEGTransformerFactory::TransformWithClause(PEGTransform
 		    transformer.Transform<pair<string, unique_ptr<CommonTableExpressionInfo>>>(with_statement_list[entry_idx]);
 
 		if (is_recursive) {
-			auto &query_node = with_entry.second->query->node;
+			auto &query_node = with_entry.second->query_node;
+			if (!query_node) {
+				throw ParserException("Recursive CTEs with DML statements are not supported");
+			}
+			if (query_node->type == QueryNodeType::INSERT_QUERY_NODE ||
+			    query_node->type == QueryNodeType::UPDATE_QUERY_NODE ||
+			    query_node->type == QueryNodeType::DELETE_QUERY_NODE) {
+				throw ParserException("Recursive CTEs with DML statements are not supported");
+			}
 			if (!query_node->modifiers.empty()) {
 				for (auto &modifier : query_node->modifiers) {
 					if (modifier->type == ResultModifierType::LIMIT_MODIFIER ||
@@ -1645,7 +1659,7 @@ PEGTransformerFactory::TransformWithStatement(PEGTransformer &transformer, optio
 	auto table_ref = transformer.Transform<unique_ptr<TableRef>>(list_pr.Child<ListParseResult>(5));
 	D_ASSERT(table_ref->type == TableReferenceType::SUBQUERY);
 	auto subquery_ref = unique_ptr_cast<TableRef, SubqueryRef>(std::move(table_ref));
-	result->query = std::move(subquery_ref->subquery);
+	result->query_node = std::move(subquery_ref->subquery->node);
 	return make_pair(cte_name, std::move(result));
 }
 
@@ -1657,10 +1671,28 @@ unique_ptr<TableRef> PEGTransformerFactory::TransformCTEBody(PEGTransformer &tra
 	// CTEBodyContent <- SelectStatementInternal / Statement
 	auto &content_list = inner->Cast<ListParseResult>();
 	auto &body_choice = content_list.Child<ChoiceParseResult>(0);
-	if (body_choice.result->name != "SelectStatementInternal") {
-		throw ParserException("A CTE needs a SELECT");
+	if (body_choice.result->name == "SelectStatementInternal") {
+		auto select_statement = transformer.Transform<unique_ptr<SelectStatement>>(body_choice.result);
+		return make_uniq<SubqueryRef>(std::move(select_statement));
+	}
+	// DML body (INSERT / UPDATE / DELETE) - transform as a Statement and extract its QueryNode
+	auto sql_stmt = transformer.Transform<unique_ptr<SQLStatement>>(body_choice.result);
+	unique_ptr<QueryNode> query_node;
+	switch (sql_stmt->type) {
+	case StatementType::INSERT_STATEMENT:
+		query_node = unique_ptr_cast<InsertQueryNode, QueryNode>(std::move(sql_stmt->Cast<InsertStatement>().node));
+		break;
+	case StatementType::UPDATE_STATEMENT:
+		query_node = unique_ptr_cast<UpdateQueryNode, QueryNode>(std::move(sql_stmt->Cast<UpdateStatement>().node));
+		break;
+	case StatementType::DELETE_STATEMENT:
+		query_node = unique_ptr_cast<DeleteQueryNode, QueryNode>(std::move(sql_stmt->Cast<DeleteStatement>().node));
+		break;
+	default:
+		throw ParserException("A CTE body must be a SELECT, INSERT, UPDATE, or DELETE statement");
 	}
-	auto select_statement = transformer.Transform<unique_ptr<SelectStatement>>(body_choice.result);
+	auto select_statement = make_uniq<SelectStatement>();
+	select_statement->node = std::move(query_node);
 	return make_uniq<SubqueryRef>(std::move(select_statement));
 }
 

src/include/duckdb/parser/common_table_expression_info.hpp

@@ -13,25 +13,32 @@
 
 namespace duckdb {
 
+class QueryNode;
 class SelectStatement;
+class Serializer;
+class Deserializer;
 
 struct CommonTableExpressionInfo {
+	CommonTableExpressionInfo() = default;
 	~CommonTableExpressionInfo();
 
+	//! Used by deserialization: prefers query_node; if null, uses query->node.
+	CommonTableExpressionInfo(unique_ptr<SelectStatement> query, unique_ptr<QueryNode> query_node);
+
 	vector<string> aliases;
 	vector<unique_ptr<ParsedExpression>> key_targets;
 	vector<unique_ptr<ParsedExpression>> payload_aggregates;
 
-	unique_ptr<SelectStatement> query;
+	//! The root QueryNode for this CTE (SELECT, INSERT, UPDATE, or DELETE)
+	unique_ptr<QueryNode> query_node;
 	CTEMaterialize materialized = CTEMaterialize::CTE_MATERIALIZE_DEFAULT;
 
-public:
+	CTEMaterialize GetMaterializedForSerialization(Serializer &serializer) const;
+	unique_ptr<SelectStatement> GetQueryForSerialization(Serializer &serializer) const;
+
 	void Serialize(Serializer &serializer) const;
 	static unique_ptr<CommonTableExpressionInfo> Deserialize(Deserializer &deserializer);
 	unique_ptr<CommonTableExpressionInfo> Copy();
-
-private:
-	CTEMaterialize GetMaterializedForSerialization(Serializer &serializer) const;
 };
 
 } // namespace duckdb

src/include/duckdb/parser/transformer.hpp

@@ -302,8 +302,8 @@ class Transformer {
 	string TransformAlias(duckdb_libpgquery::PGAlias *root, vector<string> &column_name_alias);
 	vector<string> TransformStringList(duckdb_libpgquery::PGList *list);
 	void TransformCTE(duckdb_libpgquery::PGWithClause &de_with_clause, CommonTableExpressionMap &cte_map);
-	unique_ptr<SelectStatement> TransformRecursiveCTE(duckdb_libpgquery::PGCommonTableExpr &node,
-	                                                  CommonTableExpressionInfo &info);
+	unique_ptr<QueryNode> TransformRecursiveCTE(duckdb_libpgquery::PGCommonTableExpr &node,
+	                                            CommonTableExpressionInfo &info);
 
 	unique_ptr<ParsedExpression> TransformUnaryOperator(const string &op, unique_ptr<ParsedExpression> child);
 	unique_ptr<ParsedExpression> TransformBinaryOperator(string op, unique_ptr<ParsedExpression> left,

src/include/duckdb/storage/serialization/nodes.json

@@ -13,7 +13,10 @@
       {
         "id": 101,
         "name": "query",
-        "type": "SelectStatement*"
+        "type": "SelectStatement*",
+        "version": "v2.0.0",
+        "status": "read_only",
+        "serialize_property": "GetQueryForSerialization(serializer)"
       },
       {
         "id": 102,
@@ -30,8 +33,15 @@
         "id": 104,
         "name": "payload_aggregates",
         "type": "vector<ParsedExpression*>"
+      },
+      {
+        "id": 105,
+        "name": "query_node",
+        "type": "QueryNode*",
+        "version": "v2.0.0"
       }
-    ]
+    ],
+    "constructor": ["query", "query_node"]
   },
   {
     "class": "CommonTableExpressionMap",

src/main/appender.cpp

@@ -483,7 +483,7 @@ unique_ptr<SQLStatement> BaseAppender::ParseStatement(unique_ptr<TableRef> table
 
 	// Create the CTE info.
 	auto cte_info = make_uniq<CommonTableExpressionInfo>();
-	cte_info->query = std::move(cte_select);
+	cte_info->query_node = std::move(cte_select->node);
 	cte_info->materialized = CTEMaterialize::CTE_MATERIALIZE_NEVER;
 
 	// Add the appender data as a CTE to the CTE map of the statement.

src/main/relation/query_relation.cpp

@@ -82,7 +82,7 @@ BoundStatement QueryRelation::Bind(Binder &binder) {
 			select->node = std::move(select_node);
 
 			auto cte_info = make_uniq<CommonTableExpressionInfo>();
-			cte_info->query = std::move(select);
+			cte_info->query_node = std::move(select->node);
 
 			auto subquery = make_uniq<SubqueryRef>(std::move(select_stmt), "query_relation");
 			auto top_level_select = make_uniq<SelectStatement>();

src/optimizer/cte_inlining.cpp

@@ -87,6 +87,23 @@ bool CTEInlining::EndsInAggregateOrDistinct(const LogicalOperator &op) {
 	return false;
 }
 
+static bool HasSideEffects(const LogicalOperator &op) {
+	switch (op.type) {
+	case LogicalOperatorType::LOGICAL_INSERT:
+	case LogicalOperatorType::LOGICAL_UPDATE:
+	case LogicalOperatorType::LOGICAL_DELETE:
+		return true;
+	default:
+		break;
+	}
+	for (auto &child : op.children) {
+		if (HasSideEffects(*child)) {
+			return true;
+		}
+	}
+	return false;
+}
+
 void CTEInlining::TryInlining(unique_ptr<LogicalOperator> &op) {
 	if (op->type == LogicalOperatorType::LOGICAL_PREPARE) {
 		// we are in a prepare statement, if we have to copy an operator during inlining,
@@ -104,6 +121,10 @@ void CTEInlining::TryInlining(unique_ptr<LogicalOperator> &op) {
 		auto &cte = op->Cast<LogicalMaterializedCTE>();
 		auto ref_count = CountCTEReferences(*op, cte.table_index);
 		if (ref_count == 0) {
+			if (HasSideEffects(*cte.children[0])) {
+				// DML CTEs must always execute for side effects even when unreferenced
+				return;
+			}
 			// this CTE is not referenced, we can remove it
 			op = std::move(op->children[1]);
 			return;

src/parser/parsed_expression_iterator.cpp

@@ -372,7 +372,9 @@ void ParsedExpressionIterator::EnumerateQueryNodeChildren(
 	}
 
 	for (auto &kv : node.cte_map.map) {
-		EnumerateQueryNodeChildren(*kv.second->query->node, expr_callback, ref_callback);
+		if (kv.second->query_node) {
+			EnumerateQueryNodeChildren(*kv.second->query_node, expr_callback, ref_callback);
+		}
 	}
 }
 

src/parser/query_node.cpp

@@ -5,10 +5,6 @@
 #include "duckdb/parser/query_node/recursive_cte_node.hpp"
 #include "duckdb/parser/query_node/cte_node.hpp"
 #include "duckdb/common/limits.hpp"
-#include "duckdb/common/serializer/serializer.hpp"
-#include "duckdb/common/serializer/deserializer.hpp"
-#include "duckdb/parser/statement/select_statement.hpp"
-
 namespace duckdb {
 
 CommonTableExpressionMap::CommonTableExpressionMap() {
@@ -27,7 +23,9 @@ CommonTableExpressionMap CommonTableExpressionMap::Copy() const {
 		for (auto &al : kv.second->payload_aggregates) {
 			kv_info->payload_aggregates.push_back(al->Copy());
 		}
-		kv_info->query = unique_ptr_cast<SQLStatement, SelectStatement>(kv.second->query->Copy());
+		if (kv.second->query_node) {
+			kv_info->query_node = kv.second->query_node->Copy();
+		}
 		kv_info->materialized = kv.second->materialized;
 		res.map[kv.first] = std::move(kv_info);
 	}
@@ -42,7 +40,7 @@ string CommonTableExpressionMap::ToString() const {
 	// check if there are any recursive CTEs
 	bool has_recursive = false;
 	for (auto &kv : map) {
-		if (kv.second->query->node->type == QueryNodeType::RECURSIVE_CTE_NODE) {
+		if (kv.second->query_node && kv.second->query_node->type == QueryNodeType::RECURSIVE_CTE_NODE) {
 			has_recursive = true;
 			break;
 		}
@@ -89,7 +87,8 @@ string CommonTableExpressionMap::ToString() const {
 		} else {
 			result += " AS (";
 		}
-		result += cte.query->ToString();
+		D_ASSERT(cte.query_node);
+		result += cte.query_node->ToString();
 		result += ")";
 		first_cte = false;
 	}
@@ -169,7 +168,8 @@ bool QueryNode::Equals(const QueryNode *other) const {
 		if (!ParsedExpression::ListEquals(entry.second->payload_aggregates, other_entry->second->payload_aggregates)) {
 			return false;
 		}
-		if (!entry.second->query->Equals(*other->cte_map.map.at(entry.first)->query)) {
+		if (!entry.second->query_node ||
+		    !entry.second->query_node->Equals(other->cte_map.map.at(entry.first)->query_node.get())) {
 			return false;
 		}
 	}
@@ -191,7 +191,9 @@ void QueryNode::CopyProperties(QueryNode &other) const {
 		for (auto &agg : kv.second->payload_aggregates) {
 			kv_info->payload_aggregates.push_back(agg->Copy());
 		}
-		kv_info->query = unique_ptr_cast<SQLStatement, SelectStatement>(kv.second->query->Copy());
+		if (kv.second->query_node) {
+			kv_info->query_node = kv.second->query_node->Copy();
+		}
 		kv_info->materialized = kv.second->materialized;
 		other.cte_map.map[kv.first] = std::move(kv_info);
 	}