drift3: Support LIST() queries

Author: simolus3

drift_dev/lib/src/analysis/resolver/queries/nested_queries.dart

@@ -126,15 +126,15 @@ class _AnalyzerState {
   void _process() {
     // Add necessary columns to select variables read by inner nested queries.
     for (final variable in container.variablesCapturedByChildren) {
-      container.addedColumns.add(
-        ExpressionResultColumn(
-          expression: Reference(
-            entityName: variable.reference.entityName,
-            columnName: variable.reference.columnName,
-          ),
-          as: AliasClause(variable.helperColumn),
+      final column = ExpressionResultColumn(
+        expression: Reference(
+          entityName: variable.reference.entityName,
+          columnName: variable.reference.columnName,
         ),
+        as: AliasClause(variable.helperColumn),
       );
+      container.addedColumns.add(column);
+      variable.columnIndex = container.nextColumnIndex++;
     }
 
     // Re-index variables. This is necessary for two reasons:

drift_dev/lib/src/analysis/results/query.dart

@@ -332,6 +332,7 @@ class SqlSelectQuery extends SqlQuery {
 class NestedQueriesContainer {
   final SelectStatement select;
   final Map<NestedQueryColumn, NestedQuery> nestedQueries = {};
+  int nextColumnIndex = 0;
 
   /// The nested queries transformation may change the index of variables used
   /// in a query.
@@ -347,7 +348,22 @@ class NestedQueriesContainer {
   /// variables with different indexes are considered different variables.
   final Map<Variable, int> originalIndexForVariable = {};
 
-  NestedQueriesContainer(this.select);
+  NestedQueriesContainer(this.select) {
+    for (final column in select.columns) {
+      if (column case NestedStarResultColumn(:final resultSet?)) {
+        // Nested star columns aren't resolved by the sqlparser package, so we
+        // need to count columns here.
+        nextColumnIndex += resultSet.resolvedColumns
+                ?.where((c) => c.includedInResults)
+                .length ??
+            0;
+      } else if (column.resolvedColumns case final columns?) {
+        // All other result columns should have columns set on them if they
+        // appear in the final query.
+        nextColumnIndex += columns.length;
+      }
+    }
+  }
 
   /// Columns that should be added to the [select] statement to read variables
   /// captured by children.
@@ -390,9 +406,14 @@ class CapturedVariable {
   /// This variable is not mounted to the same syntax tree as [reference], it
   /// will be mounted into the tree returned by [addHelperNodes].
   final NamedVariable introducedVariable;
+  FoundVariable? resolvedVariable;
 
   String get helperColumn => '\$n_$queryGlobalId';
 
+  /// The index at which [helperColumn] has been inserted into the outer select
+  /// statement.
+  int? columnIndex;
+
   CapturedVariable(this.reference, this.queryGlobalId)
       : introducedVariable = NamedVariable.synthetic(':', 'r$queryGlobalId') {
     introducedVariable.setMeta<CapturedVariable>(this);
@@ -1039,14 +1060,16 @@ class FoundVariable extends FoundElement implements HasType {
     required this.name,
     required this.sqlType,
     required Variable variable,
-    required this.forCaptured,
+    required CapturedVariable this.forCaptured,
   })  : originalIndex = index,
         typeConverter = null,
         nullable = false,
         isArray = false,
         isRequired = true,
         hidden = true,
-        syntacticOrigin = variable;
+        syntacticOrigin = variable {
+    forCaptured!.resolvedVariable = this;
+  }
 
   @override
   String get dartParameterName {

drift_dev/lib/src/writer/queries/query_writer.dart

@@ -28,6 +28,7 @@ typedef _ArgumentContext = ({
 /// should be included in a generated database or dao class.
 class QueryWriter {
   final Scope scope;
+  final Map<CapturedVariable, String> _outerVariables = {};
 
   late final ExplicitAliasTransformer _transformer;
   final TextEmitter _emitter;
@@ -40,6 +41,8 @@ class QueryWriter {
 
   QueryWriter(this.scope) : _emitter = scope.leaf();
 
+  QueryWriter._existingEmitter(this.scope, this._emitter);
+
   void write(SqlQuery query) {
     // Note that writing queries can have a result set if they use a RETURNING
     // clause.
@@ -88,9 +91,10 @@ class QueryWriter {
 
   /// Writes the function literal that turns a "QueryRow" into the desired
   /// custom return type of a query.
-  void _writeMappingLambda(InferredResultSet resultSet, QueryRowType rowClass) {
+  void _writeMappingLambda(InferredResultSet resultSet,
+      NestedQueriesContainer? nestedQueries, QueryRowType rowClass) {
     if (drift3) {
-      return _Drift3MappingCodeWriter(this)
+      return _Drift3MappingCodeWriter(this, nestedQueries)
           .write(_emitter, resultSet, rowClass);
     }
 
@@ -317,17 +321,20 @@ class QueryWriter {
       [QueryRowType? resultType]) {
     final resultSet = select.resultSet;
     resultType ??= select.queryRowType(options);
+    final needsAsyncMapping = resultType.requiresAsynchronousContext(options);
 
     if (drift3) {
+      final methodName =
+          needsAsyncMapping ? 'customSelectMappedAsync' : 'customSelectMapped';
       _emitter
-        ..write(' customSelectMapped<')
+        ..write(' $methodName<')
         ..writeDart(resultType.rowType)
         ..write('>(query: ${_queryCode(select)},');
       _writeVariables(select);
       _emitter.write(', ');
       _writeReadsFrom(select);
       _emitter.write(', createMapper: ');
-      _writeMappingLambda(resultSet, resultType);
+      _writeMappingLambda(resultSet, select.nestedContainer, resultType);
       _emitter.write(')');
       return;
     }
@@ -337,13 +344,13 @@ class QueryWriter {
     _buffer.write(', ');
     _writeReadsFrom(select);
 
-    if (resultType.requiresAsynchronousContext(options)) {
+    if (needsAsyncMapping) {
       _buffer.write(').asyncMap(');
     } else {
       _buffer.write(').map(');
     }
 
-    _writeMappingLambda(resultSet, resultType);
+    _writeMappingLambda(resultSet, select.nestedContainer, resultType);
     _buffer.write(')');
   }
 
@@ -374,7 +381,7 @@ class QueryWriter {
     if (drift3) {
       _buffer.writeln(').then((rows) {');
 
-      final mappingWriter = _Drift3MappingCodeWriter(this)
+      final mappingWriter = _Drift3MappingCodeWriter(this, null)
         .._writeArgumentExpression(
             rowType, resultSet, (isNullable: false, sqlPrefix: null));
       _buffer
@@ -398,11 +405,11 @@ class QueryWriter {
 
       if (rowType.requiresAsynchronousContext(options)) {
         _buffer.write('Future.wait(rows.map(');
-        _writeMappingLambda(resultSet, rowType);
+        _writeMappingLambda(resultSet, null, rowType);
         _buffer.write('))');
       } else {
         _buffer.write('rows.map(');
-        _writeMappingLambda(resultSet, rowType);
+        _writeMappingLambda(resultSet, null, rowType);
         _buffer.write(').toList()');
       }
       _buffer.write(');\n}');
@@ -545,7 +552,7 @@ class QueryWriter {
   }
 
   void _writeVariables(SqlQuery query) {
-    _ExpandedVariableWriter(query, _emitter).writeVariables();
+    _ExpandedVariableWriter(query, this).writeVariables();
   }
 
   /// Returns a Dart string literal representing the query after variables have
@@ -640,6 +647,7 @@ String readConverter(TextEmitter emitter, AppliedTypeConverter converter) {
 /// query.
 class _Drift3MappingCodeWriter {
   final QueryWriter _writer;
+  final NestedQueriesContainer? nestedQueries;
 
   final StringBuffer _outerSetup = StringBuffer();
   final StringBuffer _innerMapper = StringBuffer();
@@ -649,7 +657,7 @@ class _Drift3MappingCodeWriter {
 
   TextEmitter get _emitter => _writer._emitter;
 
-  _Drift3MappingCodeWriter(this._writer);
+  _Drift3MappingCodeWriter(this._writer, this.nestedQueries);
 
   String referenceBuiltinType(DriftSqlType type) {
     return _obtainedTypes.putIfAbsent(type, () {
@@ -690,7 +698,10 @@ class _Drift3MappingCodeWriter {
       ..writeln(_outerSetup)
       ..write('return (')
       ..writeDriftRef('RawRow')
-      ..writeln(' row) => $_innerMapper;')
+      ..write(' row) ')
+      ..write(
+          rowClass.requiresAsynchronousContext(_writer.options) ? 'async' : '')
+      ..writeln('=> $_innerMapper;')
       ..writeln('}');
   }
 
@@ -715,12 +726,23 @@ class _Drift3MappingCodeWriter {
           resultSet,
           (sqlPrefix: prefix, isNullable: argument.nullable),
         );
-      case MappedNestedListQuery():
-        _innerMapper.write("throw 'todo'");
-//        _innerMapper.write('await ');
-//        final query = argument.column.query;
-//        _writeCustomSelectStatement(query, argument.nestedType);
-//        _innerMapper.write('.get()');
+      case MappedNestedListQuery(:final column):
+        final knownVariableTypes = <CapturedVariable, String>{};
+        if (nestedQueries?.nestedQueries[column.from] case final nested?) {
+          for (final variable in nested.capturedVariables.values) {
+            if (variable.resolvedVariable case final resolved?) {
+              knownVariableTypes[variable] = referenceType(resolved.sqlType);
+            }
+          }
+        }
+
+        _innerMapper.write('await ');
+        final query = argument.column.query;
+        final innerWriter = QueryWriter._existingEmitter(
+            _writer.scope, TextEmitter(_writer.scope, buffer: _innerMapper));
+        innerWriter._outerVariables.addAll(knownVariableTypes);
+        innerWriter._writeCustomSelectStatement(query, argument.nestedType);
+        _innerMapper.write('.get()');
       case QueryRowType():
         final singleValue = argument.singleValue;
         if (singleValue != null) {
@@ -1007,10 +1029,12 @@ class _ExpandedDeclarationWriter {
 class _ExpandedVariableWriter {
   final SqlQuery query;
   final TextEmitter _emitter;
+  final QueryWriter _queryWriter;
 
   StringBuffer get _buffer => _emitter.buffer;
 
-  _ExpandedVariableWriter(this.query, this._emitter);
+  _ExpandedVariableWriter(this.query, this._queryWriter)
+      : _emitter = _queryWriter._emitter;
 
   void writeVariables() {
     _buffer.write('variables: ');
@@ -1097,6 +1121,16 @@ class _ExpandedVariableWriter {
     String constructVar(String dartExpr) {
       final capture = element.forCaptured;
       if (capture != null) {
+        if (_emitter.writer.options.drift3Preview) {
+          return _emitter.dartCode(AnnotatedDartCode.build((b) {
+            b.addSymbol('MappedValue', AnnotatedDartCode.drift);
+            b.addText('.raw(');
+            b.addText('${_queryWriter._outerVariables[capture]},');
+            b.addText('row[${capture.columnIndex}]');
+            b.addText(')');
+          }));
+        }
+
         dartExpr = ('row.read(${asDartLiteral(capture.helperColumn)})');
       }
 

drift_dev/lib/src/writer/writer.dart

@@ -277,24 +277,20 @@ abstract class _NodeOrWriter {
     return AnnotatedDartCode.build((b) {
       final converter = column.typeConverter;
 
-      if (writer.options.drift3Preview) {
-        b.addText('mapValue(');
-        b.addCode(_drift3SqlType(column.sqlType));
-        b.addText(', ');
-        if (converter != null) {
-          // apply type converter before writing the variable
-          b
-            ..addCode(readConverter(converter, forNullable: column.nullable))
-            ..addText('.toSql(')
-            ..addCode(expression)
-            ..addText(')');
-        } else {
-          b.addCode(expression);
-        }
-        b.addText(')');
-
-        return;
+      b.addText('mapValue(');
+      b.addCode(_drift3SqlType(column.sqlType));
+      b.addText(', ');
+      if (converter != null) {
+        // apply type converter before writing the variable
+        b
+          ..addCode(readConverter(converter, forNullable: column.nullable))
+          ..addText('.toSql(')
+          ..addCode(expression)
+          ..addText(')');
+      } else {
+        b.addCode(expression);
       }
+      b.addText(')');
     });
   }
 
@@ -568,14 +564,16 @@ class Scope extends _Node {
 }
 
 class TextEmitter extends _Node {
-  final StringBuffer buffer = StringBuffer();
+  final StringBuffer buffer;
   final void Function(TaggedDartLexeme, StringBuffer)? writeTaggedDartCode;
 
   @override
   final Writer writer;
 
-  TextEmitter(Scope super.parent, {this.writeTaggedDartCode})
-      : writer = parent.writer;
+  TextEmitter(Scope super.parent,
+      {this.writeTaggedDartCode, StringBuffer? buffer})
+      : writer = parent.writer,
+        buffer = buffer ?? StringBuffer();
 
   @override
   void _writeTagged(StringBuffer buffer, TaggedDartLexeme lexeme) {

drift_dev/test/writer/queries/query_writer_test.dart

@@ -268,6 +268,17 @@ FROM tbl AS parent WHERE parent.a = :a;
       );
     });
 
+    test('drift3 preview', () {
+      return runTest(
+        const DriftOptions.defaults(drift3Preview: true),
+        [
+          contains(
+            r"variables: [mapValue(BuiltinDriftType.text, a), MappedValue.raw(type$0,row[1]), mapValue(BuiltinDriftType.text, b)]",
+          ),
+        ],
+      );
+    });
+
     test('should generate correct data class', () {
       return runTest(
         const DriftOptions.defaults(),

future/drift3/lib/src/database/connection_user.dart

@@ -539,6 +539,25 @@ abstract base class DatabaseConnectionUser {
     );
   }
 
+  /// Like [customSelectMapped], but allows [createMapper] to be asynchronous
+  /// and to fire off additional queries.
+  Selectable<T> customSelectMappedAsync<T>({
+    required String query,
+    required Future<T> Function(RawRow) Function(RawResultSet) createMapper,
+    List<MappedValue> variables = const [],
+    Set<ResultSet> readsFrom = const {},
+    bool isReadOnly = true,
+  }) {
+    return AsyncCustomSelectStatement(
+      query,
+      variables,
+      readsFrom,
+      createMapper,
+      this,
+      isReadOnly,
+    );
+  }
+
   /// Map a Dart [value] into a typed [MappedValue] understood by the database
   /// driver.
   MappedValue mapValue<T extends Object>(SqlType<T> type, T? value) {