Resolve types of columns in compound select statements

sqlparser/lib/src/analysis/error.dart

@@ -59,5 +59,6 @@ enum AnalysisErrorType {
   synctactic,
 
   unknownFunction,
+  compoundColumnCountMismatch,
   other,
 }

sqlparser/lib/src/analysis/schema/column.dart

@@ -13,7 +13,8 @@ class TableColumn extends Column {
   @override
   final String name;
 
-  /// The type of this column, which is immediately available.
+  /// The type of this column, which is available before any resolution happens
+  /// (we know if from the table).
   final ResolvedType type;
 
   /// The column constraints set on this column.
@@ -39,3 +40,15 @@ class ExpressionColumn extends Column {
 
   ExpressionColumn({@required this.name, this.expression});
 }
+
+/// The result column of a [CompoundSelectStatement].
+class CompoundSelectColumn extends Column {
+  /// The column in [CompoundSelectStatement.base] each of the
+  /// [CompoundSelectStatement.additional] that contributed to this column.
+  final List<Column> columns;
+
+  CompoundSelectColumn(this.columns);
+
+  @override
+  String get name => columns.first.name;
+}

sqlparser/lib/src/analysis/steps/column_resolver.dart

@@ -15,6 +15,44 @@ class ColumnResolver extends RecursiveVisitor<void> {
     visitChildren(e);
   }
 
+  @override
+  void visitCompoundSelectStatement(CompoundSelectStatement e) {
+    // first, visit all children so that the compound parts have their columns
+    // resolved
+    visitChildren(e);
+
+    final columnSets = [
+      e.base.resolvedColumns,
+      for (var part in e.additional) part.select.resolvedColumns
+    ];
+
+    // each select statement must return the same amount of columns
+    final amount = columnSets.first.length;
+    for (var i = 1; i < columnSets.length; i++) {
+      if (columnSets[i].length != amount) {
+        context.reportError(AnalysisError(
+          type: AnalysisErrorType.compoundColumnCountMismatch,
+          relevantNode: e,
+          message: 'The parts of this compound statement return different '
+              'amount of columns',
+        ));
+        break;
+      }
+    }
+
+    final resolved = <CompoundSelectColumn>[];
+
+    // merge all columns at each position into a CompoundSelectColumn
+    for (var i = 0; i < amount; i++) {
+      final columnsAtThisIndex = [
+        for (var set in columnSets) if (set.length > i) set[i]
+      ];
+
+      resolved.add(CompoundSelectColumn(columnsAtThisIndex));
+    }
+    e.resolvedColumns = resolved;
+  }
+
   @override
   void visitUpdateStatement(UpdateStatement e) {
     final table = _resolveTableReference(e.table);

sqlparser/lib/src/analysis/types/resolver.dart

@@ -63,6 +63,9 @@ class TypeResolver {
         return ResolveResult(column.type);
       } else if (column is ExpressionColumn) {
         return resolveOrInfer(column.expression);
+      } else if (column is CompoundSelectColumn) {
+        // todo maybe use a type that matches every column in here?
+        return resolveColumn(column.columns.first);
       }
 
       throw StateError('Unknown column $column');

sqlparser/lib/src/reader/parser/expressions.dart

@@ -279,13 +279,14 @@ mixin ExpressionParser on ParserBase {
           ..token = typedToken
           ..setSpan(_previous, _previous);
       }
-    } else if (_checkLenientIdentifier()) {
+    } else if (_checkIdentifier()) {
       final first = _consumeIdentifier(
           'This error message should never be displayed. Please report.');
 
       // could be table.column, function(...) or just column
       if (_matchOne(TokenType.dot)) {
-        final second = _consumeIdentifier('Expected a column name here');
+        final second =
+            _consumeIdentifier('Expected a column name here', lenient: true);
         return Reference(
             tableName: first.identifier, columnName: second.identifier)
           ..setSpan(first, second);

sqlparser/lib/src/reader/parser/parser.dart

@@ -103,11 +103,12 @@ abstract class ParserBase {
 
   /// Returns whether the next token is an [TokenType.identifier] or a
   /// [KeywordToken]. If this method returns true, calling [_consumeIdentifier]
-  /// with the lenient parameter will now throw.
-  bool _checkLenientIdentifier() {
+  /// with same [lenient] parameter will now throw.
+  bool _checkIdentifier({bool lenient = false}) {
     final next = _peek;
-    return next.type == TokenType.identifier ||
-        (next is KeywordToken && next.canConvertToIdentifier());
+    if (next.type == TokenType.identifier) return true;
+
+    return next is KeywordToken && (next.canConvertToIdentifier() || lenient);
   }
 
   Token _advance() {
@@ -130,10 +131,11 @@ abstract class ParserBase {
   }
 
   /// Consumes an identifier.
-  IdentifierToken _consumeIdentifier(String message) {
+  IdentifierToken _consumeIdentifier(String message, {bool lenient = false}) {
     final next = _peek;
-    // non-standard keywords can be parsed as an identifier
-    if (next is KeywordToken && next.canConvertToIdentifier()) {
+    // non-standard keywords can be parsed as an identifier, we allow all
+    // keywords when lenient is true
+    if (next is KeywordToken && (next.canConvertToIdentifier() || lenient)) {
       return (_advance() as KeywordToken).convertToIdentifier();
     }
     return _consume(TokenType.identifier, message) as IdentifierToken;