From 6810cfc52d844e4dd9f0dbdfe5cf0e2df1111413 Mon Sep 17 00:00:00 2001
From: tibotix <tizian@seehaus.net>
Date: Fri, 17 Nov 2023 18:15:23 +0100
Subject: [PATCH] Refactor `TypeResolver.visitMultiColumnSetComponent`

---
 .../src/analysis/types/resolving_visitor.dart | 21 ++++++++++++-------
 sqlparser/lib/src/ast/statements/update.dart  |  5 +++++
 2 files changed, 18 insertions(+), 8 deletions(-)

diff --git a/sqlparser/lib/src/analysis/types/resolving_visitor.dart b/sqlparser/lib/src/analysis/types/resolving_visitor.dart
index b1dc77cf..c3e843bd 100644
--- a/sqlparser/lib/src/analysis/types/resolving_visitor.dart
+++ b/sqlparser/lib/src/analysis/types/resolving_visitor.dart
@@ -142,15 +142,20 @@ class TypeResolver extends RecursiveVisitor<TypeExpectation, void> {
   void visitMultiColumnSetComponent(
       MultiColumnSetComponent e, TypeExpectation arg) {
     visitList(e.columns, const NoTypeExpectation());
-    if (e.rowValue is Tuple) {
-      final expressions = (e.rowValue as Tuple).expressions;
-      for (final (idx, expression) in expressions.indexed) {
-        _lazyCopy(expression, e.columns.elementAtOrNull(idx));
-      }
-    } else if (e.rowValue is SubQuery) {
-      // TODO: handle subquery case
+
+    final targets = e.resolvedTargetColumns ?? const [];
+    for (final column in targets) {
+      _handleColumn(column, e);
     }
-    visit(e.rowValue, const NoTypeExpectation());
+
+    final expectations = targets.map((r) {
+      if (r != null && session.graph.knowsType(r)) {
+        return ExactTypeExpectation(session.typeOf(r)!);
+      }
+      return const NoTypeExpectation();
+    }).toList();
+
+    visit(e.rowValue, SelectTypeExpectation(expectations));
   }
 
   @override
diff --git a/sqlparser/lib/src/ast/statements/update.dart b/sqlparser/lib/src/ast/statements/update.dart
index bb480f4b..2d4e47c9 100644
--- a/sqlparser/lib/src/ast/statements/update.dart
+++ b/sqlparser/lib/src/ast/statements/update.dart
@@ -112,6 +112,11 @@ class MultiColumnSetComponent extends SetComponent {
   // Will be either Tuple or SubQuery
   Expression rowValue;
 
+  List<Column?>? get resolvedTargetColumns {
+    if (columns.isEmpty) return null;
+    return columns.map((c) => c.resolvedColumn).toList();
+  }
+
   MultiColumnSetComponent({required this.columns, required this.rowValue});
 
   @override