Begin with existing types for queries

drift/CHANGELOG.md

@@ -9,6 +9,8 @@
 - Add `likeExp` to generate `LIKE` expression with any comparison expression.
 - Fix `UNIQUE` keys declared in drift files being written twice.
 - Fix `customConstraints` not appearing in dumped database schema files.
+- Lazily load columns in `TypedResult.read`, increasing performance for joins
+  with lots of tables or columns.
 - Work-around an issue causing complex migrations via `Migrator.alterTable` not to
   work if a view referenced the altered table.
 

drift_dev/lib/src/analysis/resolver/drift/element_resolver.dart

@@ -1,3 +1,4 @@
+import 'package:analyzer/dart/element/nullability_suffix.dart';
 import 'package:collection/collection.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/type.dart';
@@ -48,18 +49,36 @@ abstract class DriftElementResolver<T extends DiscoveredElement>
     return null;
   }
 
-  Future<FoundDartClass?> findDartClass(String identifier) async {
-    final foundElement = await _findInDart(identifier);
-    if (foundElement is InterfaceElement) {
-      return FoundDartClass(foundElement, null);
-    } else if (foundElement is TypeAliasElement) {
-      final innerType = foundElement.aliasedType;
-      if (innerType is InterfaceType) {
-        return FoundDartClass(innerType.element, innerType.typeArguments);
-      }
-    }
+  /// Resolves [identifier] to a Dart element declaring a type, or reports an
+  /// error if this is not possible.
+  ///
+  /// The [syntacticSource] will be the base for the error's span.
+  Future<DartType?> findDartTypeOrReportError(
+      String identifier, SyntacticEntity syntacticSource) async {
+    final element = await _findInDart(identifier);
 
-    return null;
+    if (element == null) {
+      reportError(
+        DriftAnalysisError.inDriftFile(syntacticSource,
+            'Could not find `$identifier`, are you missing an import?'),
+      );
+      return null;
+    } else if (element is InterfaceElement) {
+      final library = element.library;
+      return library.typeSystem.instantiateInterfaceToBounds(
+          element: element, nullabilitySuffix: NullabilitySuffix.none);
+    } else if (element is TypeAliasElement) {
+      final library = element.library;
+      return library.typeSystem.instantiateTypeAliasToBounds(
+          element: element, nullabilitySuffix: NullabilitySuffix.none);
+    } else {
+      reportError(DriftAnalysisError.inDriftFile(
+        syntacticSource,
+        '`$identifier` does not refer to anything defining a type. Expected '
+        'a class, a mixin, an interface or a typedef.',
+      ));
+      return null;
+    }
   }
 
   /// Attempts to find a matching [ExistingRowClass] for a [DriftTableName]

drift_dev/lib/src/analysis/resolver/drift/find_dart_class.dart

@@ -1,36 +0,0 @@
-import 'package:analyzer/dart/element/element.dart';
-import 'package:analyzer/dart/element/type.dart';
-
-import '../../backend.dart';
-import '../resolver.dart';
-import '../shared/dart_types.dart';
-
-extension FindDartClass on LocalElementResolver {
-  /// Resolves a Dart class or generalized typedef pointing towards a Dart class.
-  Future<FoundDartClass?> findDartClass(
-      List<Uri> imports, String identifier) async {
-    final dartImports =
-        imports.where((importUri) => importUri.path.endsWith('.dart'));
-
-    for (final import in dartImports) {
-      LibraryElement library;
-      try {
-        library = await resolver.driver.backend.readDart(import);
-      } on NotALibraryException {
-        continue;
-      }
-
-      final foundElement = library.exportNamespace.get(identifier);
-      if (foundElement is InterfaceElement) {
-        return FoundDartClass(foundElement, null);
-      } else if (foundElement is TypeAliasElement) {
-        final innerType = foundElement.aliasedType;
-        if (innerType is InterfaceType) {
-          return FoundDartClass(innerType.element, innerType.typeArguments);
-        }
-      }
-    }
-
-    return null;
-  }
-}

drift_dev/lib/src/analysis/resolver/drift/query.dart

@@ -1,3 +1,4 @@
+import 'package:analyzer/dart/element/type.dart';
 import 'package:sqlparser/sqlparser.dart';
 
 import '../../driver/state.dart';
@@ -21,6 +22,20 @@ class DriftQueryResolver
 
     // Note: We don't analyze the query here, that happens in
     // `file_analysis.dart` after elements have been resolved.
+
+    String? resultClassName;
+    DartType? existingType;
+
+    final as = discovered.sqlNode.as;
+    if (as != null) {
+      if (as.useExistingDartClass) {
+        existingType =
+            await findDartTypeOrReportError(as.overriddenDataClassName, as);
+      } else {
+        resultClassName = as.overriddenDataClassName;
+      }
+    }
+
     return DefinedSqlQuery(
       discovered.ownId,
       DriftDeclaration.driftFile(stmt, file.ownUri),
@@ -28,7 +43,8 @@ class DriftQueryResolver
       sql: source.substring(stmt.firstPosition, stmt.lastPosition),
       sqlOffset: stmt.firstPosition,
       mode: isCreate ? QueryMode.atCreate : QueryMode.regular,
-      resultClassName: discovered.sqlNode.as,
+      resultClassName: resultClassName,
+      existingDartType: existingType,
     );
   }
 }

drift_dev/lib/src/analysis/resolver/drift/table.dart

@@ -61,19 +61,14 @@ class DriftTableResolver extends DriftElementResolver<DiscoveredDriftTable> {
           typeName != null ? _enumRegex.firstMatch(typeName) : null;
       if (enumIndexMatch != null) {
         final dartTypeName = enumIndexMatch.group(2)!;
-        final dartClass = await findDartClass(dartTypeName);
+        final dartType = await findDartTypeOrReportError(
+            dartTypeName, column.definition?.typeNames?.toSingleEntity ?? stmt);
 
-        if (dartClass == null) {
-          reportError(DriftAnalysisError.inDriftFile(
-            column.definition!.typeNames!.toSingleEntity,
-            'Type $dartTypeName could not be found. Are you missing '
-            'an import?',
-          ));
-        } else {
+        if (dartType != null) {
           converter = readEnumConverter(
             (msg) => reportError(
                 DriftAnalysisError.inDriftFile(column.definition ?? stmt, msg)),
-            dartClass.classElement.thisType,
+            dartType,
             type == DriftSqlType.int ? EnumType.intEnum : EnumType.textEnum,
             await resolver.driver.loadKnownTypes(),
           );

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

@@ -1,3 +1,4 @@
+import 'package:analyzer/dart/element/type.dart';
 import 'package:collection/collection.dart';
 import 'package:drift/drift.dart' show DriftSqlType, UpdateKind;
 import 'package:recase/recase.dart';
@@ -31,7 +32,18 @@ class DefinedSqlQuery extends DriftElement implements DriftQueryDeclaration {
   /// The unmodified source of the declared SQL statement forming this query.
   final String sql;
 
+  /// The overriden name of a result class that drift should generate for this
+  /// query.
+  ///
+  /// When multiple queries share the same result class name, drift will verify
+  /// that this is possible and map all these queries into the same generated
+  /// class.
   final String? resultClassName;
+
+  /// The existing Dart type into which a result row of this query should be
+  /// mapped.
+  final DartType? existingDartType;
+
   final QueryMode mode;
 
   /// The offset of [sql] in the source file, used to properly report errors
@@ -51,6 +63,7 @@ class DefinedSqlQuery extends DriftElement implements DriftQueryDeclaration {
     required this.sql,
     required this.sqlOffset,
     this.resultClassName,
+    this.existingDartType,
     this.mode = QueryMode.regular,
   });
 }

drift_dev/lib/src/analysis/serializer.dart

@@ -68,6 +68,7 @@ class ElementSerializer {
         'sql': element.sql,
         'offset': element.sqlOffset,
         'result_class': element.resultClassName,
+        'eixsting_type': _serializeType(element.existingDartType),
         'mode': element.mode.name,
       };
     } else if (element is DriftTrigger) {
@@ -496,6 +497,8 @@ class ElementDeserializer {
           createStmt: json['sql'] as String,
         );
       case 'query':
+        final rawExistingType = json['eixsting_type'];
+
         return DefinedSqlQuery(
           id,
           declaration,
@@ -503,6 +506,9 @@ class ElementDeserializer {
           sql: json['sql'] as String,
           sqlOffset: json['offset'] as int,
           resultClassName: json['result_class'] as String?,
+          existingDartType: rawExistingType != null
+              ? await _readDartType(id.libraryUri, rawExistingType as int)
+              : null,
           mode: QueryMode.values.byName(json['mode'] as String),
         );
       case 'trigger':

drift_dev/test/analysis/resolver/queries/existing_row_classes_test.dart

@@ -0,0 +1,49 @@
+import 'package:drift_dev/src/analysis/results/results.dart';
+import 'package:test/test.dart';
+
+import '../../test_utils.dart';
+
+void main() {
+  test('recognizes existing row classes', () async {
+    final state = TestBackend.inTest({
+      'a|lib/a.drift': '''
+import 'a.dart';
+
+foo WITH MyRow: SELECT 'hello world', 2;
+''',
+      'a|lib/a.dart': '''
+class MyRow {
+  final String a;
+  final int b;
+
+  MyRow(this.a, this.b);
+}
+''',
+    });
+
+    final uri = Uri.parse('package:a/a.drift');
+    final file = await state.driver.resolveElements(uri);
+
+    state.expectNoErrors();
+    final query = file.analyzedElements.single as DefinedSqlQuery;
+    expect(query.resultClassName, isNull);
+    expect(query.existingDartType?.getDisplayString(withNullability: true),
+        'MyRow');
+  });
+
+  test("warns if existing row classes don't exist", () async {
+    final state = TestBackend.inTest({
+      'a|lib/a.drift': '''
+import 'a.dart';
+
+foo WITH MyRow: SELECT 'hello world', 2;
+''',
+    });
+
+    final file = await state.analyze('package:a/a.drift');
+    expect(file.allErrors, [
+      isDriftError(contains('are you missing an import?'))
+          .withSpan('WITH MyRow')
+    ]);
+  });
+}

sqlparser/lib/src/ast/drift/declared_statement.dart

@@ -15,7 +15,7 @@ class DeclaredStatement extends Statement implements PartOfDriftFile {
   List<StatementParameter> parameters;
 
   /// The desired result class name, if set.
-  final String? as;
+  DriftTableName? as;
 
   Token? colon;
 
@@ -36,10 +36,12 @@ class DeclaredStatement extends Statement implements PartOfDriftFile {
   void transformChildren<A>(Transformer<A> transformer, A arg) {
     statement = transformer.transformChild(statement, this, arg);
     parameters = transformer.transformChildren(parameters, this, arg);
+    as = transformer.transformNullableChild(as, this, arg);
   }
 
   @override
-  Iterable<AstNode> get childNodes => [statement, ...parameters];
+  Iterable<AstNode> get childNodes =>
+      [statement, ...parameters, if (as != null) as!];
 }
 
 /// How a statement was declared in a drift file.

sqlparser/lib/src/reader/parser.dart

@@ -324,10 +324,7 @@ class Parser {
       _consume(TokenType.rightParen, 'Expected closing parenthesis');
     }
 
-    String? as;
-    if (_matchOne(TokenType.as)) {
-      as = _consumeIdentifier('Expected a name of the result class').identifier;
-    }
+    final as = _driftTableName();
 
     final colon = _consume(
         TokenType.colon,

sqlparser/lib/src/utils/ast_equality.dart

@@ -431,7 +431,7 @@ class EqualityEnforcingVisitor implements AstVisitor<void, void> {
 
   void visitDriftDeclaredStatement(DeclaredStatement e, void arg) {
     final current = _currentAs<DeclaredStatement>(e);
-    _assert(current.identifier == e.identifier && current.as == e.as, e);
+    _assert(current.identifier == e.identifier, e);
     _checkChildren(e);
   }
 

sqlparser/lib/utils/node_to_text.dart

@@ -448,11 +448,7 @@ class NodeSqlBuilder extends AstVisitor<void, void> {
         symbol(')');
       }
 
-      if (e.as != null) {
-        _keyword(TokenType.as);
-        identifier(e.as!);
-      }
-
+      visitNullable(e.as, arg);
       symbol(':', spaceAfter: true);
       visit(e.statement, arg);
       symbol(';');

sqlparser/test/parser/drift_file_test.dart

@@ -111,7 +111,7 @@ void main() {
             ],
             from: TableReference('tbl', as: 'foo'),
           ),
-          as: 'MyResultSet',
+          as: DriftTableName('MyResultSet', false),
         ),
         DeclaredStatement(
           SimpleName('add'),
@@ -236,4 +236,23 @@ CREATE INDEX x ON foo (a);
 
     expect(result.errors, isEmpty);
   });
+
+  test('declared statements can use existing classes syntax', () {
+    testDriftFile(
+      'foo WITH ExistingDartClass: SELECT 1;',
+      DriftFile([
+        DeclaredStatement(
+          SimpleName('foo'),
+          as: DriftTableName('ExistingDartClass', true),
+          SelectStatement(
+            columns: [
+              ExpressionResultColumn(
+                expression: NumericLiteral(1, token(TokenType.numberLiteral)),
+              ),
+            ],
+          ),
+        ),
+      ]),
+    );
+  });
 }