fix(amplify): DataStore query fix column missing issue for @connection belongsTo schema

Amplify.xcodeproj/project.pbxproj

@@ -415,6 +415,7 @@
 		B9FAA180238FBB5D009414B4 /* Model+Array.swift in Sources */ = {isa = PBXBuildFile; fileRef = B9FAA17F238FBB5D009414B4 /* Model+Array.swift */; };
 		B9FB05F82383740D00DE1FD4 /* DataStoreStatement.swift in Sources */ = {isa = PBXBuildFile; fileRef = B9FB05F72383740D00DE1FD4 /* DataStoreStatement.swift */; };
 		D83C5160248964780091548E /* ModelGraphQLTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D83C515F248964780091548E /* ModelGraphQLTests.swift */; };
+		D84CF107255C5CA8007B96A9 /* CodingKeysTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = D84CF106255C5CA8007B96A9 /* CodingKeysTests.swift */; };
 		D8DD7A1D24A1CCCD001C49FD /* QuerySortInput.swift in Sources */ = {isa = PBXBuildFile; fileRef = D8DD7A1C24A1CCCD001C49FD /* QuerySortInput.swift */; };
 		FA00F68824DA37EE003E8A71 /* AuthCategoryBehavior+Combine.swift in Sources */ = {isa = PBXBuildFile; fileRef = FA00F68724DA37EE003E8A71 /* AuthCategoryBehavior+Combine.swift */; };
 		FA00F68A24DA3A43003E8A71 /* AuthCategoryDeviceBehavior+Combine.swift in Sources */ = {isa = PBXBuildFile; fileRef = FA00F68924DA3A43003E8A71 /* AuthCategoryDeviceBehavior+Combine.swift */; };
@@ -1265,6 +1266,7 @@
 		D5363CAF9EFAA822FED56808 /* Pods_Amplify_AWSPluginsCore_AWSPluginsTestConfigs_AWSPluginsTestCommon.framework */ = {isa = PBXFileReference; explicitFileType = wrapper.framework; includeInIndex = 0; path = Pods_Amplify_AWSPluginsCore_AWSPluginsTestConfigs_AWSPluginsTestCommon.framework; sourceTree = BUILT_PRODUCTS_DIR; };
 		D5521D5FA66340943C39C451 /* Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSAPICategoryPlugin.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSAPICategoryPlugin.debug.xcconfig"; path = "Target Support Files/Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSAPICategoryPlugin/Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSAPICategoryPlugin.debug.xcconfig"; sourceTree = "<group>"; };
 		D83C515F248964780091548E /* ModelGraphQLTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ModelGraphQLTests.swift; sourceTree = "<group>"; };
+		D84CF106255C5CA8007B96A9 /* CodingKeysTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = CodingKeysTests.swift; sourceTree = "<group>"; };
 		D8DD7A1C24A1CCCD001C49FD /* QuerySortInput.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = QuerySortInput.swift; sourceTree = "<group>"; };
 		DD2486414D63230FF39130C7 /* Pods-Amplify-AWSPluginsCore.debug.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Amplify-AWSPluginsCore.debug.xcconfig"; path = "Target Support Files/Pods-Amplify-AWSPluginsCore/Pods-Amplify-AWSPluginsCore.debug.xcconfig"; sourceTree = "<group>"; };
 		DEEB82A328223C60557B75C1 /* Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSPinpointAnalyticsPlugin.release.xcconfig */ = {isa = PBXFileReference; includeInIndex = 1; lastKnownFileType = text.xcconfig; name = "Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSPinpointAnalyticsPlugin.release.xcconfig"; path = "Target Support Files/Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSPinpointAnalyticsPlugin/Pods-Amplify-AmplifyAWSPlugins-AWSPluginsCore-AWSPinpointAnalyticsPlugin.release.xcconfig"; sourceTree = "<group>"; };
@@ -3396,6 +3398,7 @@
 		FAD3937B23820CE200463F5E /* DataStore */ = {
 			isa = PBXGroup;
 			children = (
+        D84CF106255C5CA8007B96A9 /* CodingKeysTests.swift */,
 				FAD3937923820CDB00463F5E /* DataStoreCategoryClientAPITests.swift */,
 				FAD3937C23820D0200463F5E /* DataStoreCategoryConfigurationTests.swift */,
 				B4944D51251C141200BF0BFE /* JSONValueHolderTest.swift */,
@@ -4818,6 +4821,7 @@
 				B91A87A423D64B0F0049A12F /* TemporalTests.swift in Sources */,
 				FA00F69024DA3F95003E8A71 /* HubCombineTests.swift in Sources */,
 				FA5D76AF23947E9C00489864 /* Model+CodableTests.swift in Sources */,
+				D84CF107255C5CA8007B96A9 /* CodingKeysTests.swift in Sources */,
 				FA58456B24DA31370028D65A /* AmplifyInProcessReportingOperationChainedTests.swift in Sources */,
 				FA1B964E24BF5FA70002B90A /* AmplifyOperationCombineTests.swift in Sources */,
 				FAAFAF3323904BA4002CF932 /* AtomicValue+NumericTests.swift in Sources */,

Amplify/Categories/DataStore/Query/ModelKey.swift

@@ -32,24 +32,43 @@ import Foundation
 ///     post.content != nil
 /// })
 /// ```
-public protocol ModelKey: CodingKey, CaseIterable, QueryFieldOperation {}
+public protocol ModelKey: CodingKey, CaseIterable, QueryFieldOperation, DefaultLogger {
+    var modelName: String { get }
+}
+
+extension ModelKey {
+    public var modelName: String { "" }
+}
 
 extension CodingKey where Self: ModelKey {
 
+    var columnName: String {
+        guard let modelSchema: ModelSchema = ModelRegistry.modelSchema(from: modelName) else {
+            log.warn("Please upgrade to the latest version of Amplify CLI and rerun `amplify codegen models`")
+            return stringValue
+        }
+        switch modelSchema.field(withName: stringValue)?.association {
+        case .belongsTo(_, let targetName):
+            return targetName ?? stringValue
+        default:
+            return stringValue
+        }
+    }
+
     // MARK: - beginsWith
     public func beginsWith(_ value: String) -> QueryPredicateOperation {
-        return field(stringValue).beginsWith(value)
+        return field(columnName).beginsWith(value)
     }
 
     // MARK: - between
     public func between(start: Persistable, end: Persistable) -> QueryPredicateOperation {
-        return field(stringValue).between(start: start, end: end)
+        return field(columnName).between(start: start, end: end)
     }
 
     // MARK: - contains
 
     public func contains(_ value: String) -> QueryPredicateOperation {
-        return field(stringValue).contains(value)
+        return field(columnName).contains(value)
     }
 
     public static func ~= (key: Self, value: String) -> QueryPredicateOperation {
@@ -59,11 +78,11 @@ extension CodingKey where Self: ModelKey {
     // MARK: - eq
 
     public func eq(_ value: Persistable?) -> QueryPredicateOperation {
-        return field(stringValue).eq(value)
+        return field(columnName).eq(value)
     }
 
     public func eq(_ value: EnumPersistable) -> QueryPredicateOperation {
-        return field(stringValue).eq(value)
+        return field(columnName).eq(value)
     }
 
     public static func == (key: Self, value: Persistable?) -> QueryPredicateOperation {
@@ -77,7 +96,7 @@ extension CodingKey where Self: ModelKey {
     // MARK: - ge
 
     public func ge(_ value: Persistable) -> QueryPredicateOperation {
-        return field(stringValue).ge(value)
+        return field(columnName).ge(value)
     }
 
     public static func >= (key: Self, value: Persistable) -> QueryPredicateOperation {
@@ -87,7 +106,7 @@ extension CodingKey where Self: ModelKey {
     // MARK: - gt
 
     public func gt(_ value: Persistable) -> QueryPredicateOperation {
-        return field(stringValue).gt(value)
+        return field(columnName).gt(value)
     }
 
     public static func > (key: Self, value: Persistable) -> QueryPredicateOperation {
@@ -97,7 +116,7 @@ extension CodingKey where Self: ModelKey {
     // MARK: - le
 
     public func le(_ value: Persistable) -> QueryPredicateOperation {
-        return field(stringValue).le(value)
+        return field(columnName).le(value)
     }
 
     public static func <= (key: Self, value: Persistable) -> QueryPredicateOperation {
@@ -107,7 +126,7 @@ extension CodingKey where Self: ModelKey {
     // MARK: - lt
 
     public func lt(_ value: Persistable) -> QueryPredicateOperation {
-        return field(stringValue).lt(value)
+        return field(columnName).lt(value)
     }
 
     public static func < (key: Self, value: Persistable) -> QueryPredicateOperation {
@@ -117,11 +136,11 @@ extension CodingKey where Self: ModelKey {
     // MARK: - ne
 
     public func ne(_ value: Persistable?) -> QueryPredicateOperation {
-        return field(stringValue).ne(value)
+        return field(columnName).ne(value)
     }
 
     public func ne(_ value: EnumPersistable) -> QueryPredicateOperation {
-        return field(stringValue).ne(value)
+        return field(columnName).ne(value)
     }
 
     public static func != (key: Self, value: Persistable?) -> QueryPredicateOperation {

AmplifyPlugins/Core/AWSPluginsCoreTests/Model/Decorator/AuthRuleDecorator/ModelWithOwnerAuthAndGroupWithGroupClaim.swift

@@ -16,9 +16,13 @@ import XCTest
    @model
    @auth(
      rules: [
-       { allow: owner, provider: oidc, identityClaim: "sub"},
-       { allow: groups, provider: oidc, groups: ["Admins"],
-                                    groupClaim: "https://myapp.com/claims/groups"}
+       { allow: owner, provider: oidc, identityClaim: "sub" }
+       {
+         allow: groups
+         provider: oidc
+         groups: ["Admins"]
+         groupClaim: "https://myapp.com/claims/groups"
+       }
      ]
    ) {
    id: ID!
@@ -28,48 +32,48 @@ import XCTest
  */
 
 public struct OIDCGroupPost: Model {
-  public let id: String
-  public var title: String
-  public var owner: String?
-
-  public init(id: String = UUID().uuidString,
-              title: String,
-              owner: String? = nil) {
-      self.id = id
-      self.title = title
-      self.owner = owner
-  }
+    public let id: String
+    public var title: String
+    public var owner: String?
+
+    public init(id: String = UUID().uuidString,
+                title: String,
+                owner: String? = nil) {
+        self.id = id
+        self.title = title
+        self.owner = owner
+    }
 
     // MARK: - CodingKeys
-     public enum CodingKeys: String, ModelKey {
-      case id
-      case title
-      case owner
+    public enum CodingKeys: String, ModelKey {
+        case id
+        case title
+        case owner
     }
 
     public static let keys = CodingKeys.self
     public static let schema = defineSchema { model in
-      let oIDCGroupPost = OIDCGroupPost.keys
-
-      model.authRules = [
-        rule(allow: .owner,
-             ownerField: "owner",
-             identityClaim: "sub",
-             operations: [.create, .update, .delete, .read]),
-        rule(allow: .groups,
-             groupClaim: "https://myapp.com/claims/groups",
-             groups: ["Admins"],
-             operations: [.create, .update, .delete, .read])
-      ]
-
-      model.pluralName = "OIDCGroupPosts"
-
-      model.fields(
-        .id(),
-        .field(oIDCGroupPost.title, is: .required, ofType: .string),
-        .field(oIDCGroupPost.owner, is: .optional, ofType: .string)
-      )
-      }
+        let oIDCGroupPost = OIDCGroupPost.keys
+
+        model.authRules = [
+            rule(allow: .owner,
+                 ownerField: "owner",
+                 identityClaim: "sub",
+                 operations: [.create, .update, .delete, .read]),
+            rule(allow: .groups,
+                 groupClaim: "https://myapp.com/claims/groups",
+                 groups: ["Admins"],
+                 operations: [.create, .update, .delete, .read])
+        ]
+
+        model.pluralName = "OIDCGroupPosts"
+
+        model.fields(
+            .id(),
+            .field(oIDCGroupPost.title, is: .required, ofType: .string),
+            .field(oIDCGroupPost.owner, is: .optional, ofType: .string)
+        )
+    }
 }
 
 class ModelWithOwnerAuthAndGroupWithGroupClaim: XCTestCase {

AmplifyPlugins/Core/AWSPluginsTestCommon/MockAWSAuthService.swift

@@ -16,7 +16,7 @@ public class MockAWSAuthService: AWSAuthServiceBehavior {
     var getTokenClaimsError: AuthError?
     var identityId: String?
     var token: String?
-    var tokenClaims: [String : AnyObject]?
+    var tokenClaims: [String: AnyObject]?
 
     public func configure() {
     }
@@ -45,10 +45,10 @@ public class MockAWSAuthService: AWSAuthServiceBehavior {
         return .success(token ?? "token")
     }
 
-    public func getTokenClaims(tokenString: String) -> Result<[String : AnyObject], AuthError> {
+    public func getTokenClaims(tokenString: String) -> Result<[String: AnyObject], AuthError> {
         if let error = getTokenClaimsError {
             return .failure(error)
         }
-        return .success(tokenClaims ?? ["":"" as AnyObject])
+        return .success(tokenClaims ?? ["": "" as AnyObject])
     }
 }

AmplifyPlugins/DataStore/AWSDataStoreCategoryPluginTests/Core/SQLStatementTests.swift

@@ -330,6 +330,31 @@ class SQLStatementTests: XCTestCase {
         XCTAssertEqual(statement.stringValue, expectedStatement)
     }
 
+    /// - Given: a `Model` type
+    /// - When:
+    ///   - the model is of type `Comment`
+    ///   - `QueryPredicate` on a connected field is given
+    /// - Then:
+    ///   - check if the generated SQL statement is valid
+    ///   - check if an inner join is added referencing `Post`
+    ///   - check if a where with columnName `commentPostId` is generated
+    func testSelectStatementFromModelWhenQueryPredicateIsGiven() {
+        let statement = SelectStatement(from: Comment.schema, predicate: Comment.keys.post == "commentPostId")
+        let expectedStatement = """
+        select
+          "root"."id" as "id", "root"."content" as "content", "root"."createdAt" as "createdAt",
+          "root"."commentPostId" as "commentPostId", "post"."id" as "post.id", "post"."content" as "post.content",
+          "post"."createdAt" as "post.createdAt", "post"."draft" as "post.draft", "post"."rating" as "post.rating",
+          "post"."status" as "post.status", "post"."title" as "post.title", "post"."updatedAt" as "post.updatedAt"
+        from Comment as "root"
+        inner join Post as "post"
+          on "post"."id" = "root"."commentPostId"
+        where 1 = 1
+          and "root"."commentPostId" = ?
+        """
+        XCTAssertEqual(statement.stringValue, expectedStatement)
+    }
+
     // MARK: - Select Statements paginated
 
     /// - Given: a `Model` type and a `QueryPaginationInput`

AmplifyTestCommon/Models/Associations/BookAuthor+Schema.swift

@@ -15,6 +15,10 @@ extension BookAuthor {
         case id
         case book
         case author
+
+        public var modelName: String {
+            return "BookAuthor"
+        }
     }
 
     public static let keys = CodingKeys.self

AmplifyTestCommon/Models/Associations/UserProfile+Schema.swift

@@ -14,6 +14,10 @@ extension UserProfile {
     public enum CodingKeys: String, ModelKey {
         case id
         case account
+
+        public var modelName: String {
+            return "UserProfile"
+        }
     }
 
     public static let keys = CodingKeys.self

AmplifyTestCommon/Models/Comment+Schema.swift

@@ -10,27 +10,31 @@ import Amplify
 import Foundation
 
 extension Comment {
-  // MARK: - CodingKeys
-   public enum CodingKeys: String, ModelKey {
-    case id
-    case content
-    case createdAt
-    case post
-  }
+    // MARK: - CodingKeys
+    public enum CodingKeys: String, ModelKey {
+        case id
+        case content
+        case createdAt
+        case post
 
-  public static let keys = CodingKeys.self
-  //  MARK: - ModelSchema
+        public var modelName: String {
+            return "Comment"
+        }
+    }
+
+    public static let keys = CodingKeys.self
+    //  MARK: - ModelSchema
 
-  public static let schema = defineSchema { model in
-    let comment = Comment.keys
+    public static let schema = defineSchema { model in
+        let comment = Comment.keys
 
-    model.pluralName = "Comments"
+        model.pluralName = "Comments"
 
-    model.fields(
-      .id(),
-      .field(comment.content, is: .required, ofType: .string),
-      .field(comment.createdAt, is: .required, ofType: .dateTime),
-      .belongsTo(comment.post, is: .required, ofType: Post.self, targetName: "commentPostId")
-    )
+        model.fields(
+            .id(),
+            .field(comment.content, is: .required, ofType: .string),
+            .field(comment.createdAt, is: .required, ofType: .dateTime),
+            .belongsTo(comment.post, is: .required, ofType: Post.self, targetName: "commentPostId")
+        )
     }
 }