types2: Support inference for insert statements (#297)

sqlparser/lib/src/analysis/types2/expectation.dart

@@ -57,3 +57,14 @@ class RoughTypeExpectation extends TypeExpectation {
 enum _RoughType {
   numeric,
 }
+
+/// Type expectation for result columns in a select statement.
+///
+/// This can be used to set expectations in an `INSERT INTO SELECT` statement,
+/// where the column types are constrained by what's expected in the insert
+/// statement.
+class SelectTypeExpectation extends TypeExpectation {
+  final List<TypeExpectation> columnExpectations;
+
+  SelectTypeExpectation(this.columnExpectations);
+}

sqlparser/lib/src/analysis/types2/resolving_visitor.dart

@@ -10,6 +10,67 @@ class TypeResolver extends RecursiveVisitor<TypeExpectation, void> {
     session.finish();
   }
 
+  @override
+  void visitSelectStatement(SelectStatement e, TypeExpectation arg) {
+    _handleWhereClause(e);
+
+    var currentColumnIndex = 0;
+    final columnExpectations = arg is SelectTypeExpectation
+        ? arg.columnExpectations
+        : const <TypeExpectation>[];
+
+    for (final child in e.childNodes) {
+      if (child == e.where) continue; // handled above
+
+      if (child is ResultColumn) {
+        if (child is ExpressionResultColumn) {
+          final expectation = currentColumnIndex < columnExpectations.length
+              ? columnExpectations[currentColumnIndex]
+              : const NoTypeExpectation();
+          visit(child, expectation);
+
+          currentColumnIndex++;
+        } else if (child is StarResultColumn) {
+          currentColumnIndex += child.scope.availableColumns.length;
+        }
+      } else {
+        visit(child, arg);
+      }
+    }
+  }
+
+  @override
+  void visitInsertStatement(InsertStatement e, TypeExpectation arg) {
+    if (e.withClause != null) visit(e.withClause, arg);
+    visitList(e.targetColumns, const NoTypeExpectation());
+
+    final targets = e.resolvedTargetColumns ?? const [];
+    targets.forEach(_handleColumn);
+
+    final expectations = targets.map((r) {
+      if (session.graph.knowsType(r)) {
+        return ExactTypeExpectation(session.typeOf(r));
+      }
+      return const NoTypeExpectation();
+    }).toList();
+
+    e.source.when(
+      isSelect: (select) {
+        visit(select.stmt, SelectTypeExpectation(expectations));
+      },
+      isValues: (values) {
+        for (final tuple in values.values) {
+          for (var i = 0; i < tuple.expressions.length; i++) {
+            final expectation = i < expectations.length
+                ? expectations[i]
+                : const NoTypeExpectation();
+            visit(tuple.expressions[i], expectation);
+          }
+        }
+      },
+    );
+  }
+
   @override
   void visitCrudStatement(CrudStatement stmt, TypeExpectation arg) {
     if (stmt is HasWhereClause) {
@@ -172,6 +233,36 @@ class TypeResolver extends RecursiveVisitor<TypeExpectation, void> {
     visit(e.operand, const NoTypeExpectation());
   }
 
+  @override
+  void visitReference(Reference e, TypeExpectation arg) {
+    final resolved = e.resolvedColumn;
+    if (resolved == null) return;
+
+    _handleColumn(resolved);
+    _lazyCopy(e, resolved);
+  }
+
+  void _handleColumn(Column column) {
+    if (session.graph.knowsType(column)) return;
+
+    if (column is TableColumn) {
+      session.markTypeResolved(column, column.type);
+    } else if (column is ExpressionColumn) {
+      _lazyCopy(column, column.expression);
+    } else if (column is DelegatedColumn && column.innerColumn != null) {
+      _handleColumn(column.innerColumn);
+      _lazyCopy(column, column.innerColumn);
+    }
+  }
+
+  void _lazyCopy(Typeable to, Typeable from) {
+    if (session.graph.knowsType(from)) {
+      session.markTypeResolved(to, session.typeOf(from));
+    } else {
+      session.addRelationship(CopyTypeFrom(to, from));
+    }
+  }
+
   void _handleWhereClause(HasWhereClause stmt) {
     if (stmt.where != null) {
       // assume that a where statement is a boolean expression. Sqlite

sqlparser/lib/src/ast/visitor.dart

@@ -367,9 +367,12 @@ class RecursiveVisitor<A, R> implements AstVisitor<A, R> {
   R visit(AstNode e, A arg) => e.accept(this, arg);
 
   @protected
-  R visitChildren(AstNode e, A arg) {
-    for (final child in e.childNodes) {
-      child.accept(this, arg);
+  R visitChildren(AstNode e, A arg) => visitList(e.childNodes, arg);
+
+  @protected
+  R visitList(Iterable<AstNode> nodes, A arg) {
+    for (final node in nodes) {
+      node.accept(this, arg);
     }
     return null;
   }

sqlparser/test/analysis/types2/resolver_test.dart

@@ -72,4 +72,31 @@ void main() {
       expect(_resolveFirstVariable('SELECT CAST(? AS TEXT)'), null);
     });
   });
+
+  group('types in insert statements', () {
+    test('for VALUES', () {
+      final resolver =
+          _obtainResolver('INSERT INTO demo VALUES (:id, :content);');
+      final root = resolver.session.context.root;
+      final variables = root.allDescendants.whereType<Variable>();
+
+      final idVar = variables.singleWhere((v) => v.resolvedIndex == 1);
+      final contentVar = variables.singleWhere((v) => v.resolvedIndex == 2);
+
+      expect(resolver.session.typeOf(idVar), id.type);
+      expect(resolver.session.typeOf(contentVar), content.type);
+    });
+
+    test('for SELECT', () {
+      final resolver = _obtainResolver('INSERT INTO demo SELECT :id, :content');
+      final root = resolver.session.context.root;
+      final variables = root.allDescendants.whereType<Variable>();
+
+      final idVar = variables.singleWhere((v) => v.resolvedIndex == 1);
+      final contentVar = variables.singleWhere((v) => v.resolvedIndex == 2);
+
+      expect(resolver.session.typeOf(idVar), id.type);
+      expect(resolver.session.typeOf(contentVar), content.type);
+    });
+  });
 }