From e0d5520ee17331c08abb063462f1946984d1ec34 Mon Sep 17 00:00:00 2001
From: Simon Binder <oss@simonbinder.eu>
Date: Mon, 13 Mar 2023 21:51:16 +0100
Subject: [PATCH] Analyze foreign keys as custom constraint on cols

---
 .../src/analysis/resolver/dart/column.dart    | 50 ++++++++++++++++---
 .../analysis/resolver/dart/column_test.dart   | 43 +++++++++++++++-
 2 files changed, 85 insertions(+), 8 deletions(-)

diff --git a/drift_dev/lib/src/analysis/resolver/dart/column.dart b/drift_dev/lib/src/analysis/resolver/dart/column.dart
index b8a837ab..95c0a762 100644
--- a/drift_dev/lib/src/analysis/resolver/dart/column.dart
+++ b/drift_dev/lib/src/analysis/resolver/dart/column.dart
@@ -429,7 +429,7 @@ class ColumnParser {
     final docString =
         getter.documentationComment?.tokens.map((t) => t.toString()).join('\n');
 
-    foundConstraints.addAll(_driftConstraintsFromCustomConstraints(
+    foundConstraints.addAll(await _driftConstraintsFromCustomConstraints(
       isNullable: nullable,
       customConstraints: foundCustomConstraint,
       sourceForCustomConstraints: customConstraintSource,
@@ -484,12 +484,12 @@ class ColumnParser {
     return object.computeConstantValue()!.getField('key')!.toStringValue();
   }
 
-  Iterable<DriftColumnConstraint> _driftConstraintsFromCustomConstraints({
+  Future<List<DriftColumnConstraint>> _driftConstraintsFromCustomConstraints({
     required bool isNullable,
     String? customConstraints,
     AstNode? sourceForCustomConstraints,
-  }) sync* {
-    if (customConstraints == null) return;
+  }) async {
+    if (customConstraints == null) return const [];
 
     final engine = _resolver.resolver.driver.newSqlEngine();
     final parseResult = engine.parseColumnConstraints(customConstraints);
@@ -514,15 +514,51 @@ class ColumnParser {
       ));
     }
 
+    final parsedConstraints = <DriftColumnConstraint>[];
+
     for (final constraint in constraints) {
       if (constraint is sql.GeneratedAs) {
-        yield ColumnGeneratedAs.fromParser(constraint);
+        parsedConstraints.add(ColumnGeneratedAs.fromParser(constraint));
       } else if (constraint is sql.PrimaryKeyColumn) {
-        yield PrimaryKeyColumn(constraint.autoIncrement);
+        parsedConstraints.add(PrimaryKeyColumn(constraint.autoIncrement));
       } else if (constraint is sql.UniqueColumn) {
-        yield UniqueColumn();
+        parsedConstraints.add(UniqueColumn());
+      } else if (constraint is sql.ForeignKeyColumnConstraint) {
+        final clause = constraint.clause;
+
+        final table =
+            await _resolver.resolveSqlReferenceOrReportError<DriftTable>(
+          clause.foreignTable.tableName,
+          (msg) => DriftAnalysisError.inDartAst(
+            _resolver.discovered.dartElement,
+            sourceForCustomConstraints!,
+            msg,
+          ),
+        );
+
+        if (table != null) {
+          final columnName = clause.columnNames.first;
+          final column =
+              table.columnBySqlName[clause.columnNames.first.columnName];
+
+          if (column == null) {
+            _resolver.reportError(DriftAnalysisError.inDartAst(
+              _resolver.discovered.dartElement,
+              sourceForCustomConstraints!,
+              'The referenced table has no column named `$columnName`',
+            ));
+          } else {
+            parsedConstraints.add(ForeignKeyReference(
+              column,
+              constraint.clause.onUpdate,
+              constraint.clause.onDelete,
+            ));
+          }
+        }
       }
     }
+
+    return parsedConstraints;
   }
 }
 
diff --git a/drift_dev/test/analysis/resolver/dart/column_test.dart b/drift_dev/test/analysis/resolver/dart/column_test.dart
index cac5b278..3f5e8b5b 100644
--- a/drift_dev/test/analysis/resolver/dart/column_test.dart
+++ b/drift_dev/test/analysis/resolver/dart/column_test.dart
@@ -1,5 +1,5 @@
 import 'package:drift_dev/src/analysis/options.dart';
-import 'package:drift_dev/src/analysis/results/table.dart';
+import 'package:drift_dev/src/analysis/results/results.dart';
 import 'package:test/test.dart';
 
 import '../../test_utils.dart';
@@ -253,6 +253,47 @@ class TestTable extends Table {
       ]);
     });
 
+    test('resolves foreign key references', () async {
+      final state = TestBackend.inTest({
+        'a|lib/a.dart': '''
+import 'package:drift/drift.dart';
+
+class ReferencedTable extends Table {
+  TextColumn get textColumn => text()();
+}
+
+class TestTable extends Table {
+  TextColumn get a => text().customConstraint('NOT NULL REFERENCES foo (bar)')();
+  TextColumn get b => text().customConstraint('NOT NULL REFERENCES referenced_table (foo)')();
+  TextColumn get c => text().customConstraint('NOT NULL REFERENCES referenced_table (text_column)')();
+}
+''',
+      });
+
+      final file = await state.analyze('package:a/a.dart');
+      final referencedTable =
+          file.analysis[file.id('referenced_table')]!.result! as DriftTable;
+      final tableAnalysis = file.analysis[file.id('test_table')]!;
+
+      expect(tableAnalysis.errorsDuringAnalysis, [
+        isDriftError('`foo` could not be found in any import.')
+            .withSpan(contains('REFERENCES foo (bar)')),
+        isDriftError(contains('has no column named `foo`'))
+            .withSpan(contains('referenced_table (foo)')),
+      ]);
+
+      final testTable = tableAnalysis.result! as DriftTable;
+      expect(
+        testTable.columnBySqlName['c'],
+        isA<DriftColumn>().having(
+          (e) => e.constraints,
+          'constraints',
+          contains(isA<ForeignKeyReference>().having((e) => e.otherColumn,
+              'otherColumn', referencedTable.columns.single)),
+        ),
+      );
+    });
+
     test('warns about missing `NOT NULL`', () async {
       final state = TestBackend.inTest({
         'a|lib/a.dart': '''