fix(DataStore): dynamic model support for cascade delete (#1296)

* fix(DataStore): dynamic model support for cascade delete

* PR comments

Author: lawmicha

Amplify/Categories/DataStore/Subscribe/MutationEvent+Model.swift

@@ -15,12 +15,21 @@ extension MutationEvent {
             throw dataStoreError
         }
 
+        try self.init(untypedModel: model,
+                      modelName: modelType.schema.name,
+                      mutationType: mutationType,
+                      version: version)
+    }
+  
+    public init(untypedModel model: Model,
+                modelName: ModelName,
+                mutationType: MutationType,
+                version: Int? = nil) throws {
         let json = try model.toJSON()
         self.init(modelId: model.id,
-                  modelName: modelType.schema.name,
+                  modelName: modelName,
                   json: json,
                   mutationType: mutationType,
                   version: version)
     }
-
 }

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/Statement+AnyModel.swift

@@ -10,20 +10,20 @@ import SQLite
 import Foundation
 
 extension Statement {
-    func convert(toUntypedModel modelType: Model.Type,
-                 using statement: SelectStatement) throws -> [Model] {
+    func convertToUntypedModel(using modelSchema: ModelSchema,
+                               statement: SelectStatement) throws -> [Model] {
         var models = [Model]()
 
         for row in self {
-            let modelValues = try convert(row: row, to: modelType, using: statement)
-            let untypedModel = try convert(toAnyModel: modelType, modelDictionary: modelValues)
+            let modelValues = try convert(row: row, withSchema: modelSchema, using: statement)
+            let untypedModel = try convertToAnyModel(using: modelSchema, modelDictionary: modelValues)
             models.append(untypedModel)
         }
 
         return models
     }
 
-    private func convert(toAnyModel modelType: Model.Type, modelDictionary: ModelValues) throws -> Model {
+    private func convertToAnyModel(using modelSchema: ModelSchema, modelDictionary: ModelValues) throws -> Model {
         let data = try JSONSerialization.data(withJSONObject: modelDictionary)
         guard let jsonString = String(data: data, encoding: .utf8) else {
             let error = DataStoreError.decodingError(
@@ -36,7 +36,7 @@ extension Statement {
             throw error
         }
 
-        let instance = try ModelRegistry.decode(modelName: modelType.modelName, from: jsonString)
+        let instance = try ModelRegistry.decode(modelName: modelSchema.name, from: jsonString)
         return instance
     }
 }

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/Statement+Model.swift

@@ -36,9 +36,11 @@ protocol StatementModelConvertible {
     ///
     /// - Parameters:
     ///   - modelType - the target `Model` type
+    ///   - modelSchema - the schema for `Model`
     ///   - statement - the query executed that generated this result
     /// - Returns: an array of `Model` of the specified type
     func convert<M: Model>(to modelType: M.Type,
+                           withSchema modelSchema: ModelSchema,
                            using statement: SelectStatement) throws -> [M]
 
 }
@@ -51,12 +53,13 @@ extension Statement: StatementModelConvertible {
     }
 
     func convert<M: Model>(to modelType: M.Type,
+                           withSchema modelSchema: ModelSchema,
                            using statement: SelectStatement) throws -> [M] {
         var elements: [ModelValues] = []
 
         // parse each row of the result
         for row in self {
-            let modelDictionary = try convert(row: row, to: modelType, using: statement)
+            let modelDictionary = try convert(row: row, withSchema: modelSchema, using: statement)
             elements.append(modelDictionary)
         }
 
@@ -66,7 +69,7 @@ extension Statement: StatementModelConvertible {
     }
 
     func convert(row: Element,
-                 to modelType: Model.Type,
+                 withSchema modelSchema: ModelSchema,
                  using statement: SelectStatement) throws -> ModelValues {
         let columnMapping = statement.metadata.columnMapping
         let modelDictionary = ([:] as ModelValues).mutableCopy()
@@ -75,7 +78,7 @@ extension Statement: StatementModelConvertible {
             guard let (schema, field) = columnMapping[column] else {
                 logger.debug("""
                 A column named \(column) was found in the result set but no field on
-                \(modelType.modelName) could be found with that name and it will be ignored.
+                \(modelSchema.name) could be found with that name and it will be ignored.
                 """)
                 continue
             }

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/StorageEngineAdapter+SQLite.swift

@@ -233,7 +233,9 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
                                             sort: sort,
                                             paginationInput: paginationInput)
             let rows = try connection.prepare(statement.stringValue).run(statement.variables)
-            let result: [M] = try rows.convert(to: modelType, using: statement)
+            let result: [M] = try rows.convert(to: modelType,
+                                               withSchema: modelSchema,
+                                               using: statement)
             completion(.success(result))
         } catch {
             completion(.failure(causedBy: error))
@@ -286,6 +288,7 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
         let rows = try connection.prepare(sql).bind(ids)
 
         let syncMetadataList = try rows.convert(to: MutationSyncMetadata.self,
+                                                withSchema: MutationSyncMetadata.schema,
                                                 using: statement)
         let mutationSyncList = try syncMetadataList.map { syncMetadata -> MutationSync<AnyModel> in
             guard let model = modelById[syncMetadata.id] else {
@@ -304,6 +307,7 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
 
     func queryMutationSyncMetadata(for modelIds: [Model.Identifier]) throws -> [MutationSyncMetadata] {
         let modelType = MutationSyncMetadata.self
+        let modelSchema = MutationSyncMetadata.schema
         let fields = MutationSyncMetadata.keys
         var results = [MutationSyncMetadata]()
         let chunkedModelIdsArr = modelIds.chunked(into: SQLiteStorageEngineAdapter.maxNumberOfPredicates)
@@ -313,9 +317,10 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
                 queryPredicates.append(QueryPredicateOperation(field: fields.id.stringValue, operator: .equals(id)))
             }
             let groupedQueryPredicates = QueryPredicateGroup(type: .or, predicates: queryPredicates)
-            let statement = SelectStatement(from: modelType.schema, predicate: groupedQueryPredicates)
+            let statement = SelectStatement(from: modelSchema, predicate: groupedQueryPredicates)
             let rows = try connection.prepare(statement.stringValue).run(statement.variables)
             let result = try rows.convert(to: modelType,
+                                          withSchema: modelSchema,
                                           using: statement)
             results.append(contentsOf: result)
         }
@@ -327,6 +332,7 @@ final class SQLiteStorageEngineAdapter: StorageEngineAdapter {
                                         predicate: field("id").eq(modelSchema.name))
         let rows = try connection.prepare(statement.stringValue).run(statement.variables)
         let result = try rows.convert(to: ModelSyncMetadata.self,
+                                      withSchema: ModelSyncMetadata.schema,
                                       using: statement)
         return try result.unique()
     }

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/SQLite/StorageEngineAdapter+UntypedModel.swift

@@ -43,13 +43,13 @@ extension SQLiteStorageEngineAdapter {
         }
     }
 
-    func query(untypedModel modelType: Model.Type,
+    func query(modelSchema: ModelSchema,
                predicate: QueryPredicate? = nil,
                completion: DataStoreCallback<[Model]>) {
         do {
-            let statement = SelectStatement(from: modelType.schema, predicate: predicate)
+            let statement = SelectStatement(from: modelSchema, predicate: predicate)
             let rows = try connection.prepare(statement.stringValue).run(statement.variables)
-            let result: [Model] = try rows.convert(toUntypedModel: modelType, using: statement)
+            let result: [Model] = try rows.convertToUntypedModel(using: modelSchema, statement: statement)
             completion(.success(result))
         } catch {
             completion(.failure(causedBy: error))

AmplifyPlugins/DataStore/AWSDataStoreCategoryPlugin/Storage/StorageEngine+DeleteTransaction.swift

@@ -32,10 +32,10 @@ extension StorageEngine {
 
     func queryAndDeleteTransaction<M: Model>(_ modelType: M.Type,
                                              modelSchema: ModelSchema,
-                                             deleteInput: DeleteInput) -> DataStoreResult<([M], [Model])> {
+                                             deleteInput: DeleteInput) -> DataStoreResult<([M], [ModelName: [Model]])> {
         var queriedResult: DataStoreResult<[M]>?
         var deletedResult: DataStoreResult<[M]>?
-        var associatedModels: [Model] = []
+        var associatedModels: [(ModelName, Model)] = []
 
         let queryCompletionBlock: DataStoreCallback<[M]> = { queryResult in
             queriedResult = queryResult
@@ -93,11 +93,11 @@ extension StorageEngine {
 
         return collapseResults(queryResult: queriedResult,
                                deleteResult: deletedResult,
-                               associatedModels: associatedModels)
+                               associatedModelsMap: getAssociatedModelsMap(associatedModels: associatedModels))
     }
 
-    func recurseQueryAssociatedModels(modelSchema: ModelSchema, ids: [Model.Identifier]) -> [Model] {
-        var associatedModels: [Model] = []
+    func recurseQueryAssociatedModels(modelSchema: ModelSchema, ids: [Model.Identifier]) -> [(ModelName, Model)] {
+        var associatedModels: [(ModelName, Model)] = []
         for (_, modelField) in modelSchema.fields {
             guard modelField.hasAssociation,
                 modelField.isOneToOne || modelField.isOneToMany,
@@ -109,57 +109,69 @@ extension StorageEngine {
             let queriedModels = queryAssociatedModels(modelName: associatedModelName,
                                                       associatedField: associatedField,
                                                       ids: ids)
-            let associatedModelIds = queriedModels.map {$0.id}
+            let associatedModelIds = queriedModels.map { $0.1.id }
             associatedModels.append(contentsOf: queriedModels)
-            associatedModels.append(contentsOf: recurseQueryAssociatedModels(modelSchema: associatedModelSchema, ids: associatedModelIds))
+            associatedModels.append(contentsOf: recurseQueryAssociatedModels(modelSchema: associatedModelSchema,
+                                                                            ids: associatedModelIds))
         }
         return associatedModels
     }
 
-    func queryAssociatedModels(modelName: ModelName, associatedField: ModelField, ids: [Model.Identifier]) -> [Model] {
-        guard let modelType = ModelRegistry.modelType(from: modelName) else {
+    func queryAssociatedModels(modelName: ModelName,
+                               associatedField: ModelField,
+                               ids: [Model.Identifier]) -> [(ModelName, Model)] {
+        guard let modelSchema = ModelRegistry.modelSchema(from: modelName) else {
             log.error("Failed to lookup \(modelName)")
             return []
         }
 
-        var queriedModels: [Model] = []
+        var queriedModels: [(ModelName, Model)] = []
         let chunkedArrays = ids.chunked(into: SQLiteStorageEngineAdapter.maxNumberOfPredicates)
         for chunkedArray in chunkedArrays {
             // TODO: Add conveinence to queryPredicate where we have a list of items, to be all or'ed
             var queryPredicates: [QueryPredicateOperation] = []
             for id in chunkedArray {
                 queryPredicates.append(QueryPredicateOperation(field: associatedField.name, operator: .equals(id)))
             }
-            let groupedQueryPredicates =  QueryPredicateGroup(type: .or, predicates: queryPredicates)
+            let groupedQueryPredicates = QueryPredicateGroup(type: .or, predicates: queryPredicates)
 
             let sempahore = DispatchSemaphore(value: 0)
-            storageAdapter.query(untypedModel: modelType, predicate: groupedQueryPredicates) { result in
+            storageAdapter.query(modelSchema: modelSchema, predicate: groupedQueryPredicates) { result in
                 defer {
                     sempahore.signal()
                 }
                 switch result {
                 case .success(let models):
-                    queriedModels.append(contentsOf: models)
-
+                    queriedModels.append(contentsOf: models.map { model in
+                        (modelName, model)
+                    })
                 case .failure(let error):
-                    log.error("Failed to query \(modelType) on mutation event generation: \(error)")
+                    log.error("Failed to query \(modelSchema) on mutation event generation: \(error)")
                 }
             }
             sempahore.wait()
         }
         return queriedModels
     }
 
+    private func getAssociatedModelsMap(associatedModels: [(ModelName, Model)]) -> [ModelName: [Model]] {
+        var associatedModelsMap = [ModelName: [Model]]()
+        for (modelName, model) in associatedModels {
+            associatedModelsMap[modelName, default: []].append(model)
+        }
+        return associatedModelsMap
+    }
+
     private func collapseResults<M: Model>(queryResult: DataStoreResult<[M]>?,
                                            deleteResult: DataStoreResult<[M]>?,
-                                           associatedModels: [Model]) -> DataStoreResult<([M], [Model])> {
+                                           associatedModelsMap: [ModelName: [Model]]) -> DataStoreResult<([M], [ModelName: [Model]])> {
         if let queryResult = queryResult {
             switch queryResult {
             case .success(let models):
                 if let deleteResult = deleteResult {
                     switch deleteResult {
                     case .success:
-                        return .success((models, associatedModels))
+                        return .success((models, associatedModelsMap))
                     case .failure(let error):
                         return .failure(error)
                     }
@@ -174,7 +186,7 @@ extension StorageEngine {
         }
     }
 
-    func collapseMResult<M: Model>(_ result: DataStoreResult<([M], [Model])>) -> DataStoreResult<[M]> {
+    func collapseMResult<M: Model>(_ result: DataStoreResult<([M], [ModelName: [Model]])>) -> DataStoreResult<[M]> {
         switch result {
         case .success(let results):
             return .success(results.0)
@@ -183,12 +195,13 @@ extension StorageEngine {
         }
     }
 
-    func resolveAssociatedModels<M: Model>(_ result: DataStoreResult<([M], [Model])>) -> [Model] {
+    func resolveAssociatedModels<M: Model>(
+        _ result: DataStoreResult<([M], [ModelName: [Model]])>) -> [ModelName: [Model]] {
         switch result {
         case .success(let results):
             return results.1
         case .failure:
-            return []
+            return [:]
         }
     }
 }