feat(datastore): support for custom primary key (#1752)

* feat(datastore): custom pk

* chore: update API calls

* address comments

* address comments, update models

* update pk encoding

* fix(datastore): create indexes for custom primary key (#1828)

* precondition failure in case of identifier field not found

* test(datastore): custom pk tests (#1799)

* test(datastore): custom PK tests

* update models

* more SQL statements tests

* add CascadeDeleteOperation tests with composite pk

* address comments

* feat(datastore): custom pk associations (#1811)

* test(datastore): custom PK tests

* update models

* update pk encoding

* fix associations

* feat(datastore): support for CPK in associations

* restore associations test models

* rebase on tests

* feat(datastore): support for CPK in associations

* add more tests

* add schema for test models;

* rename schema.graphql to avoid build failures

* address comments

* fix tests when pk field is enum

* tests with composite keys

* resolve associated model

* update test models

* feat(datastore): custom primary key - add more tests, fix cascade delete (#1878)

* feat(datastore): add more tests, fix cascade delete

* fix typo

* precondition failure only when associated field is required

* add tests for query predicate

* fix(datastore): avoid creating index on foreign key fields (#1937)

* fix(datastore): resolve PK value in JSONValue models (#1949)

* fix quoting on table names in unit test

* fix(datastore): Fix ModelSyncedEventEmitter, add lifecycle events logging (#2022)

* try to fix ready event emitter

* lifecycle event logging

* fix(datastore): Add default value for where predicate (#2041)

* fix(datastore): Add default value for where predicate

* Indentation fixtures for where predicate

* Fix-final changes

* fix(datastore): String conform to ModelIdentifierProtocol (#2067)

Co-authored-by: Michael Law <[email protected]>
Co-authored-by: Manaswi Manthena <[email protected]>

Author: 3 people

Amplify.xcodeproj/project.pbxproj

@@ -31,6 +31,7 @@ public extension DataStoreBaseBehavior {
     ///   - modelType: The type of the model to delete
     ///   - id: The ID of the model to delete
     /// - Returns: A DataStorePublisher with the results of the operation
+    @available(*, deprecated, message: "Use delete(:withIdentifier:where:)")
     func delete<M: Model>(
         _ modelType: M.Type,
         withId id: String,
@@ -41,6 +42,40 @@ public extension DataStoreBaseBehavior {
         }.eraseToAnyPublisher()
     }
 
+    /// Deletes the model with the specified identifier from the DataStore. If sync is enabled, this will delete the
+    /// model from the remote store as well.
+    ///
+    /// - Parameters:
+    ///   - modelType: The type of the model to delete
+    ///   - identifier: The identifier of the model to delete
+    /// - Returns: A DataStorePublisher with the results of the operation
+    func delete<M: Model>(
+        _ modelType: M.Type,
+        withIdentifier identifier: String,
+        where predicate: QueryPredicate? = nil
+    )  -> DataStorePublisher<Void> where M: ModelIdentifiable, M.IdentifierFormat == ModelIdentifierFormat.Default {
+        Future { promise in
+            self.delete(modelType, withIdentifier: identifier, where: predicate) { promise($0) }
+        }.eraseToAnyPublisher()
+    }
+
+    /// Deletes the model with the specified identifier from the DataStore. If sync is enabled, this will delete the
+    /// model from the remote store as well.
+    ///
+    /// - Parameters:
+    ///   - modelType: The type of the model to delete
+    ///   - identifier: The identifier of the model to delete
+    /// - Returns: A DataStorePublisher with the results of the operation
+    func delete<M: Model>(
+        _ modelType: M.Type,
+        withIdentifier identifier: ModelIdentifier<M, M.IdentifierFormat>,
+        where predicate: QueryPredicate? = nil
+    )  -> DataStorePublisher<Void> where M: ModelIdentifiable {
+        Future { promise in
+            self.delete(modelType, withIdentifier: identifier, where: predicate) { promise($0) }
+        }.eraseToAnyPublisher()
+    }
+
     /// Deletes models matching the `predicate` from the DataStore. If sync is enabled, this will delete the
     /// model from the remote store as well.
     ///
@@ -79,6 +114,7 @@ public extension DataStoreBaseBehavior {
     ///   - modelType: The type of the model to query
     ///   - id: The ID of the model to query
     /// - Returns: A DataStorePublisher with the results of the operation
+    @available(*, deprecated, message: "Use query(:byIdentifier:)")
     func query<M: Model>(
         _ modelType: M.Type,
         byId id: String
@@ -88,6 +124,36 @@ public extension DataStoreBaseBehavior {
         }.eraseToAnyPublisher()
     }
 
+    /// Queries for a specific model instance by id
+    ///
+    /// - Parameters:
+    ///   - modelType: The type of the model to query
+    ///   - id: The ID of the model to query
+    /// - Returns: A DataStorePublisher with the results of the operation
+    func query<M: Model>(
+        _ modelType: M.Type,
+        byIdentifier identifier: String
+    ) -> DataStorePublisher<M?> where M: ModelIdentifiable, M.IdentifierFormat == ModelIdentifierFormat.Default {
+        Future { promise in
+            self.query(modelType, byIdentifier: identifier) { promise($0) }
+        }.eraseToAnyPublisher()
+    }
+
+    /// Queries for a specific model instance by id
+    ///
+    /// - Parameters:
+    ///   - modelType: The type of the model to query
+    ///   - id: The ID of the model to query
+    /// - Returns: A DataStorePublisher with the results of the operation
+    func query<M: Model>(
+        _ modelType: M.Type,
+        byIdentifier identifier: ModelIdentifier<M, M.IdentifierFormat>
+    ) -> DataStorePublisher<M?> where M: ModelIdentifiable {
+        Future { promise in
+            self.query(modelType, byIdentifier: identifier) { promise($0) }
+        }.eraseToAnyPublisher()
+    }
+
     /// Queries for any model instances that match the specified predicate
     ///
     /// - Parameters:

Amplify/Categories/DataStore/DataStoreCategory+Behavior+Combine.swift

@@ -18,6 +18,19 @@ extension DataStoreCategory: DataStoreBaseBehavior {
         plugin.query(modelType, byId: id, completion: completion)
     }
 
+    public func query<M: Model>(_ modelType: M.Type,
+                                byIdentifier id: String,
+                                completion: (DataStoreResult<M?>) -> Void) where M: ModelIdentifiable,
+                                                                                 M.IdentifierFormat == ModelIdentifierFormat.Default {
+        plugin.query(modelType, byIdentifier: id, completion: completion)
+    }
+
+    public func query<M: Model>(_ modelType: M.Type,
+                                byIdentifier identifier: ModelIdentifier<M, M.IdentifierFormat>,
+                                completion: (DataStoreResult<M?>) -> Void) where M: ModelIdentifiable {
+        plugin.query(modelType, byIdentifier: identifier, completion: completion)
+    }
+
     public func query<M: Model>(_ modelType: M.Type,
                                 where predicate: QueryPredicate? = nil,
                                 sort sortInput: QuerySortInput? = nil,
@@ -39,6 +52,22 @@ extension DataStoreCategory: DataStoreBaseBehavior {
         plugin.delete(modelType, withId: id, where: predicate, completion: completion)
     }
 
+    public func delete<M: Model>(_ modelType: M.Type,
+                                 withIdentifier id: String,
+                                 where predicate: QueryPredicate? = nil,
+                                 completion: @escaping DataStoreCallback<Void>) where M: ModelIdentifiable,
+                                                                                      M.IdentifierFormat == ModelIdentifierFormat.Default {
+
+        plugin.delete(modelType, withIdentifier: id, where: predicate, completion: completion)
+    }
+
+    public func delete<M: Model>(_ modelType: M.Type,
+                                 withIdentifier id: ModelIdentifier<M, M.IdentifierFormat>,
+                                 where predicate: QueryPredicate? = nil,
+                                 completion: @escaping DataStoreCallback<Void>) where M: ModelIdentifiable {
+        plugin.delete(modelType, withIdentifier: id, where: predicate, completion: completion)
+    }
+
     public func delete<M: Model>(_ modelType: M.Type,
                                  where predicate: QueryPredicate,
                                  completion: @escaping DataStoreCallback<Void>) {

Amplify/Categories/DataStore/DataStoreCategory+Behavior.swift

@@ -16,10 +16,20 @@ public protocol DataStoreBaseBehavior {
                         where condition: QueryPredicate?,
                         completion: @escaping DataStoreCallback<M>)
 
+    @available(*, deprecated, message: "Use query(:byIdentifier:completion)")
     func query<M: Model>(_ modelType: M.Type,
                          byId id: String,
                          completion: DataStoreCallback<M?>)
 
+    func query<M: Model>(_ modelType: M.Type,
+                         byIdentifier id: String,
+                         completion: DataStoreCallback<M?>) where M: ModelIdentifiable,
+                                                                  M.IdentifierFormat == ModelIdentifierFormat.Default
+
+    func query<M: Model>(_ modelType: M.Type,
+                         byIdentifier id: ModelIdentifier<M, M.IdentifierFormat>,
+                         completion: DataStoreCallback<M?>) where M: ModelIdentifiable
+
     func query<M: Model>(_ modelType: M.Type,
                          where predicate: QueryPredicate?,
                          sort sortInput: QuerySortInput?,
@@ -30,14 +40,26 @@ public protocol DataStoreBaseBehavior {
                           where predicate: QueryPredicate?,
                           completion: @escaping DataStoreCallback<Void>)
 
+    @available(*, deprecated, message: "Use delete(:withIdentifier:where:completion)")
     func delete<M: Model>(_ modelType: M.Type,
                           withId id: String,
                           where predicate: QueryPredicate?,
                           completion: @escaping DataStoreCallback<Void>)
 
     func delete<M: Model>(_ modelType: M.Type,
-                          where predicate: QueryPredicate,
-                          completion: @escaping DataStoreCallback<Void>)
+                          withIdentifier id: String,
+                          where predicate: QueryPredicate?,
+                          completion: @escaping DataStoreCallback<Void>) where M: ModelIdentifiable,
+                                                                               M.IdentifierFormat == ModelIdentifierFormat.Default
+
+    func delete<M: Model>(_ modelType: M.Type,
+                          withIdentifier id: ModelIdentifier<M, M.IdentifierFormat>,
+                          where predicate: QueryPredicate?,
+                          completion: @escaping DataStoreCallback<Void>) where M: ModelIdentifiable
+
+    func delete<M: Model>(_ modelType: M.Type,
+                           where predicate: QueryPredicate,
+                           completion: @escaping DataStoreCallback<Void>)
 
     /**
      Synchronization starts automatically whenever you run any DataStore operation (query(), save(), delete())

Amplify/Categories/DataStore/DataStoreCategoryBehavior.swift

@@ -88,25 +88,38 @@ import Foundation
 ///   directly by host applications. The behavior of this may change without warning.
 public enum ModelAssociation {
     case hasMany(associatedFieldName: String?)
-    case hasOne(associatedFieldName: String?, targetName: String? = nil)
-    case belongsTo(associatedFieldName: String?, targetName: String?)
+    case hasOne(associatedFieldName: String?, targetNames: [String])
+    case belongsTo(associatedFieldName: String?, targetNames: [String])
 
-    public static let belongsTo: ModelAssociation = .belongsTo(associatedFieldName: nil, targetName: nil)
+    public static let belongsTo: ModelAssociation = .belongsTo(associatedFieldName: nil, targetNames: [])
 
     public static func belongsTo(targetName: String? = nil) -> ModelAssociation {
-        return .belongsTo(associatedFieldName: nil, targetName: nil)
+        let targetNames = targetName.map { [$0] } ?? []
+        return .belongsTo(associatedFieldName: nil, targetNames: targetNames)
     }
 
     public static func hasMany(associatedWith: CodingKey?) -> ModelAssociation {
         return .hasMany(associatedFieldName: associatedWith?.stringValue)
     }
 
+    @available(*, deprecated, message: "Use hasOne(associatedWith:targetNames:)")
     public static func hasOne(associatedWith: CodingKey?, targetName: String? = nil) -> ModelAssociation {
-        return .hasOne(associatedFieldName: associatedWith?.stringValue, targetName: targetName)
+        let targetNames = targetName.map { [$0] } ?? []
+        return .hasOne(associatedWith: associatedWith, targetNames: targetNames)
     }
 
+    public static func hasOne(associatedWith: CodingKey?, targetNames: [String] = []) -> ModelAssociation {
+        return .hasOne(associatedFieldName: associatedWith?.stringValue, targetNames: targetNames)
+    }
+
+    @available(*, deprecated, message: "Use belongsTo(associatedWith:targetNames:)")
     public static func belongsTo(associatedWith: CodingKey?, targetName: String?) -> ModelAssociation {
-        return .belongsTo(associatedFieldName: associatedWith?.stringValue, targetName: targetName)
+        let targetNames = targetName.map { [$0] } ?? []
+        return .belongsTo(associatedFieldName: associatedWith?.stringValue, targetNames: targetNames)
+    }
+
+    public static func belongsTo(associatedWith: CodingKey?, targetNames: [String] = []) -> ModelAssociation {
+        return .belongsTo(associatedFieldName: associatedWith?.stringValue, targetNames: targetNames)
     }
 
 }

Amplify/Categories/DataStore/Model/Internal/Schema/ModelField+Association.swift

@@ -0,0 +1,95 @@
+//
+// Copyright Amazon.com Inc. or its affiliates.
+// All Rights Reserved.
+//
+// SPDX-License-Identifier: Apache-2.0
+//
+
+public struct ModelPrimaryKey {
+    public var fields: [ModelField] = []
+    private var fieldsLookup: Set<ModelFieldName> = []
+
+    public var isCompositeKey: Bool {
+        fields.count > 1
+    }
+
+    init?(allFields: ModelFields,
+          attributes: [ModelAttribute],
+          primaryKeyFieldKeys: [String] = []) {
+        self.fields = resolvePrimaryKeyFields(allFields: allFields,
+                                              attributes: attributes,
+                                              primaryKeyFieldKeys: primaryKeyFieldKeys)
+
+        if fields.isEmpty {
+            return nil
+        }
+
+        self.fieldsLookup = Set(fields.map { $0.name })
+    }
+
+    /// Returns the list of fields that make up the primary key for the model.
+    /// In case of a custom primary key, the model has a `@key` directive
+    /// without a name and at least 1 field
+    func primaryFieldsFromIndexes(attributes: [ModelAttribute]) -> [ModelFieldName]? {
+        attributes.compactMap {
+            if case let .index(fields, name) = $0, name == nil, fields.count >= 1 {
+                return fields
+            }
+            return nil
+        }.first
+    }
+
+    /// Resolve the model fields that are part of the primary key.
+    /// For backward compatibility with different versions of the codegen,
+    /// the algorithm tries first to resolve them using first the `primaryKeyFields`
+    /// received from `primaryKey` member set in `ModelSchemaDefinition`,
+    /// if not available tries to infer the fields using the `.indexes` and it eventually
+    /// falls back on the `.primaryKey` attribute.
+    ///
+    /// It returns an array of fields as custom and composite primary keys are supported.
+    /// - Parameter fields: schema model fields
+    /// - Returns: an array of model fields
+    func resolvePrimaryKeyFields(allFields: ModelFields,
+                                 attributes: [ModelAttribute],
+                                 primaryKeyFieldKeys: [String]) -> [ModelField] {
+        var primaryKeyFields: [ModelField] = []
+
+        if !primaryKeyFieldKeys.isEmpty {
+            primaryKeyFields = primaryKeyFieldKeys.map {
+                guard let field = allFields[$0] else {
+                    preconditionFailure("Primary key field named (\($0)) not found in schema fields.")
+                }
+                return field
+            }
+
+        /// if indexes aren't defined most likely the model has a default `id` as PK
+        /// so we have to rely on the `.primaryKey` attribute of each individual field
+        } else if attributes.indexes.filter({ $0.isPrimaryKeyIndex }).isEmpty {
+            primaryKeyFields = allFields.values.filter { $0.isPrimaryKey }
+
+        /// Use the array of fields with a primary key index
+        } else if let fieldNames = primaryFieldsFromIndexes(attributes: attributes) {
+            primaryKeyFields = fieldNames.compactMap {
+                if let field = allFields[$0] {
+                    return field
+                }
+                return nil
+            }
+        }
+        return primaryKeyFields
+    }
+
+    /// Convenience method to verify if a field is part of the primary key
+    /// - Parameter name: field name
+    /// - Returns: true if the field is part of the primary key
+    public func contains(named name: ModelFieldName) -> Bool {
+        fieldsLookup.contains(name)
+    }
+
+    /// Returns the first index in which a model field name of the collection
+    /// is equal to the provided `name`
+    /// Returns `nil` if no element was found.
+    public func indexOfField(named name: ModelFieldName) -> Int? {
+        fields.firstIndex(where: { $0.name == name })
+    }
+}