Skip to content

Commit ac23668

Browse files
5dlawmicha
5d
authored and
lawmicha
committed
fix(datastore): load hasOne and belongsTo lazy reference with composite key (#2737)
* fix(datastore): load has-one reference with composite key * fix: update descretKeys split logic * fix: add associatedFields for hashOne and belongsTo * fix: change json object seralization
1 parent bfd7882 commit ac23668

29 files changed

+251
-103
lines changed

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

Lines changed: 11 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -97,7 +97,7 @@ public enum ModelAssociation {
9797
let targetNames = targetName.map { [$0] } ?? []
9898
return .belongsTo(associatedFieldName: nil, targetNames: targetNames)
9999
}
100-
100+
101101
public static func hasMany(associatedWith: CodingKey? = nil, associatedFields: [CodingKey] = []) -> ModelAssociation {
102102
return .hasMany(associatedFieldName: associatedWith?.stringValue, associatedFieldNames: associatedFields.map { $0.stringValue })
103103
}
@@ -225,7 +225,7 @@ extension ModelField {
225225
}
226226
return true
227227
}
228-
228+
229229
/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
230230
/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
231231
/// application making any change to these `public` types should be backward compatible, otherwise it will be a
@@ -267,19 +267,23 @@ extension ModelField {
267267
/// breaking change.
268268
public var associatedFieldNames: [String] {
269269
switch association {
270-
case .belongsTo(let associatedKey, let associatedKeys),
271-
.hasOne(let associatedKey, let associatedKeys),
272-
.hasMany(let associatedKey, let associatedKeys):
270+
case .hasMany(let associatedKey, let associatedKeys):
273271
if associatedKeys.isEmpty, let associatedKey = associatedKey {
274272
return [associatedKey]
275273
}
276-
277274
return associatedKeys
275+
276+
case .hasOne, .belongsTo:
277+
return ModelRegistry.modelSchema(from: requiredAssociatedModelName)?
278+
.primaryKey
279+
.fields
280+
.map(\.name) ?? []
281+
278282
case .none:
279283
return []
280284
}
281285
}
282-
286+
283287
/// - Warning: Although this has `public` access, it is intended for internal & codegen use and should not be used
284288
/// directly by host applications. The behavior of this may change without warning. Though it is not used by host
285289
/// application making any change to these `public` types should be backward compatible, otherwise it will be a

‎Amplify/Categories/DataStore/Model/Internal/Schema/ModelSchema+Definition.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -240,7 +240,7 @@ public enum ModelFieldDefinition {
240240
ofType: .collection(of: type),
241241
association: .hasMany(associatedWith: associatedKey))
242242
}
243-
243+
244244
public static func hasMany(_ key: CodingKey,
245245
is nullability: ModelFieldNullability = .required,
246246
isReadOnly: Bool = false,
@@ -250,7 +250,7 @@ public enum ModelFieldDefinition {
250250
is: nullability,
251251
isReadOnly: isReadOnly,
252252
ofType: .collection(of: type),
253-
association: .hasMany(associatedWith: associatedKeys.first ?? nil, associatedFields: associatedKeys))
253+
association: .hasMany(associatedWith: associatedKeys.first, associatedFields: associatedKeys))
254254
}
255255

256256
public static func hasOne(_ key: CodingKey,

‎AmplifyPlugins/Core/AWSPluginsCore/Model/Support/Model+GraphQL.swift

Lines changed: 13 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -167,23 +167,18 @@ extension Model {
167167

168168
let fieldNames = getFieldNameForAssociatedModels(modelField: field)
169169
var values = getModelIdentifierValues(from: value, modelSchema: associateModelSchema)
170+
if values.count != fieldNames.count {
171+
values = [Persistable?](repeating: nil, count: fieldNames.count)
172+
}
170173

171-
if fieldNames.count != values.count {
172-
// if the field is required, the associated field keys and values should match
173-
if field.isRequired {
174-
preconditionFailure(
175-
"""
176-
Associated model target names and values for field \(field.name) of model \(modelName) mismatch.
177-
There is a possibility that is an issue with the generated models.
178-
"""
179-
)
180-
} else if mutationType == .update {
181-
// otherwise, pad the values with `nil` to account for removals of associations on updates.
182-
values = [Persistable?](repeating: nil, count: fieldNames.count)
174+
let associatedModelIdentifiers = zip(fieldNames, values).map { ($0.0, $0.1)}
175+
if mutationType != .update {
176+
return associatedModelIdentifiers.compactMap { key, value in
177+
value.map { (key, $0) }
183178
}
179+
} else {
180+
return associatedModelIdentifiers
184181
}
185-
186-
return Array(zip(fieldNames, values))
187182
}
188183

189184
/// Given a model and its schema, returns the values of its identifier (primary key).
@@ -192,7 +187,10 @@ extension Model {
192187
/// - value: model value
193188
/// - modelSchema: model's schema
194189
/// - Returns: array of values of its primary key
195-
private func getModelIdentifierValues(from value: Any, modelSchema: ModelSchema) -> [Persistable?] {
190+
private func getModelIdentifierValues(
191+
from value: Any,
192+
modelSchema: ModelSchema
193+
) -> [Persistable?] {
196194
if let modelValue = value as? Model {
197195
return modelValue.identifier(schema: modelSchema).values
198196
} else if let optionalModel = value as? Model?,

‎AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/AWSDataStorePlugin+DataStoreBaseBehavior.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ extension AWSDataStorePlugin: DataStoreBaseBehavior {
6868
storageEngine.save(model,
6969
modelSchema: modelSchema,
7070
condition: condition,
71-
eagerLoad: true,
71+
eagerLoad: isEagerLoad,
7272
completion: publishingCompletion)
7373
}
7474

‎AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Core/DataStoreModelDecoder.swift

Lines changed: 12 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -15,18 +15,25 @@ public struct DataStoreModelDecoder: ModelProviderDecoder {
1515

1616
/// Metadata that contains the foreign key value of a parent model, which is the primary key of the model to be loaded.
1717
struct Metadata: Codable {
18-
let identifier: String?
18+
let identifiers: [LazyReferenceIdentifier]
1919
let source: String
2020

21-
init(identifier: String?, source: String = DataStoreSource) {
22-
self.identifier = identifier
21+
init(identifiers: [LazyReferenceIdentifier], source: String = DataStoreSource) {
22+
self.identifiers = identifiers
2323
self.source = source
2424
}
25+
26+
func toJsonObject() -> Any? {
27+
try? JSONSerialization.jsonObject(with: JSONEncoder().encode(self))
28+
}
2529
}
2630

2731
/// Create a SQLite payload that is capable of initializting a LazyReference, by decoding to `DataStoreModelDecoder.Metadata`.
28-
static func lazyInit(identifier: Binding?) -> [String: Binding?] {
29-
return ["identifier": identifier, "source": DataStoreSource]
32+
static func lazyInit(identifiers: [LazyReferenceIdentifier]) -> Metadata? {
33+
if identifiers.isEmpty {
34+
return nil
35+
}
36+
return Metadata(identifiers: identifiers)
3037
}
3138

3239
public static func decode<ModelType: Model>(modelType: ModelType.Type, decoder: Decoder) -> AnyModelProvider<ModelType>? {

‎AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Core/DataStoreModelProvider.swift

Lines changed: 11 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -14,11 +14,7 @@ public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
1414

1515
// Create a "not loaded" model provider with the identifier metadata, useful for hydrating the model
1616
init(metadata: DataStoreModelDecoder.Metadata) {
17-
if let identifier = metadata.identifier {
18-
self.loadedState = .notLoaded(identifiers: [.init(name: ModelType.schema.primaryKey.sqlName, value: identifier)])
19-
} else {
20-
self.loadedState = .notLoaded(identifiers: nil)
21-
}
17+
self.loadedState = .notLoaded(identifiers: metadata.identifiers)
2218
}
2319

2420
// Create a "loaded" model provider with the model instance
@@ -31,10 +27,15 @@ public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
3127
public func load() async throws -> ModelType? {
3228
switch loadedState {
3329
case .notLoaded(let identifiers):
34-
guard let identifiers = identifiers, let identifier = identifiers.first else {
30+
guard let identifiers = identifiers, !identifiers.isEmpty else {
3531
return nil
3632
}
37-
let queryPredicate: QueryPredicate = field(identifier.name).eq(identifier.value)
33+
34+
let identifierValue = identifiers.count == 1
35+
? identifiers.first?.value
36+
: identifiers.map({ "\"\($0.value)\""}).joined(separator: ModelIdentifierFormat.Custom.separator)
37+
38+
let queryPredicate: QueryPredicate = field(ModelType.schema.primaryKey.sqlName).eq(identifierValue)
3839
let models = try await Amplify.DataStore.query(ModelType.self, where: queryPredicate)
3940
guard let model = models.first else {
4041
return nil
@@ -53,11 +54,9 @@ public class DataStoreModelProvider<ModelType: Model>: ModelProvider {
5354
public func encode(to encoder: Encoder) throws {
5455
switch loadedState {
5556
case .notLoaded(let identifiers):
56-
if let identifier = identifiers?.first {
57-
let metadata = DataStoreModelDecoder.Metadata(identifier: identifier.value)
58-
var container = encoder.singleValueContainer()
59-
try container.encode(metadata)
60-
}
57+
let metadata = DataStoreModelDecoder.Metadata(identifiers: identifiers ?? [])
58+
var container = encoder.singleValueContainer()
59+
try container.encode(metadata)
6160

6261
case .loaded(let element):
6362
try element.encode(to: encoder)

‎AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/CascadeDeleteOperation.swift

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -142,9 +142,9 @@ public class CascadeDeleteOperation<M: Model>: AsynchronousOperation {
142142
}
143143

144144
let modelIds = queriedModels.map { $0.identifier(schema: self.modelSchema).stringValue }
145-
145+
146146
associatedModels = await self.recurseQueryAssociatedModels(modelSchema: self.modelSchema, ids: modelIds)
147-
147+
148148
deletedResult = await withCheckedContinuation { continuation in
149149
self.storageAdapter.delete(self.modelType,
150150
modelSchema: self.modelSchema,

‎AmplifyPlugins/DataStore/Sources/AWSDataStorePlugin/Storage/SQLite/Statement+Model.swift

Lines changed: 34 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ extension Statement: StatementModelConvertible {
5353
Amplify.Logging.logger(forCategory: .dataStore)
5454
}
5555

56-
56+
5757
func convert<M: Model>(to modelType: M.Type,
5858
withSchema modelSchema: ModelSchema,
5959
using statement: SelectStatement,
@@ -66,7 +66,7 @@ extension Statement: StatementModelConvertible {
6666
let result: StatementResult<M> = try StatementResult.from(dictionary: values)
6767
return result.elements
6868
}
69-
69+
7070
func convertToModelValues<M: Model>(to modelType: M.Type,
7171
withSchema modelSchema: ModelSchema,
7272
using statement: SelectStatement,
@@ -121,13 +121,24 @@ extension Statement: StatementModelConvertible {
121121
return convertCollection(field: field, schema: schema, from: element, path: path)
122122

123123
case (.model, false):
124-
let foreignKey = field.association.map(getTargetNames).map {
125-
field.foreignKeySqlName(withAssociationTargets: $0)
126-
}
127-
let targetBinding = foreignKey.flatMap {
124+
let targetNames = getTargetNames(field: field)
125+
var associatedFieldValues = targetNames.map {
128126
getValue(from: element, by: path + [$0])
127+
}.compactMap { $0 }
128+
.map { String(describing: $0) }
129+
130+
if associatedFieldValues.isEmpty {
131+
associatedFieldValues = associatedValues(
132+
from: path + [field.foreignKeySqlName(withAssociationTargets: targetNames)],
133+
element: element
134+
)
135+
}
136+
guard associatedFieldValues.count == field.associatedFieldNames.count else {
137+
return nil
129138
}
130-
return DataStoreModelDecoder.lazyInit(identifier: targetBinding)
139+
return DataStoreModelDecoder.lazyInit(identifiers: zip(field.associatedFieldNames, associatedFieldValues).map {
140+
LazyReferenceIdentifier(name: $0.0, value: $0.1)
141+
})?.toJsonObject()
131142

132143
case let (.model(modelName), true):
133144
guard let modelSchema = getModelSchema(for: modelName, with: statement)
@@ -152,6 +163,14 @@ extension Statement: StatementModelConvertible {
152163
return statement.metadata.columnMapping.values.first { $0.0.name == modelName }.map { $0.0 }
153164
}
154165

166+
private func associatedValues(from foreignKeyPath: [String], element: Element) -> [String] {
167+
return [getValue(from: element, by: foreignKeyPath)]
168+
.compactMap({ $0 })
169+
.map({ String(describing: $0) })
170+
.flatMap({ $0.split(separator: ModelIdentifierFormat.Custom.separator.first!) })
171+
.map({ String($0).trimmingCharacters(in: .init(charactersIn: "\"")) })
172+
}
173+
155174
private func convertCollection(field: ModelField, schema: ModelSchema, from element: Element, path: [String]) -> Any? {
156175
if field.isArray && field.hasAssociation,
157176
case let .some(.hasMany(associatedFieldName: associatedFieldName, associatedFieldNames: associatedFieldNames)) = field.association
@@ -178,19 +197,19 @@ extension Statement: StatementModelConvertible {
178197
return DataStoreListDecoder.lazyInit(associatedIds: primaryKeyValues,
179198
associatedWith: associatedFieldNames)
180199
}
181-
200+
182201
}
183202

184203
return nil
185204
}
186205

187-
private func getTargetNames(assoication: ModelAssociation) -> [String] {
188-
switch assoication {
189-
case let .hasOne(associatedFieldName: _, targetNames: targets):
190-
return targets
191-
case let.belongsTo(associatedFieldName: _, targetNames: targets):
192-
return targets
193-
case .hasMany:
206+
private func getTargetNames(field: ModelField) -> [String] {
207+
switch field.association {
208+
case let .some(.hasOne(_, targetNames)):
209+
return targetNames
210+
case let .some(.belongsTo(_, targetNames)):
211+
return targetNames
212+
default:
194213
return []
195214
}
196215
}

‎AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginLazyLoadTests/LL1/AWSDataStoreLazyLoadPostComment4V2Tests.swift

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ class AWSDataStoreLazyLoadPostComment4V2Tests: AWSDataStoreLazyLoadBaseTest {
4242
let comment = Comment(content: "content", post: post)
4343
let savedPost = try await saveAndWaitForSync(post)
4444
let savedComment = try await saveAndWaitForSync(comment)
45-
try await assertComment(savedComment, hasEagerLoaded: savedPost)
45+
try await assertComment(savedComment, canLazyLoad: savedPost)
4646
try await assertPost(savedPost, canLazyLoad: savedComment)
4747
let queriedComment = try await query(for: savedComment)
4848
try await assertComment(queriedComment, canLazyLoad: savedPost)
@@ -62,7 +62,7 @@ class AWSDataStoreLazyLoadPostComment4V2Tests: AWSDataStoreLazyLoadBaseTest {
6262
XCTFail("Could not encode comment")
6363
return
6464
}
65-
try await assertComment(savedComment, hasEagerLoaded: savedPost)
65+
try await assertComment(savedComment, canLazyLoad: savedPost)
6666

6767
guard let decodedComment = try? ModelRegistry.decode(modelName: Comment.modelName,
6868
from: encodedComment) as? Comment else {
@@ -71,7 +71,7 @@ class AWSDataStoreLazyLoadPostComment4V2Tests: AWSDataStoreLazyLoadBaseTest {
7171
return
7272
}
7373

74-
try await assertComment(decodedComment, hasEagerLoaded: savedPost)
74+
try await assertComment(decodedComment, canLazyLoad: savedPost)
7575
}
7676

7777
func testLazyLoadOnQueryAfterEncodeDecoder() async throws {

‎AmplifyPlugins/DataStore/Tests/DataStoreHostApp/AWSDataStorePluginLazyLoadTests/LL10/AWSDataStoreLazyLoadPostComment7Tests.swift

Lines changed: 21 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
3636
let comment = Comment(commentId: UUID().uuidString, content: "content", post: post)
3737
let savedPost = try await saveAndWaitForSync(post)
3838
let savedComment = try await saveAndWaitForSync(comment)
39-
try await assertComment(savedComment, hasEagerLoaded: savedPost)
39+
try await assertComment(savedComment, canLazyLoad: savedPost)
4040
try await assertPost(savedPost, canLazyLoad: savedComment)
4141
let queriedComment = try await query(for: savedComment)
4242
try await assertComment(queriedComment, canLazyLoad: savedPost)
@@ -68,7 +68,10 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
6868
func assertComment(_ comment: Comment,
6969
canLazyLoad post: Post) async throws {
7070
assertLazyReference(comment._post,
71-
state: .notLoaded(identifiers: [.init(name: "@@primaryKey", value: post.identifier)]))
71+
state: .notLoaded(identifiers: [
72+
.init(name: Post7.keys.postId.stringValue, value: post.postId),
73+
.init(name: Post7.keys.title.stringValue, value: post.title)
74+
]))
7275
guard let loadedPost = try await comment.post else {
7376
XCTFail("Failed to load the post from the comment")
7477
return
@@ -94,7 +97,10 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
9497
return
9598
}
9699
assertLazyReference(comment._post,
97-
state: .notLoaded(identifiers: [.init(name: "@@primaryKey", value: post.identifier)]))
100+
state: .notLoaded(identifiers: [
101+
.init(name: Post7.keys.postId.stringValue, value: post.postId),
102+
.init(name: Post7.keys.title.stringValue, value: post.title)
103+
]))
98104
}
99105

100106
func testSaveWithoutPost() async throws {
@@ -120,7 +126,10 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
120126
let savedComment = try await saveAndWaitForSync(comment)
121127
let queriedComment = try await query(for: savedComment)
122128
assertLazyReference(queriedComment._post,
123-
state: .notLoaded(identifiers: [.init(name: "@@primaryKey", value: post.identifier)]))
129+
state: .notLoaded(identifiers: [
130+
.init(name: Post.keys.postId.stringValue, value: post.postId),
131+
.init(name: Post.keys.title.stringValue, value: post.title)
132+
]))
124133
let savedQueriedComment = try await saveAndWaitForSync(queriedComment, assertVersion: 2)
125134
let queriedComment2 = try await query(for: savedQueriedComment)
126135
try await assertComment(queriedComment2, canLazyLoad: savedPost)
@@ -135,7 +144,10 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
135144
let savedComment = try await saveAndWaitForSync(comment)
136145
var queriedComment = try await query(for: savedComment)
137146
assertLazyReference(queriedComment._post,
138-
state: .notLoaded(identifiers: [.init(name: "@@primaryKey", value: post.identifier)]))
147+
state: .notLoaded(identifiers: [
148+
.init(name: Post.keys.postId.stringValue, value: post.postId),
149+
.init(name: Post.keys.title.stringValue, value: post.title)
150+
]))
139151

140152
let newPost = Post(postId: UUID().uuidString, title: "title")
141153
_ = try await saveAndWaitForSync(newPost)
@@ -154,7 +166,10 @@ final class AWSDataStoreLazyLoadPostComment7Tests: AWSDataStoreLazyLoadBaseTest
154166
let savedComment = try await saveAndWaitForSync(comment)
155167
var queriedComment = try await query(for: savedComment)
156168
assertLazyReference(queriedComment._post,
157-
state: .notLoaded(identifiers: [.init(name: "@@primaryKey", value: post.identifier)]))
169+
state: .notLoaded(identifiers: [
170+
.init(name: Post.keys.postId.stringValue, value: post.postId),
171+
.init(name: Post.keys.title.stringValue, value: post.title)
172+
]))
158173

159174
queriedComment.setPost(nil)
160175
let saveCommentRemovePost = try await saveAndWaitForSync(queriedComment, assertVersion: 2)

0 commit comments

Comments
 (0)