Handle common table expressions in the analyzer

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

@@ -106,8 +106,16 @@ class ReferenceExpressionColumn extends ExpressionColumn {
       : super(name: null, expression: ref);
 }
 
+/// A column that wraps another column.
+mixin DelegatedColumn on Column {
+  Column get innerColumn;
+
+  @override
+  String get name => innerColumn.name;
+}
+
 /// The result column of a [CompoundSelectStatement].
-class CompoundSelectColumn extends Column {
+class CompoundSelectColumn extends Column with DelegatedColumn {
   /// The column in [CompoundSelectStatement.base] each of the
   /// [CompoundSelectStatement.additional] that contributed to this column.
   final List<Column> columns;
@@ -115,5 +123,14 @@ class CompoundSelectColumn extends Column {
   CompoundSelectColumn(this.columns);
 
   @override
-  String get name => columns.first.name;
+  Column get innerColumn => columns.first;
+}
+
+class CommonTableExpressionColumn extends Column with DelegatedColumn {
+  @override
+  final String name;
+  @override
+  final Column innerColumn;
+
+  CommonTableExpressionColumn(this.name, this.innerColumn);
 }

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

@@ -44,6 +44,10 @@ class ReferenceScope {
     return ReferenceScope(this, root: effectiveRoot);
   }
 
+  ReferenceScope createSibling() {
+    return parent.createChild();
+  }
+
   /// Registers something that can be referenced in this and child scopes.
   void register(String identifier, Referencable ref) {
     _references.putIfAbsent(identifier.toUpperCase(), () => []).add(ref);

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

@@ -11,8 +11,9 @@ class ColumnResolver extends RecursiveVisitor<void> {
 
   @override
   void visitSelectStatement(SelectStatement e) {
-    _resolveSelect(e);
+    // visit children first so that common table expressions are resolved
     visitChildren(e);
+    _resolveSelect(e);
   }
 
   @override
@@ -168,9 +169,9 @@ class ColumnResolver extends RecursiveVisitor<void> {
     return span;
   }
 
-  Table _resolveTableReference(TableReference r) {
+  ResultSet _resolveTableReference(TableReference r) {
     final scope = r.scope;
-    final resolvedTable = scope.resolve<Table>(r.tableName, orElse: () {
+    final resolvedTable = scope.resolve<ResultSet>(r.tableName, orElse: () {
       final available = scope.allOf<Table>().map((t) => t.name);
 
       context.reportError(UnresolvedReferenceError(

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

@@ -38,7 +38,9 @@ class AstPreparingVisitor extends RecursiveVisitor<void> {
     final scope = e.scope;
 
     if (isInFROM) {
-      final forked = scope.effectiveRoot.createChild();
+      final surroundingSelect =
+          e.parents.firstWhere((node) => node is BaseSelectStatement).scope;
+      final forked = surroundingSelect.createSibling();
       e.scope = forked;
     } else {
       final forked = scope.createChild();
@@ -125,6 +127,12 @@ class AstPreparingVisitor extends RecursiveVisitor<void> {
     visitChildren(e);
   }
 
+  @override
+  void visitCommonTableExpression(CommonTableExpression e) {
+    e.scope.register(e.cteTableName, e);
+    visitChildren(e);
+  }
+
   @override
   void visitNumberedVariable(NumberedVariable e) {
     _foundVariables.add(e);
@@ -170,4 +178,18 @@ class AstPreparingVisitor extends RecursiveVisitor<void> {
       }
     }
   }
+
+  void _forkScope(AstNode node) {
+    node.scope = node.scope.createChild();
+  }
+
+  @override
+  void visitChildren(AstNode e) {
+    // hack to fork scopes on statements (selects are handled above)
+    if (e is Statement && e is! SelectStatement) {
+      _forkScope(e);
+    }
+
+    super.visitChildren(e);
+  }
 }

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

@@ -82,9 +82,8 @@ 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);
+      } else if (column is DelegatedColumn) {
+        return resolveColumn(column.innerColumn);
       }
 
       throw StateError('Unknown column $column');

sqlparser/lib/src/ast/ast.dart

@@ -138,6 +138,14 @@ abstract class AstNode with HasMetaMixin {
   /// type. The "content" refers to anything stored only in this node, children
   /// are ignored.
   bool contentEquals(covariant AstNode other);
+
+  @override
+  String toString() {
+    if (hasSpan) {
+      return '$runtimeType: ${span.text}';
+    }
+    return super.toString();
+  }
 }
 
 abstract class AstVisitor<T> {

sqlparser/lib/src/ast/clauses/with.dart

@@ -20,7 +20,7 @@ class WithClause extends AstNode {
   bool contentEquals(WithClause other) => other.recursive == recursive;
 }
 
-class CommonTableExpression extends AstNode {
+class CommonTableExpression extends AstNode with ResultSet {
   final String cteTableName;
 
   /// If this common table expression has explicit column names, e.g. with
@@ -47,4 +47,22 @@ class CommonTableExpression extends AstNode {
   bool contentEquals(CommonTableExpression other) {
     return other.cteTableName == cteTableName;
   }
+
+  @override
+  List<Column> get resolvedColumns {
+    final columnsOfSelect = as.resolvedColumns;
+    if (columnsOfSelect == null || columnNames == null) return columnsOfSelect;
+
+    // adapt names of result columns to the [columnNames] declared here
+    final mappedColumns = <Column>[];
+    for (var i = 0; i < columnNames.length; i++) {
+      final name = columnNames[i];
+
+      if (i < columnsOfSelect.length) {
+        final selectColumn = columnsOfSelect[i];
+        mappedColumns.add(CommonTableExpressionColumn(name, selectColumn));
+      }
+    }
+    return mappedColumns;
+  }
 }

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

@@ -10,7 +10,7 @@ mixin CrudParser on ParserBase {
       return _deleteStmt(withClause);
     } else if (_check(TokenType.update)) {
       return _update(withClause);
-    } else if (_check(TokenType.insert)) {
+    } else if (_check(TokenType.insert) || _check(TokenType.replace)) {
       return _insertStmt(withClause);
     }
     return null;

sqlparser/test/analysis/common_table_expressions.dart

@@ -0,0 +1,26 @@
+import 'package:sqlparser/sqlparser.dart';
+import 'package:test/test.dart';
+
+import 'data.dart';
+
+void main() {
+  test('resolves columns from CTEs', () {
+    final engine = SqlEngine()..registerTable(demoTable);
+
+    final context = engine.analyze('''
+      WITH
+        cte (foo, bar) AS (SELECT * FROM demo)
+      SELECT * FROM cte;
+    ''');
+
+    expect(context.errors, isEmpty);
+    final select = context.root as SelectStatement;
+    final types = context.types;
+
+    expect(select.resolvedColumns.map((c) => c.name), ['foo', 'bar']);
+    expect(
+      select.resolvedColumns.map((c) => types.resolveColumn(c).type),
+      [id.type, content.type],
+    );
+  });
+}