feat: DataStore and API Lazy Loading with Custom Selection Set

[lawmicha]

Issue #, if available:

Table of Contents

• Custom Selection Set -> feat(api): adds selection set customization #2526
• LazyReference -> feat: DataStore and API Lazy Loading with Custom Selection Set #2583 (You are here)
• Codegen -> feat(amplify-codegen): iOS LazyReference and ModelPath amplify-codegen#504
• CLI feature flag -> chore: feature flag for Lazy Loading and Custom Selection Set support amplify-cli#11566
• Docs -> feat(cli): migration guide for lazy loading and custom selection set feature flag docs#4934

Instructions for installing this PR

• See codegen PR (feat(amplify-codegen): iOS LazyReference and ModelPath amplify-codegen#504) to set up amplify-dev to generate the models used by this PR
• To install this PR in your iOS App, Add Package and use
  • repo: https://github.com/aws-amplify/amplify-swift.git
  • branch: data-dev-preview

Description of changes

This PR introduces the LazyReference type and plugin implementations to extend scenarios in API and DataStore Model decoding. Without this PR and related codegen changes, a belongs-to/has-one association of a model must exist when it is a required association and must be in the data payload for it to be decoded successfully into the Model with that association. For example, a Comment belongs-to a Post, when retrieving the Comment, the data payload being decoded must include the Post because the model type defintion of the Comment includes the Post as a reference:

struct Comment: Model {
  let id: String
  var content: String?
  var post: Post
}

The data payload must include post for it to decode, (shown as json below):

data: {
   id: "commentId",
   content: "content",
   post: {
      postId: "postId123",
      title: "postTitle"
   }
}

This is known as eager loading the associated data. All retrievals of the Comment requires that the post data is retrieved at the same time as well- in DataStore, this is done using a SQL join statement. In API plugin, this is specified in a nested selection set containing the post fields.

This PR introduces the LazyReference type. The post reference in the Comment type is replaced with an instance of LazyReference<Post>.

struct Comment: Model {
  let _post: LazyReference<Post>
  public let post: Post { 
     get async throws {
         try await _post.require()
     }
  }
}

The computed property post is backed by the internal _post of type LazyReference. The computed property allows the developer to make an async call to lazy load the post. Alternative solutions to having the computed property backed by the internal property is to have a property wrapper instead. The main disadvantage of the property wrapper is the inability to perform an async operations, like how we are achieving it with the async throws getter of the computed property.

With the computed property and internal property, the developer call pattern becomes:

let post = try await comment.post

When the association is required, the underlying API used is _post.require(), if the association is optional, then _post.get(). _post.require() will throw if the model cannot be loaded, while get() will return nil. The LazyReference type allows the data payloads to include all fields (eager loaded) or include only the primary keys of post (lazy loaded), used for decoding to the type into one of the two states.

When a LazyReference instance is created during data decoding, a model provider is created using the registered plugin's model provider registry.

When using API plugin:

flowchart TB
  A["Amplify.addPlugin(plugin)"] --> B["Amplify.configure()"]
  B --> C["APIPlugin.configure()"]
  C --> D["ModelProviderRegistry.registry(AppSyncModelDecoder.self)"]
  E["Amplify.API.query(.get(Comment.self, byId: 'commentId'))"] -- response payload --> F["LazyReference decoder"]
  F --> G["ModelProviderRegistry.decode()"]
  G --> I["LazyReference's modelProvider = AppSyncModelProvider()"]

Loading

When using DataStore plugin:

flowchart TB
  A["Amplify.addPlugin(plugin)"] --> B["Amplify.configure()"]
  B --> C["DataStorePlugin.configure()"]
  C --> D["ModelProviderRegistry.registry(DataStoreModelDecoder.self)"]
  E["Amplify.DataStore.query(Comment.self, byId: 'commentId')"] -- response payload --> F["LazyReference decoder"]
  F --> G["ModelProviderRegistry.decode()"]
  G --> I["LazyReference's modelProvider = DataStoreModelProvider()"]

Loading

The LazyReference instance delegates its lazy loading capabilities to the underlying plugin's model provider, which performs requests to its data source (either SQLite or AppSync directly)

classDiagram
    class LazyReference
    LazyReference : +ModelProvider modelProvider
    LazyReference : +LazyReferenceValueState _state
    LazyReference : +get()
    LazyReference : +required()

Loading

For example, API plugin will make a GraphQL get request using the model schema and identifiers of the model to retrieve the entire model and transition the LazyReference to a loaded state.

flowchart TB
  A["let post = try await comment.post"] --> B["try await _post.require()"]
  B --> C["try await modelProvider.load()"]
  C --> D["try await Amplify.API.query(.get(Post.self, byId: 'postId'))"]

Loading

The selection set created above includes only the primary keys of the post when using the latest codegen that includes the rootPath, to control the selection set customizations from @drochetti's PR #2526. This is a pre-requiste PR to allow selection set customizations to create "not loaded" LazyReference instances.

DataStore will perform a SQL select operation.

flowchart TB
  A["let post = try await comment.post"] --> B["try await _post.require()"]
  B --> C["try await modelProvider.load()"]
  C --> D["try await Amplify.DataStore.query(Post.self, byId: 'postId')"]

Loading

The plugin will manipulate the data payload to always add the model's identifier as metadata to the last level of the model graph. When dealing with has-many associations, it will add the target's associated identifier (comments hold the post identifier as metadata). When dealing with the belongs-to, it will add ithe target's identifier (post will hold the post's identifier). The added metadata gets decoded to a LazyReference and List types to hold the metadata in its model provider. This allows the developer to tranverse the model graph indefinitely using the existing lazy List implementation and the new LazyReference implementation, regardless of the initial loaded state of the model.

Blog-Post-Comment example

flowchart TB
  A["let post = try await comment.post"] --> B["let blog = try await post.blog"]
  B -->  C["try await blog.posts?.fetch()"] 
  C --> D["let posts = blog.posts"]
  D --> E["try await post.comments?.fetch()"] 
  E --> F["let comments = post.comments"]

Loading

AppSync LazyReference metadata

AppSyncModelProvider's idenifier data is an ordered dictionary of its identifiers, the ordering matters to specify the ID! type and String! types when making the AppSync get request. For default identifier, this is ["id":"idValue"]. For a composite key, with a renamed primary key for example, will be ["postId":"postIdValue", "title":"titleValue"]

DataStore LazyReference metadata

DataStoreModelProvider's identifier is the parent object's foreign key value. The column name is either the single column name like "id" or "postId" or for composite keys, it is called @@primaryKey.

Codegen Changes

PR with related codegen changes aws-amplify/amplify-codegen#504

Selection Set changes

• DataStore establishes subscriptions with a reduced selection set, the nested selection set only contains primary keys.
• The codegen models have to be the ones with the LazyReference type since it can decode the reduced selection set
• DataStore's mutations do not change, since older iOS clients or iOS clients using the previously codegenerated models will continue to decode the response payload successfully.
• (iOS DataStore imeplementation detail) The __typename is injected at runtime from the request to the response that allows response payloads missing the __typename to be decoded successfully.

Developer Experience / Migration Guide

aws-amplify/docs#4934

Issues

Use case: Has-One Bidirectional belongs to

See #504 (comment)

See #505 (comment)

Use case: Cross platform DataStore events (Android)

DataStore on Android syncs data to AppSync using a selection set in the mutation to AppSync differently from iOS. For the existing issue, see x. The problem is that there is not enough data in the subscription event, due to Android's selection set missing fields that are required for iOS DataStore. For example, in the Comment Post example, Android may sync a comment with the selection set

comment {
  id
  post {
    id
  }
}

This will fail to decode to iOS without the LazyReference codegen changes because the swift model type Post that is a field on the Comment may have required fields. With the LazyReference encapsulation, the post field on the Comment will be decoded as a "not loaded" Post with the identifiers. In DataStore's reconciliation logic, it only needs to extract the identifiers of the Post from the comment and store that as the FK values in the Comment table. It doesn't need the entire Post object.

See #1753 (comment)

Use case: Cross platform DataStore events (Studio/JS)

Studio/JS sends a mutation events with a selection set of this structure:

mutation MyMutation2 {
  createBookAuthor(input: {authorId: "b87ebfd4-04e7-45aa-b54a-6916928bb5c8", bookId: "532078b3-376a-451d-a470-608cc7e462a5"}) {
    _deleted
    _lastChangedAt
    _version
    authorId
    bookId
    createdAt
    id
    updatedAt
    author {
      id
      _deleted
    }
    book {
      id
      _deleted
    }
  }
}

• no __typename
• version metadata fields at the top level, _deleted, _lastChangedAt, _version
• nested models only containing its primary keys, in this case id
• _deleted

These events are successfully reconciled by the iOS client using this implementation because __typename is injected at runtime by DataStore to decode successfully to the AnyModel type used by DataStore. The author and book in this example is codegenerated to be encapsulated by the LazyReference type and decodes to the "not loaded" states. Their identifiers are retrieved as the FK for saving the bookAuthor join model

See #1601 (comment)

Check points: (check or cross out if not relevant)

• Added new tests to cover change, if needed
• Build succeeds with all target using Swift Package Manager
• All unit tests pass
• All integration tests pass
• Security oriented best practices and standards are followed (e.g. using input sanitization, principle of least privilege, etc)
• Documentation update for the change if required
• PR title conforms to conventional commit style
• If breaking change, documentation/changelog update with migration instructions

DataStore checkpoints (check when completed)

• Ran AWSDataStorePluginIntegrationTests
• Ran AWSDataStorePluginV2Tests
• Ran AWSDataStorePluginMultiAuthTests
• Ran AWSDataStorePluginCPKTests
• Ran AWSDataStorePluginAuthCognitoTests
• Ran AWSDataStorePluginAuthIAMTests

By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.

[drochetti]

Some initial comments, I'm still going through it so... more to come!

[drochetti]

Is this custom key/value pair required? Is [String: String] not compatible with Codable?

[lawmicha]

I don't believe the swift dictionary of String is codable, but I can check again to see if there are alternatives to this implementation

[atierian]

[String: String] does indeed conform to Codable. Are we doing key based lookups on an [LazyReferenceIdentifier] anywhere? Do we need to support (or explicitly disallow) duplicate keys?

[lawmicha]

ah, the reason was the need for an ordered dictionary, that's why it's array of codable structs right now. We are not doing key based lookups anywhere since

• datastore model provider extracts a single item from the array and uses that as the identifier
• api model provider requires the ordered list of identifiers, it translates it to the GraphQL request in ModelIdDecorator

[atierian]

[String: String] does indeed conform to Codable. Are we doing key based lookups on an [LazyReferenceIdentifier] anywhere? Do we need to support (or explicitly disallow) duplicate keys?

[lawmicha]

Error Handling

Lazy List APIs were launched in Amplify v1 with a return type that is shared across DataStore and API users. Because, they pivot on behavior based on the underlying plugin being used, error scenarios are the plugin's list/model provider's responsiblity to return the correct error case. However, in this architecture, there's also a default list/model provider which is used when neither plugin specific providers can be instantaited. The explicit return type on the APIs off the List type cannot be either DataStoreError or APIError. The default providers also should not return either DataStoreError or APIError since it is a plugin agnostic core type. Because of these reasons we introduced the CoreError, with the API signature like this:

public func fetch(_ completion: @escaping (Result<Void, CoreError>) -> Void)

With the launch of Amplify v2, the API signatures have been updated to the Async APIs here:

public func fetch() async throws

The Async API removed the requirement to explicitly define the CoreError type on the API, except that we still need to return something in the implementation of the core types, such as in the default providers or the code within the Lazy types that manage creating plugin specific providers. This leaves us a path forward in this PR for the minor version bump:

1. In this PR, the new plugin providers will return their respective errors, APIError and DataStoreError. The existing list providers will continue to return CoreError.listOperation
2. CoreError is still needed for the default providers and some other places in the implementation of Amplify core types. We shouldn't remove CoreError, nor can we extend CoreError to new cases that we may want to communicate. For now, we'll return clientValidation error until this changes.

For a major version bump, we can change the behavior of the plugin providers

1. list providers to return APIError or DataStore instead of CoreError. Extend APIError/DataStore (breaking change) to what is required to accomplish this.
2. CoreError can be extended to improve the cases it can communicate. There are two cases are
3. dataUnavailable indicates some data marked as required was unavailable
4. dataIntegrity indicates expected data could not be loaded due to the lack of integrity.

Ref Issue: #2654

[atierian]

formatting nit: leading whitespace

[lawmicha]

Thanks for reviewing again, as this PR is getting pretty large, i'll address this in the next PR with a couple of other things as well

[lostdawg3]

I'm not sure if I can respond on this thread but is there any update to this in regards to the main. Currently this branch is the only one that I can get @auth owner to work on both physical device and ios simulator (M1). When trying to use amplify 2.0.0, updates do not work.

[lostdawg3]

Sorry it worked for a few times then stopped. It does work when I put @auth private for both the physical device and ios simulator (M1).

[lawmicha]

Hi @lostdawg3, thanks for trying out DataStore! Feel free to provide us more information in a GitHub Issue . Your schema will be helpful and steps you took to reproduce the issue. That way we'll have a dedicated thread to dig deeper into issues you are experiencing. If there's something that's working on this branch but not on the current main (v2), we can pinpoint the changes necessary and apply them outside of this PR to get you unblocked faster.