checkpoint; initial fix... needs more testing

packages/datastore/__tests__/DataStore.ts

@@ -537,6 +537,59 @@ describe('DataStore observeQuery, with fake-indexeddb and fake sync', () => {
 			}
 		})();
 	});
+
+	test('attaches related belongsTo properties records just like query() on inserts', async done => {
+		try {
+			const expecteds = [5, 15];
+
+			for (let i = 0; i < 5; i++) {
+				await DataStore.save(
+					new Comment({
+						content: `comment content ${i}`,
+						post: await DataStore.save(
+							new Post({
+								title: `new post ${i}`,
+							})
+						),
+					})
+				);
+			}
+
+			const sub = DataStore.observeQuery(Comment).subscribe(
+				({ items, isSynced }) => {
+					const expected = expecteds.shift() || 0;
+					expect(items.length).toBe(expected);
+
+					for (let i = 0; i < expected; i++) {
+						expect(items[i].content).toEqual(`comment content ${i}`);
+						expect(items[i].post.title).toEqual(`new post ${i}`);
+					}
+
+					if (expecteds.length === 0) {
+						sub.unsubscribe();
+						done();
+					}
+				}
+			);
+
+			setTimeout(async () => {
+				for (let i = 5; i < 15; i++) {
+					await DataStore.save(
+						new Comment({
+							content: `comment content ${i}`,
+							post: await DataStore.save(
+								new Post({
+									title: `new post ${i}`,
+								})
+							),
+						})
+					);
+				}
+			}, 1);
+		} catch (error) {
+			done(error);
+		}
+	});
 });
 
 describe('DataStore tests', () => {

packages/datastore/src/datastore/datastore.ts

@@ -1145,24 +1145,70 @@ class DataStore {
 				handle = this.storage
 					.observe(modelConstructor, predicate)
 					.filter(({ model }) => namespaceResolver(model) === USER)
-					.map(
-						(event: InternalSubscriptionMessage<T>): SubscriptionMessage<T> => {
-							// The `element` returned by storage only contains updated fields.
-							// Intercept the event to send the `savedElement` so that the first
-							// snapshot returned to the consumer contains all fields.
-							// In the event of a delete we return `element`, as `savedElement`
-							// here is undefined.
-							const { opType, model, condition, element, savedElement } = event;
-
-							return {
-								opType,
-								element: savedElement || element,
-								model,
-								condition,
-							};
-						}
-					)
-					.subscribe(observer);
+					.subscribe({
+						next: async item => {
+							// the `element` for UPDATE events isn't a fully fleshed out instance of `modelConstructor`.
+
+							let message = item;
+
+							// as lnog 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 freshElement = await this.query(
+									item.model,
+									item.element.id
+								);
+								message = {
+									...message,
+									element: freshElement as T,
+								};
+							}
+
+							observer.next(message as SubscriptionMessage<T>);
+						},
+						error: err => observer.error(err),
+						complete: () => observer.complete(),
+					});
+
+				// .map(
+				// 	(item: InternalSubscriptionMessage<T>): SubscriptionMessage<T> => {
+
+				// 		// the `element` for UPDATE events isn't an instance of `modelConstructor`.
+				// 		// however, `executivePredicate` expects an instance that supports lazy loaded
+				// 		// associations. customers will presumably expect the same!
+				// 		let message = item;
+				// 		if (
+				// 			isModelConstructor(modelConstructor) &&
+				// 			!(item.element instanceof modelConstructor)
+				// 		) {
+				// 			message = {
+				// 				...message,
+				// 				element: modelInstanceCreator(modelConstructor, item.element),
+				// 			};
+				// 		}
+				// 		if (
+				// 			!executivePredicate ||
+				// 			(await executivePredicate.matches(message.element))
+				// 		) {
+				// 			observer.next(message as SubscriptionMessage<T>);
+				// 		}
+
+				// 		// The `element` returned by storage only contains updated fields.
+				// 		// Intercept the event to send the `savedElement` so that the first
+				// 		// snapshot returned to the consumer contains all fields.
+				// 		// In the event of a delete we return `element`, as `savedElement`
+				// 		// here is undefined.
+				// 		const { opType, model, condition, element, savedElement } = event;
+
+				// 		return {
+				// 			opType,
+				// 			element: savedElement || element,
+				// 			model,
+				// 			condition,
+				// 		};
+				// 	}
+				// )
+				// .subscribe(observer);
 			})();
 
 			return () => {