@aws-amplify/datastore 3.7.3 → 3.7.4-custom-pk.86

package/src/datastore/datastore.ts

@@ -1,4 +1,7 @@
+import API from '@aws-amplify/api';
 import { Amplify, ConsoleLogger as Logger, Hub, JS } from '@aws-amplify/core';
+import { Auth } from '@aws-amplify/auth';
+import Cache from '@aws-amplify/cache';
 import {
 	Draft,
 	immerable,
@@ -54,9 +57,17 @@ import {
 	isNonModelFieldType,
 	isModelFieldType,
 	ObserveQueryOptions,
+	ManagedIdentifier,
+	PersistentModelMetaData,
+	IdentifierFieldOrIdentifierObject,
+	__modelMeta__,
+	isIdentifierObject,
+	AmplifyContext,
+	isModelAttributePrimaryKey,
 } from '../types';
 import {
 	DATASTORE,
+	errorMessages,
 	establishRelationAndKeys,
 	exhaustiveCheck,
 	isModelConstructor,
@@ -69,7 +80,14 @@ import {
 	registerNonModelClass,
 	sortCompareFunction,
 	DeferredCallbackResolver,
+	extractPrimaryKeyFieldNames,
+	extractPrimaryKeysAndValues,
+	isIdManaged,
+	isIdOptionallyManaged,
+	validatePredicate,
+	mergePatches,
 } from '../util';
+import { getIdentifierValue } from '../sync/utils';
 
 setAutoFreeze(true);
 enablePatches();
@@ -79,11 +97,16 @@ const logger = new Logger('DataStore');
 const ulid = monotonicUlidFactory(Date.now());
 const { isNode } = JS.browserOrNode();
 
+type SettingMetaData = {
+	identifier: ManagedIdentifier<Setting, 'id'>;
+	readOnlyFields: never;
+};
 declare class Setting {
-	constructor(init: ModelInit<Setting>);
+	public readonly [__modelMeta__]: SettingMetaData;
+	constructor(init: ModelInit<Setting, SettingMetaData>);
 	static copyOf(
 		src: Setting,
-		mutator: (draft: MutableModel<Setting>) => void | Setting
+		mutator: (draft: MutableModel<Setting, SettingMetaData>) => void | Setting
 	): Setting;
 	public readonly id: string;
 	public readonly key: string;
@@ -222,6 +245,20 @@ const initSchema = (userSchema: Schema) => {
 	return userClasses;
 };
 
+/* Checks if the schema has been initialized by initSchema().
+ *
+ * Call this function before accessing schema.
+ * Currently this only needs to be called in start() and clear() because all other functions will call start first.
+ */
+const checkSchemaInitialized = () => {
+	if (schema === undefined) {
+		const message =
+			'Schema is not initialized. DataStore will not function as expected. This could happen if you have multiple versions of DataStore installed. Please see https://docs.amplify.aws/lib/troubleshooting/upgrading/q/platform/js/#check-for-duplicate-versions';
+		logger.error(message);
+		throw new Error(message);
+	}
+};
+
 const createTypeClasses: (
 	namespace: SchemaNamespace
 ) => TypeConstructorMap = namespace => {
@@ -246,16 +283,19 @@ const createTypeClasses: (
 
 export declare type ModelInstanceCreator = typeof modelInstanceCreator;
 
-const instancesMetadata = new WeakSet<
-	ModelInit<PersistentModel & Partial<ModelInstanceMetadata>>
->();
-function modelInstanceCreator<T extends PersistentModel = PersistentModel>(
+const instancesMetadata = new WeakSet<ModelInit<unknown, unknown>>();
+
+function modelInstanceCreator<T extends PersistentModel>(
 	modelConstructor: PersistentModelConstructor<T>,
-	init: ModelInit<T> & Partial<ModelInstanceMetadata>
+	init: Partial<T>
 ): T {
 	instancesMetadata.add(init);
 
-	return <T>new modelConstructor(init);
+	return new modelConstructor(<ModelInit<T, PersistentModelMetaData<T>>>init);
+}
+
+function isSchemaModel(m: SchemaModel | SchemaNonModel): m is SchemaModel {
+	return (m as SchemaModel).attributes !== undefined;
 }
 
 const validateModelFields =
@@ -273,6 +313,14 @@ const validateModelFields =
 				throw new Error(`Field ${name} is required`);
 			}
 
+			if (isSchemaModel(modelDefinition) && !isIdManaged(modelDefinition)) {
+				const keys = extractPrimaryKeyFieldNames(modelDefinition);
+				if (keys.includes(k) && v === '') {
+					logger.error(errorMessages.idEmptyString, { k, value: v });
+					throw new Error(errorMessages.idEmptyString);
+				}
+			}
+
 			if (isGraphQLScalarType(type)) {
 				const jsType = GraphQLScalarType.getJSType(type);
 				const validateScalar = GraphQLScalarType.getValidationFunction(type);
@@ -383,7 +431,7 @@ const castInstanceType = (
 	return v;
 };
 
-const initializeInstance = <T>(
+const initializeInstance = <T extends PersistentModel>(
 	init: ModelInit<T>,
 	modelDefinition: SchemaModel | SchemaNonModel,
 	draft: Draft<T & ModelInstanceMetadata>
@@ -407,31 +455,39 @@ const createModelClass = <T extends PersistentModel>(
 				(draft: Draft<T & ModelInstanceMetadata>) => {
 					initializeInstance(init, modelDefinition, draft);
 
+					// model is initialized inside a DataStore component (e.g. by Sync Engine, Storage Engine, etc.)
+					const isInternallyInitialized = instancesMetadata.has(init);
+
 					const modelInstanceMetadata: ModelInstanceMetadata =
-						instancesMetadata.has(init)
+						isInternallyInitialized
 							? <ModelInstanceMetadata>(<unknown>init)
 							: <ModelInstanceMetadata>{};
-					const {
-						id: _id,
-						_version,
-						_lastChangedAt,
-						_deleted,
-					} = modelInstanceMetadata;
-
-					// instancesIds are set by modelInstanceCreator, it is accessible only internally
-					const isInternal = _id !== null && _id !== undefined;
-
-					const id = isInternal
-						? _id
-						: modelDefinition.syncable
-						? uuid4()
-						: ulid();
-
-					if (!isInternal) {
+
+					type ModelWithIDIdentifier = { id: string };
+
+					const { id: _id } =
+						modelInstanceMetadata as unknown as ModelWithIDIdentifier;
+
+					if (isIdManaged(modelDefinition)) {
+						const isInternalModel = _id !== null && _id !== undefined;
+
+						const id = isInternalModel
+							? _id
+							: modelDefinition.syncable
+							? uuid4()
+							: ulid();
+
+						(<ModelWithIDIdentifier>(<unknown>draft)).id = id;
+					} else if (isIdOptionallyManaged(modelDefinition)) {
+						// only auto-populate if the id was not provided
+						(<ModelWithIDIdentifier>(<unknown>draft)).id = draft.id || uuid4();
+					}
+
+					if (!isInternallyInitialized) {
 						checkReadOnlyPropertyOnCreate(draft, modelDefinition);
 					}
 
-					draft.id = id;
+					const { _version, _lastChangedAt, _deleted } = modelInstanceMetadata;
 
 					if (modelDefinition.syncable) {
 						draft._version = _version;
@@ -456,8 +512,13 @@ const createModelClass = <T extends PersistentModel>(
 			const model = produce(
 				source,
 				draft => {
-					fn(<MutableModel<T>>(draft as unknown));
-					draft.id = source.id;
+					fn(<MutableModel<T>>draft);
+
+					const keyNames = extractPrimaryKeyFieldNames(modelDefinition);
+					// Keys are immutable
+					// @ts-ignore TODO: fix type
+					keyNames.forEach(key => (draft[key] = source[key]));
+
 					const modelValidator = validateModelFields(modelDefinition);
 					Object.entries(draft).forEach(([k, v]) => {
 						const parsedValue = castInstanceType(modelDefinition, k, v);
@@ -468,9 +529,21 @@ const createModelClass = <T extends PersistentModel>(
 				p => (patches = p)
 			);
 
-			if (patches.length) {
-				modelPatchesMap.set(model, [patches, source]);
-				checkReadOnlyPropertyOnUpdate(patches, modelDefinition);
+			const hasExistingPatches = modelPatchesMap.has(source);
+			if (patches.length || hasExistingPatches) {
+				if (hasExistingPatches) {
+					const [existingPatches, existingSource] = modelPatchesMap.get(source);
+					const mergedPatches = mergePatches(
+						existingSource,
+						existingPatches,
+						patches
+					);
+					modelPatchesMap.set(model, [mergedPatches, existingSource]);
+					checkReadOnlyPropertyOnUpdate(mergedPatches, modelDefinition);
+				} else {
+					modelPatchesMap.set(model, [patches, source]);
+					checkReadOnlyPropertyOnUpdate(patches, modelDefinition);
+				}
 			}
 
 			return model;
@@ -484,6 +557,7 @@ const createModelClass = <T extends PersistentModel>(
 			}
 
 			const instance = modelInstanceCreator(clazz, json);
+
 			const modelValidator = validateModelFields(modelDefinition);
 
 			Object.entries(instance).forEach(([k, v]) => {
@@ -531,7 +605,9 @@ const checkReadOnlyPropertyOnUpdate = (
 	});
 };
 
-const createNonModelClass = <T>(typeDefinition: SchemaNonModel) => {
+const createNonModelClass = <T extends PersistentModel>(
+	typeDefinition: SchemaNonModel
+) => {
 	const clazz = <NonModelTypeConstructor<T>>(<unknown>class Model {
 		constructor(init: ModelInit<T>) {
 			const instance = produce(
@@ -564,7 +640,7 @@ function defaultConflictHandler(conflictData: SyncConflict): PersistentModel {
 	return modelInstanceCreator(modelConstructor, { ...localModel, _version });
 }
 
-function defaultErrorHandler(error: SyncError) {
+function defaultErrorHandler(error: SyncError<PersistentModel>): void {
 	logger.warn(error);
 }
 
@@ -615,7 +691,6 @@ async function checkSchemaVersion(
 		const [schemaVersionSetting] = await s.query(
 			Setting,
 			ModelPredicateCreator.createFromExisting(modelDefinition, c =>
-				// @ts-ignore Argument of type '"eq"' is not assignable to parameter of type 'never'.
 				c.key('eq', SETTING_SCHEMA_VERSION)
 			),
 			{ page: 0, limit: 1 }
@@ -682,23 +757,34 @@ function getNamespace(): SchemaNamespace {
 }
 
 class DataStore {
+	// reference to configured category instances. Used for preserving SSR context
+	Auth = Auth;
+	API = API;
+	Cache = Cache;
+
 	private amplifyConfig: Record<string, any> = {};
 	private authModeStrategy: AuthModeStrategy;
 	private conflictHandler: ConflictHandler;
-	private errorHandler: (error: SyncError) => void;
+	private errorHandler: (error: SyncError<PersistentModel>) => void;
 	private fullSyncInterval: number;
-	private initialized: Promise<void>;
+	private initialized?: Promise<void>;
 	private initReject: Function;
 	private initResolve: Function;
 	private maxRecordsToSync: number;
-	private storage: Storage;
-	private sync: SyncEngine;
+	private storage?: Storage;
+	private sync?: SyncEngine;
 	private syncPageSize: number;
 	private syncExpressions: SyncExpression[];
 	private syncPredicates: WeakMap<SchemaModel, ModelPredicate<any>> =
 		new WeakMap<SchemaModel, ModelPredicate<any>>();
 	private sessionId: string;
 	private storageAdapter: Adapter;
+	// object that gets passed to descendent classes. Allows us to pass these down by reference
+	private amplifyContext: AmplifyContext = {
+		Auth: this.Auth,
+		API: this.API,
+		Cache: this.Cache,
+	};
 
 	getModuleName() {
 		return 'DataStore';
@@ -728,6 +814,7 @@ class DataStore {
 
 		await this.storage.init();
 
+		checkSchemaInitialized();
 		await checkSchemaVersion(this.storage, schema.version);
 
 		const { aws_appsync_graphqlEndpoint } = this.amplifyConfig;
@@ -744,13 +831,12 @@ class DataStore {
 				userClasses,
 				this.storage,
 				modelInstanceCreator,
-				this.maxRecordsToSync,
-				this.syncPageSize,
 				this.conflictHandler,
 				this.errorHandler,
 				this.syncPredicates,
 				this.amplifyConfig,
-				this.authModeStrategy
+				this.authModeStrategy,
+				this.amplifyContext
 			);
 
 			// tslint:disable-next-line:max-line-length
@@ -796,7 +882,10 @@ class DataStore {
 	query: {
 		<T extends PersistentModel>(
 			modelConstructor: PersistentModelConstructor<T>,
-			id: string
+			identifier: IdentifierFieldOrIdentifierObject<
+				T,
+				PersistentModelMetaData<T>
+			>
 		): Promise<T | undefined>;
 		<T extends PersistentModel>(
 			modelConstructor: PersistentModelConstructor<T>,
@@ -805,7 +894,10 @@ class DataStore {
 		): Promise<T[]>;
 	} = async <T extends PersistentModel>(
 		modelConstructor: PersistentModelConstructor<T>,
-		idOrCriteria?: string | ProducerModelPredicate<T> | typeof PredicateAll,
+		identifierOrCriteria?:
+			| IdentifierFieldOrIdentifierObject<T, PersistentModelMetaData<T>>
+			| ProducerModelPredicate<T>
+			| typeof PredicateAll,
 		paginationProducer?: ProducerPaginationInput<T>
 	): Promise<T | T[] | undefined> => {
 		await this.start();
@@ -819,28 +911,37 @@ class DataStore {
 			throw new Error(msg);
 		}
 
-		if (typeof idOrCriteria === 'string') {
+		if (typeof identifierOrCriteria === 'string') {
 			if (paginationProducer !== undefined) {
 				logger.warn('Pagination is ignored when querying by id');
 			}
 		}
 
 		const modelDefinition = getModelDefinition(modelConstructor);
+		const keyFields = extractPrimaryKeyFieldNames(modelDefinition);
+
 		let predicate: ModelPredicate<T>;
 
-		if (isQueryOne(idOrCriteria)) {
-			predicate = ModelPredicateCreator.createForId<T>(
+		if (isQueryOne(identifierOrCriteria)) {
+			predicate = ModelPredicateCreator.createForSingleField<T>(
 				modelDefinition,
-				idOrCriteria
+				keyFields[0],
+				identifierOrCriteria
 			);
 		} else {
-			if (isPredicatesAll(idOrCriteria)) {
+			// Object is being queried using object literal syntax
+			if (isIdentifierObject(<T>identifierOrCriteria, modelDefinition)) {
+				predicate = ModelPredicateCreator.createForPk<T>(
+					modelDefinition,
+					<T>identifierOrCriteria
+				);
+			} else if (isPredicatesAll(identifierOrCriteria)) {
 				// Predicates.ALL means "all records", so no predicate (undefined)
 				predicate = undefined;
 			} else {
 				predicate = ModelPredicateCreator.createFromExisting(
 					modelDefinition,
-					idOrCriteria
+					<any>identifierOrCriteria
 				);
 			}
 		}
@@ -870,7 +971,13 @@ class DataStore {
 			pagination
 		);
 
-		return isQueryOne(idOrCriteria) ? result[0] : result;
+		const isSingleFieldIdentifier = keyFields.length === 1;
+		const returnOne =
+			isSingleFieldIdentifier &&
+			(isQueryOne(identifierOrCriteria) ||
+				isIdentifierObject(identifierOrCriteria, modelDefinition));
+
+		return returnOne ? result[0] : result;
 	};
 
 	save = async <T extends PersistentModel>(
@@ -883,7 +990,7 @@ class DataStore {
 		// Allows us to only include changed fields for updates
 		const patchesTuple = modelPatchesMap.get(model);
 
-		const modelConstructor: PersistentModelConstructor<T> = model
+		const modelConstructor: PersistentModelConstructor<T> | undefined = model
 			? <PersistentModelConstructor<T>>model.constructor
 			: undefined;
 
@@ -898,15 +1005,15 @@ class DataStore {
 
 		const producedCondition = ModelPredicateCreator.createFromExisting(
 			modelDefinition,
-			condition
+			condition!
 		);
 
 		const [savedModel] = await this.storage.runExclusive(async s => {
 			await s.save(model, producedCondition, undefined, patchesTuple);
 
-			return s.query(
+			return s.query<T>(
 				modelConstructor,
-				ModelPredicateCreator.createForId(modelDefinition, model.id)
+				ModelPredicateCreator.createForPk(modelDefinition, model)
 			);
 		});
 
@@ -946,22 +1053,28 @@ class DataStore {
 	};
 
 	delete: {
-		<T extends PersistentModel>(
-			model: T,
-			condition?: ProducerModelPredicate<T>
-		): Promise<T>;
 		<T extends PersistentModel>(
 			modelConstructor: PersistentModelConstructor<T>,
-			id: string
+			identifier: IdentifierFieldOrIdentifierObject<
+				T,
+				PersistentModelMetaData<T>
+			>
 		): Promise<T[]>;
 		<T extends PersistentModel>(
 			modelConstructor: PersistentModelConstructor<T>,
 			condition: ProducerModelPredicate<T> | typeof PredicateAll
 		): Promise<T[]>;
+		<T extends PersistentModel>(
+			model: T,
+			condition?: ProducerModelPredicate<T>
+		): Promise<T>;
 	} = async <T extends PersistentModel>(
 		modelOrConstructor: T | PersistentModelConstructor<T>,
-		idOrCriteria?: string | ProducerModelPredicate<T> | typeof PredicateAll
-	) => {
+		identifierOrCriteria?:
+			| IdentifierFieldOrIdentifierObject<T, PersistentModelMetaData<T>>
+			| ProducerModelPredicate<T>
+			| typeof PredicateAll
+	): Promise<T | T[]> => {
 		await this.start();
 
 		let condition: ModelPredicate<T>;
@@ -973,31 +1086,43 @@ class DataStore {
 			throw new Error(msg);
 		}
 
-		if (isValidModelConstructor(modelOrConstructor)) {
+		if (isValidModelConstructor<T>(modelOrConstructor)) {
 			const modelConstructor = modelOrConstructor;
 
-			if (!idOrCriteria) {
+			if (!identifierOrCriteria) {
 				const msg =
 					'Id to delete or criteria required. Do you want to delete all? Pass Predicates.ALL';
-				logger.error(msg, { idOrCriteria });
+				logger.error(msg, { identifierOrCriteria });
 
 				throw new Error(msg);
 			}
 
-			if (typeof idOrCriteria === 'string') {
-				condition = ModelPredicateCreator.createForId<T>(
+			const modelDefinition = getModelDefinition(modelConstructor);
+
+			if (typeof identifierOrCriteria === 'string') {
+				const [keyField] = extractPrimaryKeyFieldNames(modelDefinition);
+
+				condition = ModelPredicateCreator.createForSingleField<T>(
 					getModelDefinition(modelConstructor),
-					idOrCriteria
+					keyField,
+					identifierOrCriteria
 				);
 			} else {
-				condition = ModelPredicateCreator.createFromExisting(
-					getModelDefinition(modelConstructor),
-					/**
-					 * idOrCriteria is always a ProducerModelPredicate<T>, never a symbol.
-					 * The symbol is used only for typing purposes. e.g. see Predicates.ALL
-					 */
-					idOrCriteria as ProducerModelPredicate<T>
-				);
+				if (isIdentifierObject(identifierOrCriteria, modelDefinition)) {
+					condition = ModelPredicateCreator.createForPk<T>(
+						modelDefinition,
+						<T>identifierOrCriteria
+					);
+				} else {
+					condition = ModelPredicateCreator.createFromExisting(
+						modelDefinition,
+						/**
+						 * idOrCriteria is always a ProducerModelPredicate<T>, never a symbol.
+						 * The symbol is used only for typing purposes. e.g. see Predicates.ALL
+						 */
+						identifierOrCriteria as ProducerModelPredicate<T>
+					);
+				}
 
 				if (!condition || !ModelPredicateCreator.isValidPredicate(condition)) {
 					const msg =
@@ -1025,22 +1150,24 @@ class DataStore {
 
 			const modelDefinition = getModelDefinition(modelConstructor);
 
-			const idPredicate = ModelPredicateCreator.createForId<T>(
+			const pkPredicate = ModelPredicateCreator.createForPk<T>(
 				modelDefinition,
-				model.id
+				model
 			);
 
-			if (idOrCriteria) {
-				if (typeof idOrCriteria !== 'function') {
+			if (identifierOrCriteria) {
+				if (typeof identifierOrCriteria !== 'function') {
 					const msg = 'Invalid criteria';
-					logger.error(msg, { idOrCriteria });
+					logger.error(msg, { identifierOrCriteria });
 
 					throw new Error(msg);
 				}
 
-				condition = idOrCriteria(idPredicate);
+				condition = (<ProducerModelPredicate<T>>identifierOrCriteria)(
+					pkPredicate
+				);
 			} else {
-				condition = idPredicate;
+				condition = pkPredicate;
 			}
 
 			const [[deleted]] = await this.storage.delete(model, condition);
@@ -1052,20 +1179,28 @@ class DataStore {
 	observe: {
 		(): Observable<SubscriptionMessage<PersistentModel>>;
 
-		<T extends PersistentModel>(model: T): Observable<SubscriptionMessage<T>>;
+		<T extends PersistentModel>(
+			modelConstructor: PersistentModelConstructor<T>,
+			identifier: string
+		): Observable<SubscriptionMessage<T>>;
 
 		<T extends PersistentModel>(
 			modelConstructor: PersistentModelConstructor<T>,
-			criteria?: string | ProducerModelPredicate<T>
+			criteria?: ProducerModelPredicate<T> | typeof PredicateAll
 		): Observable<SubscriptionMessage<T>>;
-	} = <T extends PersistentModel = PersistentModel>(
+
+		<T extends PersistentModel>(model: T): Observable<SubscriptionMessage<T>>;
+	} = <T extends PersistentModel>(
 		modelOrConstructor?: T | PersistentModelConstructor<T>,
-		idOrCriteria?: string | ProducerModelPredicate<T>
+		identifierOrCriteria?:
+			| string
+			| ProducerModelPredicate<T>
+			| typeof PredicateAll
 	): Observable<SubscriptionMessage<T>> => {
 		let predicate: ModelPredicate<T>;
 
-		const modelConstructor: PersistentModelConstructor<T> =
-			modelOrConstructor && isValidModelConstructor(modelOrConstructor)
+		const modelConstructor: PersistentModelConstructor<T> | undefined =
+			modelOrConstructor && isValidModelConstructor<T>(modelOrConstructor)
 				? modelOrConstructor
 				: undefined;
 
@@ -1075,10 +1210,10 @@ class DataStore {
 				model && (<Object>Object.getPrototypeOf(model)).constructor;
 
 			if (isValidModelConstructor<T>(modelConstructor)) {
-				if (idOrCriteria) {
+				if (identifierOrCriteria) {
 					logger.warn('idOrCriteria is ignored when using a model instance', {
 						model,
-						idOrCriteria,
+						identifierOrCriteria,
 					});
 				}
 
@@ -1092,9 +1227,9 @@ class DataStore {
 			}
 		}
 
-		if (idOrCriteria !== undefined && modelConstructor === undefined) {
+		if (identifierOrCriteria !== undefined && modelConstructor === undefined) {
 			const msg = 'Cannot provide criteria without a modelConstructor';
-			logger.error(msg, idOrCriteria);
+			logger.error(msg, identifierOrCriteria);
 			throw new Error(msg);
 		}
 
@@ -1105,18 +1240,26 @@ class DataStore {
 			throw new Error(msg);
 		}
 
-		if (typeof idOrCriteria === 'string') {
-			predicate = ModelPredicateCreator.createForId<T>(
+		if (typeof identifierOrCriteria === 'string') {
+			const modelDefinition = getModelDefinition(modelConstructor);
+			const [keyField] = extractPrimaryKeyFieldNames(modelDefinition);
+
+			predicate = ModelPredicateCreator.createForSingleField<T>(
 				getModelDefinition(modelConstructor),
-				idOrCriteria
+				keyField,
+				identifierOrCriteria
 			);
 		} else {
-			predicate =
-				modelConstructor &&
-				ModelPredicateCreator.createFromExisting<T>(
-					getModelDefinition(modelConstructor),
-					idOrCriteria
-				);
+			if (isPredicatesAll(identifierOrCriteria)) {
+				predicate = undefined;
+			} else {
+				predicate =
+					modelConstructor &&
+					ModelPredicateCreator.createFromExisting<T>(
+						getModelDefinition(modelConstructor),
+						identifierOrCriteria
+					);
+			}
 		}
 
 		return new Observable<SubscriptionMessage<T>>(observer => {
@@ -1125,10 +1268,43 @@ class DataStore {
 			(async () => {
 				await this.start();
 
+				// Filter the events returned by Storage according to namespace,
+				// append original element data, and subscribe to the observable
 				handle = this.storage
 					.observe(modelConstructor, predicate)
 					.filter(({ model }) => namespaceResolver(model) === USER)
-					.subscribe(observer);
+					.subscribe({
+						next: async item => {
+							// the `element` doesn't necessarily contain all item details or
+							// have related records attached consistently with that of a query()
+							// result item. for consistency, we attach them here.
+
+							let message = item;
+
+							// as long as we're not dealing with a DELETE, we need to fetch a fresh
+							// item from storage to ensure it's fully populated.
+							if (item.opType !== 'DELETE') {
+								const modelDefinition = getModelDefinition(item.model);
+								const keyFields = extractPrimaryKeyFieldNames(modelDefinition);
+								const primaryKeysAndValues = extractPrimaryKeysAndValues(
+									item.element,
+									keyFields
+								);
+								const freshElement = await this.query(
+									item.model,
+									primaryKeysAndValues
+								);
+								message = {
+									...message,
+									element: freshElement as T,
+								};
+							}
+
+							observer.next(message as SubscriptionMessage<T>);
+						},
+						error: err => observer.error(err),
+						complete: () => observer.complete(),
+					});
 			})();
 
 			return () => {
@@ -1145,7 +1321,7 @@ class DataStore {
 			criteria?: ProducerModelPredicate<T> | typeof PredicateAll,
 			paginationProducer?: ObserveQueryOptions<T>
 		): Observable<DataStoreSnapshot<T>>;
-	} = <T extends PersistentModel = PersistentModel>(
+	} = <T extends PersistentModel>(
 		model: PersistentModelConstructor<T>,
 		criteria?: ProducerModelPredicate<T> | typeof PredicateAll,
 		options?: ObserveQueryOptions<T>
@@ -1155,7 +1331,18 @@ class DataStore {
 			const itemsChanged = new Map<string, T>();
 			let deletedItemIds: string[] = [];
 			let handle: ZenObservable.Subscription;
-
+			let predicate: ModelPredicate<T>;
+
+			/**
+			 * As the name suggests, this geneates a snapshot in the form of
+			 * 	`{items: T[], isSynced: boolean}`
+			 * and sends it to the observer.
+			 *
+			 * SIDE EFFECT: The underlying generation and emission methods may touch:
+			 * `items`, `itemsChanged`, and `deletedItemIds`.
+			 *
+			 * Refer to `generateSnapshot` and `emitSnapshot` for more details.
+			 */
 			const generateAndEmitSnapshot = (): void => {
 				const snapshot = generateSnapshot();
 				emitSnapshot(snapshot);
@@ -1172,41 +1359,101 @@ class DataStore {
 			const { sort } = options || {};
 			const sortOptions = sort ? { sort } : undefined;
 
+			const modelDefinition = getModelDefinition(model);
+			const keyFields = extractPrimaryKeyFieldNames(modelDefinition);
+
+			if (isQueryOne(criteria)) {
+				predicate = ModelPredicateCreator.createForSingleField<T>(
+					modelDefinition,
+					keyFields[0],
+					criteria
+				);
+			} else {
+				if (isPredicatesAll(criteria)) {
+					// Predicates.ALL means "all records", so no predicate (undefined)
+					predicate = undefined;
+				} else {
+					predicate = ModelPredicateCreator.createFromExisting(
+						modelDefinition,
+						criteria
+					);
+				}
+			}
+
+			const { predicates, type: predicateGroupType } =
+				ModelPredicateCreator.getPredicates(predicate, false) || {};
+			const hasPredicate = !!predicates;
+
 			(async () => {
 				try {
 					// first, query and return any locally-available records
-					(await this.query(model, criteria, sortOptions)).forEach(item =>
-						items.set(item.id, item)
-					);
-
-					// observe the model and send a stream of updates (debounced)
-					handle = this.observe(
-						model,
-						// @ts-ignore TODO: fix this TSlint error
-						criteria
-					).subscribe(({ element, model, opType }) => {
-						// Flag items which have been recently deleted
-						// NOTE: Merging of separate operations to the same model instance is handled upstream
-						// in the `mergePage` method within src/sync/merger.ts. The final state of a model instance
-						// depends on the LATEST record (for a given id).
-						if (opType === 'DELETE') {
-							deletedItemIds.push(element.id);
-						} else {
-							itemsChanged.set(element.id, element);
+					(await this.query(model, criteria, sortOptions)).forEach(item => {
+						let record = item;
+						// TODO: fix query
+						if (Array.isArray(item)) {
+							record = item[0];
 						}
+						const itemModelDefinition = getModelDefinition(model);
+						const idOrPk = getIdentifierValue(itemModelDefinition, record);
+						items.set(idOrPk, record);
+					});
+
+					// Observe the model and send a stream of updates (debounced).
+					// We need to post-filter results instead of passing criteria through
+					// to have visibility into items that move from in-set to out-of-set.
+					// We need to explicitly remove those items from the existing snapshot.
+					handle = this.observe(model).subscribe(
+						({ element, model, opType }) => {
+							let record = element;
+
+							// TODO: fix query
+							if (Array.isArray(element)) {
+								record = element[0];
+							}
+							const itemModelDefinition = getModelDefinition(model);
+							const idOrPk = getIdentifierValue(itemModelDefinition, record);
+							if (
+								hasPredicate &&
+								!validatePredicate(record, predicateGroupType, predicates)
+							) {
+								if (
+									opType === 'UPDATE' &&
+									(items.has(idOrPk) || itemsChanged.has(idOrPk))
+								) {
+									// tracking as a "deleted item" will include the item in
+									// page limit calculations and ensure it is removed from the
+									// final items collection, regardless of which collection(s)
+									// it is currently in. (I mean, it could be in both, right!?)
+									deletedItemIds.push(idOrPk);
+								} else {
+									// ignore updates for irrelevant/filtered items.
+									return;
+								}
+							}
 
-						const isSynced = this.sync?.getModelSyncedStatus(model) ?? false;
+							// Flag items which have been recently deleted
+							// NOTE: Merging of separate operations to the same model instance is handled upstream
+							// in the `mergePage` method within src/sync/merger.ts. The final state of a model instance
+							// depends on the LATEST record (for a given id).
+							if (opType === 'DELETE') {
+								deletedItemIds.push(idOrPk);
+							} else {
+								itemsChanged.set(idOrPk, record);
+							}
 
-						const limit =
-							itemsChanged.size - deletedItemIds.length >= this.syncPageSize;
+							const isSynced = this.sync?.getModelSyncedStatus(model) ?? false;
 
-						if (limit || isSynced) {
-							limitTimerRace.resolve();
-						}
+							const limit =
+								itemsChanged.size - deletedItemIds.length >= this.syncPageSize;
 
-						// kicks off every subsequent race as results sync down
-						limitTimerRace.start();
-					});
+							if (limit || isSynced) {
+								limitTimerRace.resolve();
+							}
+
+							// kicks off every subsequent race as results sync down
+							limitTimerRace.start();
+						}
+					);
 
 					// returns a set of initial/locally-available results
 					generateAndEmitSnapshot();
@@ -1215,7 +1462,12 @@ class DataStore {
 				}
 			})();
 
-			// TODO: abstract this function into a util file to be able to write better unit tests
+			/**
+			 * Combines the `items`, `itemsChanged`, and `deletedItemIds` collections into
+			 * a snapshot in the form of `{ items: T[], isSynced: boolean}`.
+			 *
+			 * SIDE EFFECT: The shared `items` collection is recreated.
+			 */
 			const generateSnapshot = (): DataStoreSnapshot<T> => {
 				const isSynced = this.sync?.getModelSyncedStatus(model) ?? false;
 				const itemsArray = [
@@ -1228,10 +1480,21 @@ class DataStore {
 				}
 
 				items.clear();
-				itemsArray.forEach(item => items.set(item.id, item));
+				itemsArray.forEach(item => {
+					// CPK TODO: fix query
+					let record = item;
+
+					if (Array.isArray(item)) {
+						record = item[0];
+					}
+
+					const itemModelDefinition = getModelDefinition(model);
+					const idOrPk = getIdentifierValue(itemModelDefinition, record);
+					items.set(idOrPk, record);
+				});
 
 				// remove deleted items from the final result set
-				deletedItemIds.forEach(id => items.delete(id));
+				deletedItemIds.forEach(idOrPk => items.delete(idOrPk));
 
 				return {
 					items: Array.from(items.values()),
@@ -1239,6 +1502,14 @@ class DataStore {
 				};
 			};
 
+			/**
+			 * Emits the list of items to the observer.
+			 *
+			 * SIDE EFFECT: `itemsChanged` and `deletedItemIds` are cleared to prepare
+			 * for the next snapshot.
+			 *
+			 * @param snapshot The generated items data to emit.
+			 */
 			const emitSnapshot = (snapshot: DataStoreSnapshot<T>): void => {
 				// send the generated snapshot to the primary subscription
 				observer.next(snapshot);
@@ -1248,6 +1519,12 @@ class DataStore {
 				deletedItemIds = [];
 			};
 
+			/**
+			 * Sorts an `Array` of `T` according to the sort instructions given in the
+			 * original  `observeQuery()` call.
+			 *
+			 * @param itemsToSort A array of model type.
+			 */
 			const sortItems = (itemsToSort: T[]): void => {
 				const modelDefinition = getModelDefinition(model);
 				const pagination = this.processPagination(modelDefinition, options);
@@ -1262,7 +1539,14 @@ class DataStore {
 				}
 			};
 
-			// send one last snapshot when the model is fully synced
+			/**
+			 * Force one last snapshot when the model is fully synced.
+			 *
+			 * This reduces latency for that last snapshot, which will otherwise
+			 * wait for the configured timeout.
+			 *
+			 * @param payload The payload from the Hub event.
+			 */
 			const hubCallback = ({ payload }): void => {
 				const { event, data } = payload;
 				if (
@@ -1284,6 +1568,10 @@ class DataStore {
 	};
 
 	configure = (config: DataStoreConfig = {}) => {
+		this.amplifyContext.Auth = this.Auth;
+		this.amplifyContext.API = this.API;
+		this.amplifyContext.Cache = this.Cache;
+
 		const {
 			DataStore: configDataStore,
 			authModeStrategyType: configAuthModeStrategyType,
@@ -1313,7 +1601,7 @@ class DataStore {
 
 		switch (authModeStrategyType) {
 			case AuthModeStrategyType.MULTI_AUTH:
-				this.authModeStrategy = multiAuthStrategy;
+				this.authModeStrategy = multiAuthStrategy(this.amplifyContext);
 				break;
 			case AuthModeStrategyType.DEFAULT:
 				this.authModeStrategy = defaultAuthStrategy;
@@ -1329,12 +1617,13 @@ class DataStore {
 
 		this.syncExpressions =
 			(configDataStore && configDataStore.syncExpressions) ||
-			this.syncExpressions ||
-			configSyncExpressions;
+			configSyncExpressions ||
+			this.syncExpressions;
 
 		this.maxRecordsToSync =
 			(configDataStore && configDataStore.maxRecordsToSync) ||
 			configMaxRecordsToSync ||
+			this.maxRecordsToSync ||
 			10000;
 
 		// store on config object, so that Sync, Subscription, and Mutation processors can have access
@@ -1343,6 +1632,7 @@ class DataStore {
 		this.syncPageSize =
 			(configDataStore && configDataStore.syncPageSize) ||
 			configSyncPageSize ||
+			this.syncPageSize ||
 			1000;
 
 		// store on config object, so that Sync, Subscription, and Mutation processors can have access
@@ -1350,22 +1640,32 @@ class DataStore {
 
 		this.fullSyncInterval =
 			(configDataStore && configDataStore.fullSyncInterval) ||
-			this.fullSyncInterval ||
 			configFullSyncInterval ||
+			this.fullSyncInterval ||
 			24 * 60; // 1 day
 
 		this.storageAdapter =
 			(configDataStore && configDataStore.storageAdapter) ||
-			this.storageAdapter ||
 			configStorageAdapter ||
+			this.storageAdapter ||
 			undefined;
 
-		this.sessionId = this.retrieveSessionId();
+		this.sessionId = this.retrieveSessionId()!;
 	};
 
 	clear = async function clear() {
+		checkSchemaInitialized();
 		if (this.storage === undefined) {
-			return;
+			// connect to storage so that it can be cleared without fully starting DataStore
+			this.storage = new Storage(
+				schema,
+				namespaceResolver,
+				getModelConstructorByModelName,
+				modelInstanceCreator,
+				this.storageAdapter,
+				this.sessionId
+			);
+			await this.storage.init();
 		}
 
 		if (syncSubscription && !syncSubscription.closed) {
@@ -1384,7 +1684,7 @@ class DataStore {
 		this.syncPredicates = new WeakMap<SchemaModel, ModelPredicate<any>>();
 	};
 
-	stop = async function stop() {
+	stop = async function stop(this: InstanceType<typeof DataStore>) {
 		if (this.initialized !== undefined) {
 			await this.start();
 		}
@@ -1548,9 +1848,9 @@ class DataStore {
 
 				return `${sessionId}-${appSyncId}`;
 			}
-		} catch {
-			return undefined;
-		}
+		} catch {}
+
+		return undefined;
 	}
 }