Analyze index annotations for tables

drift/lib/src/drift_dev_helper.dart

@@ -5,4 +5,5 @@ export 'runtime/types/converters.dart' show TypeConverter, JsonTypeConverter2;
 export 'runtime/types/mapping.dart' show DriftAny;
 export 'runtime/query_builder/query_builder.dart' show TableInfo;
 
-export 'dsl/dsl.dart' show Table, View, DriftDatabase, DriftAccessor;
+export 'dsl/dsl.dart'
+    show Table, TableIndex, View, DriftDatabase, DriftAccessor;

drift/lib/src/dsl/table.dart

@@ -238,6 +238,25 @@ abstract class View extends HasResultSet {
   Query as();
 }
 
+@Target({TargetKind.classType})
+final class TableIndex {
+  /// The name of the index in SQL.
+  final String name;
+
+  /// Whether this index is `UNIQUE`, meaning that the database will forbid
+  /// multiple rows in the annotated table from having the same values in the
+  /// indexed columns.
+  final bool unique;
+
+  final Set<Symbol> columns;
+
+  const TableIndex({
+    required this.name,
+    required this.columns,
+    this.unique = false,
+  });
+}
+
 /// A class to be used as an annotation on [Table] classes to customize the
 /// name for the data class that will be generated for the table class. The data
 /// class is a dart object that will be used to represent a row in the table.
@@ -295,7 +314,6 @@ class DataClassName {
 
 /// An annotation specifying an existing class to be used as a data class.
 @Target({TargetKind.classType})
-@experimental
 class UseRowClass {
   /// The existing class
   ///

drift_dev/lib/src/analysis/resolver/dart/helper.dart

@@ -22,6 +22,7 @@ class KnownDriftTypes {
   final LibraryElement helperLibrary;
   final ClassElement tableElement;
   final InterfaceType tableType;
+  final InterfaceType tableIndexType;
   final InterfaceType viewType;
   final InterfaceType tableInfoType;
   final InterfaceType driftDatabase;
@@ -35,6 +36,7 @@ class KnownDriftTypes {
     this.helperLibrary,
     this.tableElement,
     this.tableType,
+    this.tableIndexType,
     this.viewType,
     this.tableInfoType,
     this.typeConverter,
@@ -57,6 +59,7 @@ class KnownDriftTypes {
       helper,
       tableElement,
       tableElement.defaultInstantiation,
+      (exportNamespace.get('TableIndex') as InterfaceElement).thisType,
       (exportNamespace.get('View') as InterfaceElement).thisType,
       (exportNamespace.get('TableInfo') as InterfaceElement).thisType,
       exportNamespace.get('TypeConverter') as InterfaceElement,
@@ -207,7 +210,7 @@ class DataClassInformation {
 
     for (final annotation in element.metadata) {
       final computed = annotation.computeConstantValue();
-      final annotationClass = computed!.type!.nameIfInterfaceType;
+      final annotationClass = computed?.type?.nameIfInterfaceType;
 
       if (annotationClass == 'DataClassName') {
         dataClassName = computed;

drift_dev/lib/src/analysis/resolver/dart/index.dart

@@ -0,0 +1,52 @@
+import 'package:analyzer/dart/constant/value.dart';
+import 'package:collection/collection.dart';
+
+import '../../driver/error.dart';
+import '../../results/results.dart';
+import '../intermediate_state.dart';
+import '../resolver.dart';
+
+class DartIndexResolver extends LocalElementResolver<DiscoveredDartIndex> {
+  DartIndexResolver(super.file, super.discovered, super.resolver, super.state);
+
+  @override
+  Future<DriftIndex> resolve() async {
+    // Revive the annotation by parsing values from the computed constant
+    // value.
+    final computed = discovered.annotation.computeConstantValue();
+    final unique = computed?.getField('unique')?.toBoolValue() ?? false;
+
+    final tableResult = await resolver.resolveReferencedElement(
+        discovered.ownId, discovered.onTable);
+    final table = handleReferenceResult<DriftTable>(
+      tableResult,
+      (msg) => DriftAnalysisError.forDartElement(discovered.dartElement, msg),
+    );
+    final columns = <DriftColumn>[];
+
+    final referencedColumns = computed?.getField('columns')?.toSetValue();
+    for (final column in referencedColumns ?? const <DartObject>{}) {
+      final columnName = column.toSymbolValue();
+      final tableColumn =
+          table?.columns.firstWhereOrNull((c) => c.nameInDart == columnName);
+
+      if (tableColumn != null) {
+        columns.add(tableColumn);
+      } else {
+        reportError(DriftAnalysisError.forDartElement(
+          discovered.dartElement,
+          'Column `$columnName`, referenced in index `${discovered.ownId.name}`, was not found in the table.',
+        ));
+      }
+    }
+
+    return DriftIndex(
+      discovered.ownId,
+      DriftDeclaration.dartElement(discovered.dartElement),
+      table: table,
+      indexedColumns: columns,
+      unique: unique,
+      createStmt: null,
+    );
+  }
+}

drift_dev/lib/src/analysis/resolver/discover.dart

@@ -2,6 +2,7 @@ import 'package:analyzer/dart/ast/ast.dart' as dart;
 import 'package:analyzer/dart/constant/value.dart';
 import 'package:analyzer/dart/element/element.dart';
 import 'package:analyzer/dart/element/visitor.dart';
+import 'package:drift/drift.dart' show TableIndex;
 import 'package:source_gen/source_gen.dart';
 import 'package:sqlparser/sqlparser.dart' hide AnalysisError;
 
@@ -159,7 +160,12 @@ class _FindDartElements extends RecursiveElementVisitor<void> {
 
   final List<DriftImport> imports = [];
 
-  final TypeChecker _isTable, _isView, _isTableInfo, _isDatabase, _isDao;
+  final TypeChecker _isTable,
+      _isTableIndex,
+      _isView,
+      _isTableInfo,
+      _isDatabase,
+      _isDao;
 
   final List<Future<void>> _pendingWork = [];
 
@@ -169,6 +175,7 @@ class _FindDartElements extends RecursiveElementVisitor<void> {
   _FindDartElements(
       this._discoverStep, this._library, KnownDriftTypes knownTypes)
       : _isTable = TypeChecker.fromStatic(knownTypes.tableType),
+        _isTableIndex = TypeChecker.fromStatic(knownTypes.tableIndexType),
         _isView = TypeChecker.fromStatic(knownTypes.viewType),
         _isTableInfo = TypeChecker.fromStatic(knownTypes.tableInfoType),
         _isDatabase = TypeChecker.fromStatic(knownTypes.driftDatabase),
@@ -209,8 +216,11 @@ class _FindDartElements extends RecursiveElementVisitor<void> {
         _pendingWork.add(Future.sync(() async {
           final name = await _sqlNameOfTable(element);
           final id = _discoverStep._id(name);
-
           found.add(DiscoveredDartTable(id, element));
+
+          for (final (annotation, indexId) in _tableIndexAnnotation(element)) {
+            found.add(DiscoveredDartIndex(indexId, element, id, annotation));
+          }
         }));
       }
     } else if (_isDslView(element)) {
@@ -261,6 +271,24 @@ class _FindDartElements extends RecursiveElementVisitor<void> {
         .apply(definingElement.name);
   }
 
+  /// Finds a [TableIndex] annotations on the [table].
+  Iterable<(ElementAnnotation, DriftElementId)> _tableIndexAnnotation(
+      ClassElement table) sync* {
+    for (final annotation in table.metadata) {
+      final computed = annotation.computeConstantValue();
+      final type = computed?.type;
+
+      if (computed != null &&
+          type != null &&
+          _isTableIndex.isExactlyType(type)) {
+        yield (
+          annotation,
+          _discoverStep._id(computed.getField('name')?.toStringValue() ?? '')
+        );
+      }
+    }
+  }
+
   DartObject? _driftViewAnnotation(ClassElement view) {
     for (final annotation in view.metadata) {
       final computed = annotation.computeConstantValue();

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

@@ -20,16 +20,29 @@ class DriftIndexResolver extends DriftElementResolver<DiscoveredDriftIndex> {
 
     final onTable = stmt.on.resolved;
     DriftTable? target;
+    List<DriftColumn> indexedColumns = [];
+
     if (onTable is Table) {
       target = references
           .whereType<DriftTable>()
           .firstWhere((e) => e.schemaName == onTable.name);
+
+      for (final indexedColumn in stmt.columns) {
+        final name = (indexedColumn.expression as Reference).columnName;
+        final tableColumn = target.columnBySqlName[name];
+
+        if (tableColumn != null) {
+          indexedColumns.add(tableColumn);
+        }
+      }
     }
 
     return DriftIndex(
       discovered.ownId,
       DriftDeclaration.driftFile(stmt, file.ownUri),
       table: target,
+      indexedColumns: indexedColumns,
+      unique: stmt.unique,
       createStmt: source.substring(stmt.firstPosition, stmt.lastPosition),
     );
   }

drift_dev/lib/src/analysis/resolver/intermediate_state.dart

@@ -51,6 +51,21 @@ class DiscoveredDartView extends DiscoveredDartElement<ClassElement> {
   DiscoveredDartView(super.ownId, super.dartElement, this.viewAnnotation);
 }
 
+class DiscoveredDartIndex extends DiscoveredDartElement<ClassElement> {
+  final DriftElementId onTable;
+
+  ElementAnnotation annotation;
+
+  @override
+  DriftElementKind get kind => DriftElementKind.dbIndex;
+
+  @override
+  String? get dartElementName => null;
+
+  DiscoveredDartIndex(
+      super.ownId, super.dartElement, this.onTable, this.annotation);
+}
+
 class DiscoveredBaseAccessor extends DiscoveredDartElement<ClassElement> {
   final bool isDatabase;
   final DartObject annotation;

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

@@ -8,6 +8,7 @@ import '../results/element.dart';
 
 import '../serializer.dart';
 import 'dart/accessor.dart' as dart_accessor;
+import 'dart/index.dart' as dart_index;
 import 'dart/table.dart' as dart_table;
 import 'dart/view.dart' as dart_view;
 import 'drift/index.dart' as drift_index;
@@ -58,14 +59,13 @@ class DriftResolver {
 
   /// Resolves a discovered element by analyzing it and its dependencies.
   Future<DriftElement> _resolveDiscovered(DiscoveredElement discovered) async {
-    LocalElementResolver resolver;
-
     final fileState = driver.cache.knownFiles[discovered.ownId.libraryUri]!;
     final elementState = fileState.analysis.putIfAbsent(
         discovered.ownId, () => ElementAnalysisState(discovered.ownId));
 
     elementState.errorsDuringAnalysis.clear();
 
+    LocalElementResolver resolver;
     if (discovered is DiscoveredDriftTable) {
       resolver = drift_table.DriftTableResolver(
           fileState, discovered, this, elementState);
@@ -87,6 +87,9 @@ class DriftResolver {
     } else if (discovered is DiscoveredDartView) {
       resolver =
           dart_view.DartViewResolver(fileState, discovered, this, elementState);
+    } else if (discovered is DiscoveredDartIndex) {
+      resolver = dart_index.DartIndexResolver(
+          fileState, discovered, this, elementState);
     } else if (discovered is DiscoveredBaseAccessor) {
       resolver = dart_accessor.DartAccessorResolver(
           fileState, discovered, this, elementState);
@@ -239,7 +242,7 @@ abstract class LocalElementResolver<T extends DiscoveredElement> {
     DriftAnalysisError Function(String msg) createError,
   ) async {
     final result = await resolver.resolveReference(discovered.ownId, reference);
-    return _handleReferenceResult(result, createError);
+    return handleReferenceResult(result, createError);
   }
 
   Future<E?> resolveDartReferenceOrReportError<E extends DriftElement>(
@@ -248,10 +251,10 @@ abstract class LocalElementResolver<T extends DiscoveredElement> {
   ) async {
     final result =
         await resolver.resolveDartReference(discovered.ownId, reference);
-    return _handleReferenceResult(result, createError);
+    return handleReferenceResult(result, createError);
   }
 
-  E? _handleReferenceResult<E extends DriftElement>(
+  E? handleReferenceResult<E extends DriftElement>(
     ResolveReferencedElementResult result,
     DriftAnalysisError Function(String msg) createError,
   ) {

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

@@ -1,7 +1,6 @@
 import 'package:sqlparser/sqlparser.dart';
 
-import 'element.dart';
-import 'table.dart';
+import 'results.dart';
 
 /// An index on a drift table.
 ///
@@ -13,17 +12,25 @@ class DriftIndex extends DriftSchemaElement {
   /// This may be null if the table couldn't be resolved.
   DriftTable? table;
 
-  /// The `CREATE INDEX` SQL statement creating this index, as written down by
-  /// the user.
+  /// Columns of [table] that have been indexed.
+  List<DriftColumn> indexedColumns;
+
+  /// Whethet the index has been declared to be unique.
+  final bool unique;
+
+  /// For indices created in drift files, the `CREATE INDEX` SQL statements as
+  /// written by the user in the drift file.
   ///
-  /// In generated code, another step will reforma this string to strip out
+  /// In generated code, another step will reformat this string to strip out
   /// comments and unncecessary whitespace.
-  final String createStmt;
+  final String? createStmt;
 
   DriftIndex(
     super.id,
     super.declaration, {
     required this.table,
+    required this.indexedColumns,
+    required this.unique,
     required this.createStmt,
   });
 
@@ -42,3 +49,5 @@ class DriftIndex extends DriftSchemaElement {
   /// analysis.
   CreateIndexStatement? parsedStatement;
 }
+
+sealed class DriftIndexDefintion {}

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

@@ -1,4 +1,5 @@
 import 'package:analyzer/dart/element/type.dart';
+import 'package:sqlparser/utils/case_insensitive_map.dart';
 
 import 'element.dart';
 import 'column.dart';
@@ -27,9 +28,9 @@ abstract class DriftElementWithResultSet extends DriftSchemaElement {
   String get nameOfRowClass;
 
   /// All [columns] of this table, indexed by their name in SQL.
-  late final Map<String, DriftColumn> columnBySqlName = {
+  late final Map<String, DriftColumn> columnBySqlName = CaseInsensitiveMap.of({
     for (final column in columns) column.nameInSql: column,
-  };
+  });
 
   /// All type converter applied to columns on this table.
   Iterable<AppliedTypeConverter> get appliedConverters sync* {

drift_dev/lib/src/analysis/serializer.dart

@@ -67,6 +67,10 @@ class ElementSerializer {
       additionalInformation = {
         'type': 'index',
         'sql': element.createStmt,
+        'columns': [
+          for (final column in element.indexedColumns) column.nameInSql,
+        ],
+        'unique': element.unique,
       };
     } else if (element is DefinedSqlQuery) {
       final existingDartType = element.existingDartType;
@@ -531,11 +535,18 @@ class ElementDeserializer {
 
         return table;
       case 'index':
+        final onTable = references.whereType<DriftTable>().firstOrNull;
+
         return DriftIndex(
           id,
           declaration,
-          table: references.whereType<DriftTable>().firstOrNull,
-          createStmt: json['sql'] as String,
+          table: onTable,
+          createStmt: json['sql'] as String?,
+          indexedColumns: [
+            for (final entry in json['columns'] as List)
+              onTable!.columnBySqlName[entry as String]!,
+          ],
+          unique: json['unique'] as bool,
         );
       case 'query':
         final types = <String, DartType>{};

drift_dev/lib/src/writer/schema_version_writer.dart

@@ -5,6 +5,7 @@ import 'package:sqlparser/utils/node_to_text.dart';
 
 import '../analysis/results/results.dart';
 import '../utils/string_escaper.dart';
+import 'database_writer.dart';
 import 'tables/table_writer.dart';
 import 'writer.dart';
 
@@ -274,21 +275,15 @@ class SchemaVersionWriter {
       final index = definition.drift('Index');
 
       definition
-        ..write('final $index $name = $index(')
-        ..write(asDartLiteral(element.schemaName))
-        ..write(',')
-        ..write(asDartLiteral(element.createStmt))
-        ..write(')');
+        ..write('final $index $name = ')
+        ..writeln(DatabaseWriter.createIndex(definition.parent!, element));
     } else if (element is DriftTrigger) {
       name = element.dbGetterName;
       final trigger = definition.drift('Trigger');
 
       definition
-        ..write('final $trigger $name = $trigger(')
-        ..write(asDartLiteral(element.createStmt))
-        ..write(',')
-        ..write(asDartLiteral(element.schemaName))
-        ..write(')');
+        ..write('final $trigger $name = ')
+        ..writeln(DatabaseWriter.createTrigger(definition.parent!, element));
     } else {
       throw ArgumentError('Unhandled element type $element');
     }

drift_dev/test/analysis/resolver/dart/index_test.dart

@@ -0,0 +1,61 @@
+import 'package:drift_dev/src/analysis/results/results.dart';
+import 'package:test/test.dart';
+
+import '../../test_utils.dart';
+
+void main() {
+  test('resolves index', () async {
+    final backend = TestBackend.inTest({
+      'a|lib/a.dart': '''
+import 'package:drift/drift.dart';
+
+@TableIndex(columns: {#a}, name: 'tbl_a')
+@TableIndex(columns: {#b, #c}, name: 'tbl_bc', unique: true)
+class MyTable extends Table {
+  IntColumn get a => integer()();
+  TextColumn get b => text()();
+  TextColumn get c => text()();
+}
+''',
+    });
+
+    final file = await backend.analyze('package:a/a.dart');
+    backend.expectNoErrors();
+
+    final elements = file.analyzedElements;
+    final table = elements.whereType<DriftTable>().first;
+
+    final indexA = file.analysis[file.id('tbl_a')]!.result as DriftIndex;
+    final indexBC = file.analysis[file.id('tbl_bc')]!.result as DriftIndex;
+
+    expect(indexA.table, table);
+    expect(indexA.unique, false);
+    expect(indexA.indexedColumns, [table.columnBySqlName['a']]);
+
+    expect(indexBC.table, table);
+    expect(indexBC.unique, true);
+    expect(indexBC.indexedColumns, [
+      table.columnBySqlName['b'],
+      table.columnBySqlName['c'],
+    ]);
+  });
+
+  test('warns about missing columns', () async {
+    final backend = TestBackend.inTest({
+      'a|lib/a.dart': '''
+import 'package:drift/drift.dart';
+
+@TableIndex(columns: {#foo}, name: 'tbl_a')
+class MyTable extends Table {
+  IntColumn get a => integer()();
+}
+''',
+    });
+
+    final file = await backend.analyze('package:a/a.dart');
+    expect(file.allErrors, [
+      isDriftError(
+          'Column `foo`, referenced in index `tbl_a`, was not found in the table.')
+    ]);
+  });
+}

sqlparser/lib/utils/case_insensitive_map.dart

@@ -4,6 +4,17 @@ import 'dart:collection';
 class CaseInsensitiveMap<K extends String?, T> extends MapBase<K, T> {
   final Map<K, T> _normalized = {};
 
+  CaseInsensitiveMap();
+
+  factory CaseInsensitiveMap.of(Map<K, T> other) {
+    final map = CaseInsensitiveMap<K, T>();
+    other.forEach((key, value) {
+      map[key] = value;
+    });
+
+    return map;
+  }
+
   @override
   T? operator [](Object? key) {
     if (key is String?) {