fix: merge patches for consecutive copyOf (aws-amplify#9936)

* test: add unit test for consecutive copyOf

* fix: merge patches for consecutive copyOf

packages/datastore/__tests__/DataStore.ts

@@ -862,6 +862,77 @@ describe('DataStore tests', () => {
 			expect(model2.optionalField1).toBeNull();
 		});
 
+		test('multiple copyOf operations carry all changes on save', async () => {
+			let model: Model;
+			const save = jest.fn(() => [model]);
+			const query = jest.fn(() => [model]);
+
+			jest.resetModules();
+			jest.doMock('../src/storage/storage', () => {
+				const mock = jest.fn().mockImplementation(() => {
+					const _mock = {
+						init: jest.fn(),
+						save,
+						query,
+						runExclusive: jest.fn(fn => fn.bind(this, _mock)()),
+					};
+
+					return _mock;
+				});
+
+				(<any>mock).getNamespace = () => ({ models: {} });
+
+				return { ExclusiveStorage: mock };
+			});
+
+			({ initSchema, DataStore } = require('../src/datastore/datastore'));
+
+			const classes = initSchema(testSchema());
+
+			const { Model } = classes as { Model: PersistentModelConstructor<Model> };
+
+			const model1 = new Model({
+				dateCreated: new Date().toISOString(),
+				field1: 'original',
+				optionalField1: 'original',
+			});
+			model = model1;
+
+			await DataStore.save(model1);
+
+			const model2 = Model.copyOf(model1, draft => {
+				(<any>draft).field1 = 'field1Change1';
+				(<any>draft).optionalField1 = 'optionalField1Change1';
+			});
+
+			const model3 = Model.copyOf(model2, draft => {
+				(<any>draft).field1 = 'field1Change2';
+			});
+			model = model3;
+
+			await DataStore.save(model3);
+
+			const [settingsSave, saveOriginalModel, saveModel3] = <any>(
+				save.mock.calls
+			);
+
+			const [_model, _condition, _mutator, [patches]] = saveModel3;
+
+			const expectedPatches = [
+				{
+					op: 'replace',
+					path: ['field1'],
+					value: 'field1Change2',
+				},
+				{
+					op: 'replace',
+					path: ['optionalField1'],
+					value: 'optionalField1Change1',
+				},
+			];
+			expect(patches).toMatchObject(expectedPatches);
+		});
+
 		test('Non @model - Field cannot be changed', () => {
 			const { Metadata } = initSchema(testSchema()) as {
 				Metadata: NonModelTypeConstructor<Metadata>;
@@ -1109,9 +1180,9 @@ describe('DataStore tests', () => {
 
 			const expectedPatches2 = [
 				{
-					op: 'add',
-					path: ['emails', 3],
-					value: 'joe@doe.com',
+					op: 'replace',
+					path: ['emails'],
+					value: ['john@doe.com', 'jane@doe.com', 'joe@doe.com', 'joe@doe.com'],
 				},
 			];
 

packages/datastore/__tests__/util.test.ts

@@ -1,3 +1,4 @@
+import { enablePatches, produce, Patch } from 'immer';
 import {
 	isAWSDate,
 	isAWSDateTime,
@@ -11,6 +12,7 @@ import {
 	validatePredicateField,
 	valuesEqual,
 	processCompositeKeys,
+	mergePatches,
 } from '../src/util';
 
 describe('datastore util', () => {
@@ -592,4 +594,128 @@ describe('datastore util', () => {
 			expect(isAWSIPAddress(test)).toBe(false);
 		});
 	});
+
+	describe('mergePatches', () => {
+		enablePatches();
+		test('merge patches with no conflict', () => {
+			const modelA = {
+				foo: 'originalFoo',
+				bar: 'originalBar',
+			};
+			let patchesAB;
+			let patchesBC;
+			const modelB = produce(
+				modelA,
+				draft => {
+					draft.foo = 'newFoo';
+				},
+				patches => {
+					patchesAB = patches;
+				}
+			);
+			const modelC = produce(
+				modelB,
+				draft => {
+					draft.bar = 'newBar';
+				},
+				patches => {
+					patchesBC = patches;
+				}
+			);
+
+			const mergedPatches = mergePatches(modelA, patchesAB, patchesBC);
+			expect(mergedPatches).toEqual([
+				{
+					op: 'replace',
+					path: ['foo'],
+					value: 'newFoo',
+				},
+				{
+					op: 'replace',
+					path: ['bar'],
+					value: 'newBar',
+				},
+			]);
+		});
+		test('merge patches with conflict', () => {
+			const modelA = {
+				foo: 'originalFoo',
+				bar: 'originalBar',
+			};
+			let patchesAB;
+			let patchesBC;
+			const modelB = produce(
+				modelA,
+				draft => {
+					draft.foo = 'newFoo';
+					draft.bar = 'newBar';
+				},
+				patches => {
+					patchesAB = patches;
+				}
+			);
+			const modelC = produce(
+				modelB,
+				draft => {
+					draft.bar = 'newestBar';
+				},
+				patches => {
+					patchesBC = patches;
+				}
+			);
+
+			const mergedPatches = mergePatches(modelA, patchesAB, patchesBC);
+			expect(mergedPatches).toEqual([
+				{
+					op: 'replace',
+					path: ['foo'],
+					value: 'newFoo',
+				},
+				{
+					op: 'replace',
+					path: ['bar'],
+					value: 'newestBar',
+				},
+			]);
+		});
+		test('merge patches with conflict - list', () => {
+			const modelA = {
+				foo: [1, 2, 3],
+			};
+			let patchesAB;
+			let patchesBC;
+			const modelB = produce(
+				modelA,
+				draft => {
+					draft.foo.push(4);
+				},
+				patches => {
+					patchesAB = patches;
+				}
+			);
+			const modelC = produce(
+				modelB,
+				draft => {
+					draft.foo.push(5);
+				},
+				patches => {
+					patchesBC = patches;
+				}
+			);
+
+			const mergedPatches = mergePatches(modelA, patchesAB, patchesBC);
+			expect(mergedPatches).toEqual([
+				{
+					op: 'add',
+					path: ['foo', 3],
+					value: 4,
+				},
+				{
+					op: 'add',
+					path: ['foo', 4],
+					value: 5,
+				},
+			]);
+		});
+	});
 });

packages/datastore/src/datastore/datastore.ts

@@ -71,6 +71,7 @@ import {
 	sortCompareFunction,
 	DeferredCallbackResolver,
 	validatePredicate,
+	mergePatches,
 } from '../util';
 
 setAutoFreeze(true);
@@ -484,9 +485,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;

packages/datastore/src/util.ts

@@ -1,6 +1,7 @@
 import { Buffer } from 'buffer';
 import { monotonicFactory, ULID } from 'ulid';
 import { v4 as uuid } from 'uuid';
+import { produce, applyPatches, Patch } from 'immer';
 import { ModelInstanceCreator } from './datastore/datastore';
 import {
 	AllOperators,
@@ -146,7 +147,7 @@ export const isNonModelConstructor = (
 	return nonModelClasses.has(obj);
 };
 
-/* 
+/*
   When we have GSI(s) with composite sort keys defined on a model
 	There are some very particular rules regarding which fields must be included in the update mutation input
 	The field selection becomes more complex as the number of GSIs with composite sort keys grows
@@ -156,7 +157,7 @@ export const isNonModelConstructor = (
 	 2. all of the fields from any other composite sort key that intersect with the fields from 1.
 
 	 E.g.,
-	 Model @model 
+	 Model @model
 		@key(name: 'key1' fields: ['hk', 'a', 'b', 'c'])
 		@key(name: 'key2' fields: ['hk', 'a', 'b', 'd'])
 		@key(name: 'key3' fields: ['hk', 'x', 'y', 'z'])
@@ -192,7 +193,7 @@ export const processCompositeKeys = (
 		.filter(isModelAttributeCompositeKey)
 		.map(extractCompositeSortKey);
 
-	/* 
+	/*
 		if 2 sets of fields have any intersecting fields => combine them into 1 union set
 		e.g., ['a', 'b', 'c'] and ['a', 'b', 'd'] => ['a', 'b', 'c', 'd']
 	*/
@@ -773,3 +774,39 @@ export class DeferredCallbackResolver {
 		this.limitPromise.resolve(LimitTimerRaceResolvedValues.LIMIT);
 	}
 }
+
+/**
+ * merge two sets of patches created by immer produce.
+ * newPatches take precedent over oldPatches for patches modifying the same path.
+ * In the case many consecutive pathces are merged the original model should
+ * always be the root model.
+ *
+ * Example:
+ * A -> B, patches1
+ * B -> C, patches2
+ *
+ * mergePatches(A, patches1, patches2) to get patches for A -> C
+ *
+ * @param originalSource the original Model the patches should be applied to
+ * @param oldPatches immer produce patch list
+ * @param newPatches immer produce patch list (will take precedence)
+ * @return merged patches
+ */
+export function mergePatches<T>(
+	originalSource: T,
+	oldPatches: Patch[],
+	newPatches: Patch[]
+): Patch[] {
+	const patchesToMerge = oldPatches.concat(newPatches);
+	let patches: Patch[];
+	produce(
+		originalSource,
+		draft => {
+			applyPatches(draft, patchesToMerge);
+		},
+		p => {
+			patches = p;
+		}
+	);
+	return patches;
+}