fix selection set generation (#12324)

.vscode/launch.json

@@ -10,12 +10,12 @@
 			"type": "node",
 			"request": "launch",
 			// The debugger will only run tests for the package specified here:
-			"cwd": "${workspaceFolder}/packages/datastore",
+			"cwd": "${workspaceFolder}/packages/api",
 			"runtimeArgs": [
 				"--inspect-brk",
 				"${workspaceRoot}/node_modules/.bin/jest",
 				// Optionally specify a single test file to run/debug:
-				"sync.test.ts",
+				"APIClient.test.ts",
 				"--runInBand",
 				"--testTimeout",
 				"600000", // 10 min timeout so jest doesn't error while we're stepping through code

packages/api/__tests__/models/APIClient.test.ts

@@ -1,4 +1,9 @@
-import { normalizeMutationInput, flattenItems } from '../../src/APIClient';
+import {
+	normalizeMutationInput,
+	flattenItems,
+	generateSelectionSet,
+	customSelectionSetToIR,
+} from '../../src/APIClient';
 import modelIntroSchema from '../assets/model-introspection';
 
 describe('APIClient', () => {
@@ -155,3 +160,117 @@ describe('flattenItems', () => {
 		expect(flattened).toEqual(expected);
 	});
 });
+
+describe('customSelectionSetToIR', () => {
+	test('1', () => {
+		const selSet = customSelectionSetToIR(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+		]);
+
+		const expected = {
+			postId: '',
+			title: '',
+		};
+
+		expect(selSet).toEqual(expected);
+	});
+
+	test('2', () => {
+		const selSet = customSelectionSetToIR(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+			'comments.id',
+			'comments.createdAt',
+		]);
+
+		const expected = {
+			postId: '',
+			title: '',
+			comments: {
+				items: {
+					id: '',
+					createdAt: '',
+				},
+			},
+		};
+
+		expect(selSet).toEqual(expected);
+	});
+
+	test('3', () => {
+		const selSet = customSelectionSetToIR(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+			'comments.*',
+		]);
+
+		const expected = {
+			postId: '',
+			title: '',
+			comments: {
+				items: {
+					id: '',
+					bingo: '',
+					anotherField: '',
+					createdAt: '',
+					updatedAt: '',
+					postCommentsPostId: '',
+					postCommentsTitle: '',
+					postComments2PostId: '',
+					postComments2Title: '',
+				},
+			},
+		};
+
+		expect(selSet).toEqual(expected);
+	});
+});
+
+describe('generateSelectionSet', () => {
+	test('it should generate default selection set', () => {
+		const selSet = generateSelectionSet(modelIntroSchema.models, 'Post');
+
+		const expected =
+			'postId title summary viewCount createdAt updatedAt postAuthorId';
+
+		expect(selSet).toEqual(expected);
+	});
+
+	test('it should generate custom selection set - top-level fields', () => {
+		const selSet = generateSelectionSet(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+		]);
+
+		const expected = 'postId title';
+
+		expect(selSet).toEqual(expected);
+	});
+
+	test('it should generate custom selection set - specific nested fields', () => {
+		const selSet = generateSelectionSet(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+			'comments.id',
+			'comments.createdAt',
+		]);
+
+		const expected = 'postId title comments { items { id createdAt } }';
+
+		expect(selSet).toEqual(expected);
+	});
+
+	test('it should generate custom selection set - all nested fields', () => {
+		const selSet = generateSelectionSet(modelIntroSchema.models, 'Post', [
+			'postId',
+			'title',
+			'comments.*',
+		]);
+
+		const expected =
+			'postId title comments { items { id bingo anotherField createdAt updatedAt postCommentsPostId postCommentsTitle postComments2PostId postComments2Title } }';
+
+		expect(selSet).toEqual(expected);
+	});
+});

packages/api/src/API.ts

@@ -141,7 +141,8 @@ export class APIClass extends InternalAPIClass {
 							const query = generateGraphQLDocument(
 								modelIntrospection.models,
 								name,
-								operation
+								operation,
+								options
 							);
 							const variables = buildGraphQLVariables(
 								model,
@@ -160,16 +161,21 @@ export class APIClass extends InternalAPIClass {
 							// flatten response
 							if (res.data !== undefined) {
 								const [key] = Object.keys(res.data);
-
-								// TODO: refactor to avoid destructuring here
-								const [initialized] = initializeModel(
-									client,
-									name,
-									[res.data[key]],
-									modelIntrospection
-								);
-
-								return initialized;
+								const flattenedResult = flattenItems(res.data)[key];
+
+								if (options?.selectionSet) {
+									return flattenedResult;
+								} else {
+									// TODO: refactor to avoid destructuring here
+									const [initialized] = initializeModel(
+										client,
+										name,
+										[flattenedResult],
+										modelIntrospection
+									);
+
+									return initialized;
+								}
 							}
 
 							return res;

packages/api/src/APIClient.ts

@@ -161,27 +161,22 @@ type OperationPrefix =
 const graphQLDocumentsCache = new Map<string, Map<ModelOperation, string>>();
 const SELECTION_SET_ALL_NESTED = '*';
 
-function defaultSelectionSetForModel(modelDefinition: any): string {
+function defaultSelectionSetForModel(modelDefinition: any): string[] {
 	const { fields } = modelDefinition;
 	return Object.values<any>(fields)
 		.map(({ type, name }) => typeof type === 'string' && name) // Default selection set omits model fields
-		.filter(Boolean)
-		.join(' ');
+		.filter(Boolean);
 }
 
-function generateSelectionSet(
+export function customSelectionSetToIR(
 	modelIntrospection: any,
 	modelName: string,
-	selectionSet?: string[]
+	selectionSet: string[]
 ) {
 	const modelDefinition = modelIntrospection[modelName];
 	const { fields } = modelDefinition;
 
-	if (!selectionSet) {
-		return defaultSelectionSetForModel(modelDefinition);
-	}
-
-	const selSet: string[] = [];
+	const intermediateSelectionSet: Record<string, string | object> = {};
 
 	for (const f of selectionSet) {
 		const nested = f.includes('.');
@@ -198,14 +193,43 @@ function generateSelectionSet(
 
 			if (selectedField === SELECTION_SET_ALL_NESTED) {
 				const relatedModelDefinition = modelIntrospection[relatedModel];
+
 				const defaultSelectionSet = defaultSelectionSetForModel(
 					relatedModelDefinition
 				);
 
+				const reduced = defaultSelectionSet.reduce((acc, curVal) => {
+					acc[curVal] = '';
+					return acc;
+				}, {});
+
+				if (fields[modelFieldName]?.isArray) {
+					intermediateSelectionSet[modelFieldName] = {
+						items: reduced,
+					};
+				} else {
+					intermediateSelectionSet[modelFieldName] = reduced;
+				}
+			} else {
+				const getNestedSelSet = customSelectionSetToIR(
+					modelIntrospection,
+					relatedModel,
+					[selectedField]
+				);
+
 				if (fields[modelFieldName]?.isArray) {
-					selSet.push(`${modelFieldName} { items { ${defaultSelectionSet} } }`);
+					const existing = (intermediateSelectionSet as any)[
+						modelFieldName
+					] || { items: {} };
+					const merged = { ...existing.items, ...getNestedSelSet };
+
+					intermediateSelectionSet[modelFieldName] = { items: merged };
 				} else {
-					selSet.push(`${modelFieldName} { ${defaultSelectionSet} }`);
+					const existingItems =
+						(intermediateSelectionSet as any)[modelFieldName] || {};
+					const merged = { ...existingItems, ...getNestedSelSet };
+
+					intermediateSelectionSet[modelFieldName] = merged;
 				}
 			}
 		} else {
@@ -215,11 +239,62 @@ function generateSelectionSet(
 				throw Error(`${f} is not a field of model ${modelName}`);
 			}
 
-			selSet.push(f);
+			intermediateSelectionSet[f] = '';
 		}
 	}
 
-	return selSet.join(' ');
+	return intermediateSelectionSet;
+}
+
+const FIELD = '';
+
+export function selectionSetIRToString(
+	obj: Record<string, string | any>
+): string {
+	const res: string[] = [];
+
+	Object.entries(obj).forEach(([fieldName, value]) => {
+		if (value === FIELD) {
+			res.push(fieldName);
+		} else if (typeof value === 'object' && value !== null) {
+			if (value?.items) {
+				res.push(
+					fieldName,
+					'{',
+					'items',
+					'{',
+					selectionSetIRToString(value.items),
+					'}',
+					'}'
+				);
+			} else {
+				res.push(fieldName, '{', selectionSetIRToString(value), '}');
+			}
+		}
+	});
+
+	return res.join(' ');
+}
+
+export function generateSelectionSet(
+	modelIntrospection: any,
+	modelName: string,
+	selectionSet?: string[]
+) {
+	const modelDefinition = modelIntrospection[modelName];
+
+	if (!selectionSet) {
+		return defaultSelectionSetForModel(modelDefinition).join(' ');
+	}
+
+	const selSetIr = customSelectionSetToIR(
+		modelIntrospection,
+		modelName,
+		selectionSet
+	);
+	const selSetString = selectionSetIRToString(selSetIr);
+
+	return selSetString;
 }
 
 export function generateGraphQLDocument(