chore(@aws-amplify/datastore): docstrings, small refactors, cleanup unused predicate function (#11040)

* added a few targeted TODO's from PR; in general, look for ops to cleanup as anys and bangs

* datastore.ts refactors and docstrings

* docstrings, small refactor, remove unused method

Author: svidgen

packages/datastore/src/datastore/datastore.ts

@@ -139,9 +139,13 @@ const modelNamespaceMap = new WeakMap<
 	PersistentModelConstructor<any>,
 	string
 >();
-// stores data for crafting the correct update mutation input for a model
-// Patch[] - array of changed fields and metadata
-// PersistentModel - the source model, used for diffing object-type fields
+
+/**
+ * Stores data for crafting the correct update mutation input for a model.
+ *
+ * - `Patch[]` - array of changed fields and metadata.
+ * - `PersistentModel` - the source model, used for diffing object-type fields.
+ */
 const modelPatchesMap = new WeakMap<
 	PersistentModel,
 	[Patch[], PersistentModel]
@@ -158,17 +162,6 @@ const getModelDefinition = (
 	return definition;
 };
 
-const getModelPKFieldName = (
-	modelConstructor: PersistentModelConstructor<any>
-) => {
-	const namespace = modelNamespaceMap.get(modelConstructor);
-	return (
-		(namespace &&
-			schema.namespaces?.[namespace]?.keys?.[modelConstructor.name]
-				.primaryKey) || ['id']
-	);
-};
-
 /**
  * Determine what the managed timestamp field names are for the given model definition
  * and return the mapping.
@@ -202,14 +195,16 @@ const getTimestampFields = (
 	};
 };
 
+/**
+ * Determines whether the given object is a Model Constructor that DataStore can
+ * safely use to construct objects and discover related metadata.
+ *
+ * @param obj The object to test.
+ */
 const isValidModelConstructor = <T extends PersistentModel>(
 	obj: any
 ): obj is PersistentModelConstructor<T> => {
-	if (isModelConstructor(obj) && modelNamespaceMap.has(obj)) {
-		return true;
-	} else {
-		return false;
-	}
+	return isModelConstructor(obj) && modelNamespaceMap.has(obj);
 };
 
 const namespaceResolver: NamespaceResolver = modelConstructor => {
@@ -222,6 +217,25 @@ const namespaceResolver: NamespaceResolver = modelConstructor => {
 	return resolver;
 };
 
+/**
+ * Creates a predicate without any conditions that can be passed to customer
+ * code to have conditions added to it.
+ *
+ * For example, in this query:
+ *
+ * ```ts
+ * await DataStore.query(
+ * 	Model,
+ * 	item => item.field.eq('value')
+ * );
+ * ```
+ *
+ * `buildSeedPredicate(Model)` is used to create `item`, which is passed to the
+ * predicate function, which in turn uses that "seed" predicate (`item`) to build
+ * a predicate tree.
+ *
+ * @param modelConstructor The model the predicate will query.
+ */
 const buildSeedPredicate = <T extends PersistentModel>(
 	modelConstructor: PersistentModelConstructor<T>
 ) => {
@@ -232,9 +246,7 @@ const buildSeedPredicate = <T extends PersistentModel>(
 	);
 	if (!modelSchema) throw new Error('Missing modelSchema');
 
-	const pks = getModelPKFieldName(
-		modelConstructor as PersistentModelConstructor<T>
-	);
+	const pks = extractPrimaryKeyFieldNames(modelSchema);
 	if (!pks) throw new Error('Could not determine PK');
 
 	return recursivePredicateFor<T>({
@@ -387,27 +399,23 @@ const initSchema = (userSchema: Schema) => {
 
 			modelAssociations.set(model.name, connectedModels);
 
+			// Precompute model info (such as pk fields) so that downstream schema consumers
+			// (such as predicate builders) don't have to reach back into "DataStore" space
+			// to go looking for it.
 			Object.values(model.fields).forEach(field => {
-				if (
-					typeof field.type === 'object' &&
-					!Object.getOwnPropertyDescriptor(
-						<ModelFieldType>field.type,
-						'modelConstructor'
-					)
-				) {
+				const relatedModel = userClasses[(<ModelFieldType>field.type).model];
+				if (isModelConstructor(relatedModel)) {
 					Object.defineProperty(field.type, 'modelConstructor', {
 						get: () => {
+							const relatedModelDefinition = getModelDefinition(relatedModel);
+							if (!relatedModelDefinition)
+								throw new Error(
+									`Could not find model definition for ${relatedModel.name}`
+								);
 							return {
-								builder: userClasses[(<ModelFieldType>field.type).model],
-								schema:
-									schema.namespaces[namespace].models[
-										(<ModelFieldType>field.type).model
-									],
-								pkField: getModelPKFieldName(
-									userClasses[
-										(<ModelFieldType>field.type).model
-									] as PersistentModelConstructor<any>
-								),
+								builder: relatedModel,
+								schema: relatedModelDefinition,
+								pkField: extractPrimaryKeyFieldNames(relatedModelDefinition),
 							};
 						},
 					});
@@ -548,7 +556,7 @@ const createTypeClasses: (
 
 	Object.entries(namespace.nonModels || {}).forEach(
 		([typeName, typeDefinition]) => {
-			const clazz = createNonModelClass(typeDefinition) as any;
+			const clazz = createNonModelClass(typeDefinition);
 			classes[typeName] = clazz;
 		}
 	);
@@ -965,25 +973,16 @@ const createModelClass = <T extends PersistentModel>(
 
 	Object.defineProperty(clazz, 'name', { value: modelDefinition.name });
 
-	for (const field in modelDefinition.fields) {
-		if (!isFieldAssociation(modelDefinition, field)) {
-			continue;
-		}
-
-		const {
-			type,
-			association: localAssociation,
-			association: { targetName, targetNames },
-		} = modelDefinition.fields[field] as Required<ModelField>;
-
-		const relationship = new ModelRelationship(
-			{
-				builder: clazz,
-				schema: modelDefinition,
-				pkField: extractPrimaryKeyFieldNames(modelDefinition),
-			},
-			field
-		);
+	// Add getters/setters for relationship fields.
+	//  getter - for lazy loading
+	//  setter - for FK management
+	const allModelRelationships = ModelRelationship.allFrom({
+		builder: clazz,
+		schema: modelDefinition,
+		pkField: extractPrimaryKeyFieldNames(modelDefinition),
+	});
+	for (const relationship of allModelRelationships) {
+		const field = relationship.field;
 
 		Object.defineProperty(clazz.prototype, modelDefinition.fields[field].name, {
 			set(model: PersistentModel | undefined | null) {
@@ -1018,6 +1017,7 @@ const createModelClass = <T extends PersistentModel>(
 					}
 				}
 
+				// if the relationship can be managed automagically, set the FK's
 				if (relationship.isComplete) {
 					for (let i = 0; i < relationship.localJoinFields.length; i++) {
 						this[relationship.localJoinFields[i]] =
@@ -1030,27 +1030,46 @@ const createModelClass = <T extends PersistentModel>(
 				}
 			},
 			get() {
+				/**
+				 * Bucket for holding related models instances specific to `this` instance.
+				 */
 				const instanceMemos = modelInstanceAssociationsMap.has(this)
 					? modelInstanceAssociationsMap.get(this)!
 					: modelInstanceAssociationsMap.set(this, {}).get(this)!;
 
+				// if the memos already has a result for this field, we'll use it.
+				// there is no "cache" invalidation of any kind; memos are permanent to
+				// keep an immutable perception of the instance.
 				if (!instanceMemos.hasOwnProperty(field)) {
+					// before we populate the memo, we need to know where to look for relatives.
+					// today, this only supports DataStore. Models aren't managed elsewhere in Amplify.
 					if (getAttachment(this) === ModelAttachment.DataStore) {
+						// when we fetch the results using a query constructed under the guidance
+						// of the relationship metadata, we DO NOT AWAIT resolution. we want to
+						// drop the promise into the memo's synchronously, eliminating the chance
+						// for a race.
 						const resultPromise = instance.query(
 							relationship.remoteModelConstructor as PersistentModelConstructor<T>,
 							base =>
 								base.and(q => {
 									return relationship.remoteJoinFields.map((field, index) => {
+										// TODO: anything we can use instead of `any` here?
 										return (q[field] as any).eq(
 											this[relationship.localJoinFields[index]]
 										);
 									});
 								})
 						);
 
+						// results in hand, how we return them to the caller depends on the relationship type.
 						if (relationship.type === 'HAS_MANY') {
+							// collections should support async iteration, even though we don't
+							// leverage it fully [yet].
 							instanceMemos[field] = new AsyncCollection(resultPromise);
 						} else {
+							// non-collections should only ever return 1 value *or nothing*.
+							// if we have more than 1 record, something's amiss. it's not our job
+							// pick a result for the customer. it's our job to say "something's wrong."
 							instanceMemos[field] = resultPromise.then(rows => {
 								if (rows.length > 1) {
 									// should never happen for a HAS_ONE or BELONGS_TO.
@@ -1084,15 +1103,34 @@ const createModelClass = <T extends PersistentModel>(
 	return clazz;
 };
 
+/**
+ * An eventually loaded related model instance.
+ */
 export class AsyncItem<T> extends Promise<T> {}
 
+/**
+ * A collection of related model instances.
+ *
+ * This collection can be async-iterated or turned directly into an array using `toArray()`.
+ */
 export class AsyncCollection<T> implements AsyncIterable<T> {
 	private values: Array<any> | Promise<Array<any>>;
 
 	constructor(values: Array<any> | Promise<Array<any>>) {
 		this.values = values;
 	}
 
+	/**
+	 * Facilitates async iteration.
+	 *
+	 * ```ts
+	 * for await (const item of collection) {
+	 *   handle(item)
+	 * }
+	 * ```
+	 *
+	 * @see https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Statements/for-await...of
+	 */
 	[Symbol.asyncIterator](): AsyncIterator<T> {
 		let values;
 		let index = 0;
@@ -1115,6 +1153,14 @@ export class AsyncCollection<T> implements AsyncIterable<T> {
 		};
 	}
 
+	/**
+	 * Turns the collection into an array, up to the amount specified in `max` param.
+	 *
+	 * ```ts
+	 * const all = await collection.toArray();
+	 * const first100 = await collection.toArray({max: 100});
+	 * ```
+	 */
 	async toArray({
 		max = Number.MAX_SAFE_INTEGER,
 	}: { max?: number } = {}): Promise<T[]> {
@@ -1332,6 +1378,8 @@ enum DataStoreState {
 	Clearing = 'Clearing',
 }
 
+// TODO: How can we get rid of the non-null assertions?
+// https://github.com/aws-amplify/amplify-js/pull/10477/files#r1007363485
 class DataStore {
 	// reference to configured category instances. Used for preserving SSR context
 	private Auth = Auth;
@@ -1656,7 +1704,7 @@ class DataStore {
 						const seedPredicate = recursivePredicateFor<T>({
 							builder: modelConstructor,
 							schema: modelDefinition,
-							pkField: getModelPKFieldName(modelConstructor),
+							pkField: extractPrimaryKeyFieldNames(modelDefinition),
 						});
 						const predicate = internals(
 							(identifierOrCriteria as RecursiveModelPredicateExtender<T>)(