Handle recursive CTEs in analyzer

sqlparser/CHANGELOG.md

@@ -1,4 +1,5 @@
 ## unreleased
+- Support common table expressions
 - Handle special `rowid`, `oid`, `__rowid__` references
 - Support references to `sqlite_master` and `sqlite_sequence` tables
 

sqlparser/lib/src/analysis/error.dart

@@ -60,5 +60,6 @@ enum AnalysisErrorType {
 
   unknownFunction,
   compoundColumnCountMismatch,
+  cteColumnCountMismatch,
   other,
 }

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

@@ -129,8 +129,12 @@ class CompoundSelectColumn extends Column with DelegatedColumn {
 class CommonTableExpressionColumn extends Column with DelegatedColumn {
   @override
   final String name;
+
   @override
-  final Column innerColumn;
+  Column innerColumn;
+
+  // note that innerColumn is mutable because the column might not be known
+  // during all analysis phases.
 
   CommonTableExpressionColumn(this.name, this.innerColumn);
 }

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

@@ -54,6 +54,22 @@ class ColumnResolver extends RecursiveVisitor<void> {
     e.resolvedColumns = resolved;
   }
 
+  @override
+  void visitCommonTableExpression(CommonTableExpression e) {
+    visitChildren(e);
+
+    final resolved = e.as.resolvedColumns;
+    final names = e.columnNames;
+    if (names != null && resolved != null && names.length != resolved.length) {
+      context.reportError(AnalysisError(
+        type: AnalysisErrorType.cteColumnCountMismatch,
+        message: 'This CTE declares ${names.length} columns, but its select '
+            'statement actually returns ${resolved.length}.',
+        relevantNode: e,
+      ));
+    }
+  }
+
   @override
   void visitUpdateStatement(UpdateStatement e) {
     final table = _resolveTableReference(e.table);
@@ -129,26 +145,28 @@ class ColumnResolver extends RecursiveVisitor<void> {
       } else if (resultColumn is ExpressionResultColumn) {
         final expression = resultColumn.expression;
         Column column;
-        String name;
 
         if (expression is Reference) {
           column = ReferenceExpressionColumn(expression,
               overriddenName: resultColumn.as);
-          if (resultColumn.as != null) name = resultColumn.as;
         } else {
-          name = _nameOfResultColumn(resultColumn);
+          final name = _nameOfResultColumn(resultColumn);
           column =
               ExpressionColumn(name: name, expression: resultColumn.expression);
         }
 
         usedColumns.add(column);
 
-        // make this column available if there is no other with the same name
-        if (name != null && !availableColumns.any((c) => c.name == name)) {
+        if (resultColumn.as != null) {
+          // make this column available for references if there is no other
+          // column with the same name
+          final name = resultColumn.as;
+          if (!availableColumns.any((c) => c.name == name)) {
             availableColumns.add(column);
           }
         }
       }
+    }
 
     s.resolvedColumns = usedColumns;
     scope.availableColumns = availableColumns;

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 with ResultSet {
+class CommonTableExpression extends AstNode with ResultSet, VisibleToChildren {
   final String cteTableName;
 
   /// If this common table expression has explicit column names, e.g. with
@@ -32,6 +32,8 @@ class CommonTableExpression extends AstNode with ResultSet {
   Token asToken;
   IdentifierToken tableNameToken;
 
+  List<CommonTableExpressionColumn> _cachedColumns;
+
   CommonTableExpression(
       {@required this.cteTableName, this.columnNames, @required this.as});
 
@@ -51,18 +53,27 @@ class CommonTableExpression extends AstNode with ResultSet {
   @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];
+    // we don't override column names, so just return the columns declared by
+    // the select statement
+    if (columnNames == null) return columnsOfSelect;
 
+    _cachedColumns ??= columnNames
+        .map((name) => CommonTableExpressionColumn(name, null))
+        .toList();
+
+    if (columnsOfSelect != null) {
+      // bind the CommonTableExpressionColumn to the real underlying column
+      // returned by the select statement
+
+      for (var i = 0; i < _cachedColumns.length; i++) {
         if (i < columnsOfSelect.length) {
           final selectColumn = columnsOfSelect[i];
-        mappedColumns.add(CommonTableExpressionColumn(name, selectColumn));
+          _cachedColumns[i].innerColumn = selectColumn;
         }
       }
-    return mappedColumns;
+    }
+
+    return _cachedColumns;
   }
 }

sqlparser/test/analysis/common_table_expressions.dart

@@ -23,4 +23,40 @@ void main() {
       [id.type, content.type],
     );
   });
+
+  test('warns on column count mismatch', () {
+    final engine = SqlEngine()..registerTable(demoTable);
+
+    final context = engine.analyze(''' 
+      WITH
+        cte (foo, bar, baz) AS (SELECT * FROM demo)
+      SELECT 1;
+    ''');
+
+    expect(context.errors, hasLength(1));
+    final error = context.errors.single;
+    expect(error.type, AnalysisErrorType.cteColumnCountMismatch);
+    expect(error.message, stringContainsInOrder(['3', '2']));
+  });
+
+  test('handles recursive CTEs', () {
+    final engine = SqlEngine();
+
+    final context = engine.analyze('''
+      WITH RECURSIVE
+        cnt(x) AS (
+          SELECT 1
+          UNION ALL
+          SELECT x+1 FROM cnt
+            LIMIT 1000000
+        )
+      SELECT x FROM cnt;
+    ''');
+
+    expect(context.errors, isEmpty);
+    final select = context.root as SelectStatement;
+    final column = context.typeOf(select.resolvedColumns.single);
+
+    expect(column.type, const ResolvedType(type: BasicType.int));
+  });
 }