Read primary key in drift files

drift/test/generated/custom_tables.g.dart

@@ -853,11 +853,11 @@ class MytableCompanion extends UpdateCompanion<MytableData> {
     this.somedate = const Value.absent(),
   });
   MytableCompanion.insert({
-    required int someid,
+    this.someid = const Value.absent(),
     this.sometext = const Value.absent(),
     this.isInserting = const Value.absent(),
     this.somedate = const Value.absent(),
-  }) : someid = Value(someid);
+  });
   static Insertable<MytableData> custom({
     Expression<int>? someid,
     Expression<String>? sometext,
@@ -923,7 +923,7 @@ class $MytableTable extends Table with TableInfo<$MytableTable, MytableData> {
   final VerificationMeta _someidMeta = const VerificationMeta('someid');
   late final GeneratedColumn<int> someid = GeneratedColumn<int>(
       'someid', aliasedName, false,
-      type: DriftSqlType.int, requiredDuringInsert: true);
+      type: DriftSqlType.int, requiredDuringInsert: false);
   final VerificationMeta _sometextMeta = const VerificationMeta('sometext');
   late final GeneratedColumn<String> sometext = GeneratedColumn<String>(
       'sometext', aliasedName, true,
@@ -954,8 +954,6 @@ class $MytableTable extends Table with TableInfo<$MytableTable, MytableData> {
     if (data.containsKey('someid')) {
       context.handle(_someidMeta,
           someid.isAcceptableOrUnknown(data['someid']!, _someidMeta));
-    } else if (isInserting) {
-      context.missing(_someidMeta);
     }
     if (data.containsKey('sometext')) {
       context.handle(_sometextMeta,
@@ -976,7 +974,7 @@ class $MytableTable extends Table with TableInfo<$MytableTable, MytableData> {
   }
 
   @override
-  Set<GeneratedColumn> get $primaryKey => const <GeneratedColumn>{};
+  Set<GeneratedColumn> get $primaryKey => {someid};
   @override
   MytableData map(Map<String, dynamic> data, {String? tablePrefix}) {
     final effectivePrefix = tablePrefix != null ? '$tablePrefix.' : '';

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

@@ -56,8 +56,8 @@ class DartTableResolver extends LocalElementResolver<DiscoveredDartTable> {
       existingRowClass: dataClassInfo.existingClass,
       customParentClass: dataClassInfo.extending,
       baseDartName: element.name,
-      primaryKeyFromTableConstraint: primaryKey,
       tableConstraints: [
+        if (primaryKey != null) PrimaryKeyColumns(primaryKey),
         for (final uniqueKey in uniqueKeys ?? const <Set<DriftColumn>>[])
           UniqueColumns(uniqueKey),
       ],

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

@@ -176,6 +176,22 @@ class DriftTableResolver extends LocalElementResolver<DiscoveredDriftTable> {
               onDelete: constraint.clause.onDelete,
             ));
           }
+        } else if (constraint is KeyClause) {
+          final keyColumns = <DriftColumn>{};
+
+          for (final keyColumn in constraint.columns) {
+            final expression = keyColumn.expression;
+            if (expression is Reference) {
+              keyColumns.add(columns
+                  .firstWhere((e) => e.nameInSql == expression.columnName));
+            }
+          }
+
+          if (constraint.isPrimaryKey) {
+            tableConstraints.add(PrimaryKeyColumns(keyColumns));
+          } else {
+            tableConstraints.add(UniqueColumns(keyColumns));
+          }
         }
       }
     } else if (stmt is CreateVirtualTableStatement) {

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

@@ -1,3 +1,4 @@
+import 'package:collection/collection.dart';
 import 'package:drift/drift.dart' show DriftSqlType;
 import 'package:json_annotation/json_annotation.dart';
 import 'package:sqlparser/sqlparser.dart' as sql;
@@ -16,8 +17,6 @@ class DriftTable extends DriftElementWithResultSet {
 
   final List<DriftTableConstraint> tableConstraints;
 
-  final Set<DriftColumn>? primaryKeyFromTableConstraint;
-
   @override
   final List<DriftElement> references;
 
@@ -69,7 +68,6 @@ class DriftTable extends DriftElementWithResultSet {
     this.fixedEntityInfoName,
     this.withoutRowId = false,
     this.strict = false,
-    this.primaryKeyFromTableConstraint,
     this.tableConstraints = const [],
     this.virtualTableData,
     this.writeDefaultConstraints = true,
@@ -83,12 +81,14 @@ class DriftTable extends DriftElementWithResultSet {
   @override
   String get dbGetterName => DriftSchemaElement.dbFieldName(baseDartName);
 
-  /// The primary key for this table, computed by looking at the
-  /// [primaryKeyFromTableConstraint] and primary key constraints applied to
-  /// individiual columns.
+  /// The primary key for this table, computed by looking at the primary key
+  /// defined as a table constraint or as a column constraint.
   Set<DriftColumn> get fullPrimaryKey {
-    if (primaryKeyFromTableConstraint != null) {
-      return primaryKeyFromTableConstraint!;
+    final fromTable =
+        tableConstraints.whereType<PrimaryKeyColumns>().firstOrNull;
+
+    if (fromTable != null) {
+      return fromTable.primaryKey;
     }
 
     return columns
@@ -149,6 +149,12 @@ class UniqueColumns extends DriftTableConstraint {
   UniqueColumns(this.uniqueSet);
 }
 
+class PrimaryKeyColumns extends DriftTableConstraint {
+  final Set<DriftColumn> primaryKey;
+
+  PrimaryKeyColumns(this.primaryKey);
+}
+
 class ForeignKeyTable extends DriftTableConstraint {
   final List<DriftColumn> localColumns;
   final DriftTable otherTable;

drift_dev/lib/src/analysis/serializer.dart

@@ -27,9 +27,6 @@ class ElementSerializer {
           for (final column in element.columns) _serializeColumn(column),
         ],
         'existing_data_class': element.existingRowClass?.toJson(),
-        'primary_key_table_constraint': element.primaryKeyFromTableConstraint
-            ?.map((e) => e.nameInSql)
-            .toList(),
         'table_constraints': [
           for (final constraint in element.tableConstraints)
             _serializeTableConstraint(constraint),
@@ -203,6 +200,13 @@ class ElementSerializer {
         'type': 'unique',
         'columns': [for (final column in constraint.uniqueSet) column.nameInSql]
       };
+    } else if (constraint is PrimaryKeyColumns) {
+      return {
+        'type': 'primary_key',
+        'columns': [
+          for (final column in constraint.primaryKey) column.nameInSql,
+        ],
+      };
     } else if (constraint is ForeignKeyTable) {
       return {
         'type': 'foreign',
@@ -411,14 +415,6 @@ class ElementDeserializer {
           for (final column in columns) column.nameInSql: column,
         };
 
-        Set<DriftColumn>? primaryKeyFromTableConstraint;
-        final serializedPk = json['primary_key_table_constraint'];
-        if (serializedPk != null) {
-          primaryKeyFromTableConstraint = {
-            for (final entry in serializedPk) columnByName[entry]!,
-          };
-        }
-
         return DriftTable(
           id,
           declaration,
@@ -427,7 +423,6 @@ class ElementDeserializer {
           existingRowClass: json['existing_data_class'] != null
               ? ExistingRowClass.fromJson(json['existing_data_class'] as Map)
               : null,
-          primaryKeyFromTableConstraint: primaryKeyFromTableConstraint,
           tableConstraints: [
             for (final constraint in json['table_constraints'])
               await _readTableConstraint(constraint as Map, columnByName),
@@ -667,6 +662,10 @@ class ElementDeserializer {
         return UniqueColumns({
           for (final ref in json['columns']) localColumns[ref]!,
         });
+      case 'primary_key':
+        return PrimaryKeyColumns(
+          {for (final ref in json['columns']) localColumns[ref]!},
+        );
       case 'foreign':
         return ForeignKeyTable(
           localColumns: [