@aws-amplify/api-graphql 4.0.29-unstable.2a7b3ac.0 → 4.0.29-unstable.f37faeb.0

package/src/internals/APIClient.ts

@@ -318,10 +318,10 @@ export type ModelOperation = keyof typeof graphQLOperationsInfo;
 
 const SELECTION_SET_WILDCARD = '*';
 
-export function defaultSelectionSetForNonModelWithIR(
+export const getDefaultSelectionSetForNonModelWithIR = (
 	nonModelDefinition: SchemaNonModel,
 	modelIntrospection: ModelIntrospectionSchema,
-): Record<string, unknown> {
+): Record<string, unknown> => {
 	const { fields } = nonModelDefinition;
 	const mappedFields = Object.values(fields)
 		.map(({ type, name }) => {
@@ -332,7 +332,7 @@ export function defaultSelectionSetForNonModelWithIR(
 			if (typeof (type as NonModelFieldType).nonModel === 'string') {
 				return [
 					name,
-					defaultSelectionSetForNonModelWithIR(
+					getDefaultSelectionSetForNonModelWithIR(
 						modelIntrospection.nonModels[(type as NonModelFieldType).nonModel],
 						modelIntrospection,
 					),
@@ -352,7 +352,47 @@ export function defaultSelectionSetForNonModelWithIR(
 		);
 
 	return Object.fromEntries(mappedFields);
-}
+};
+
+const getDefaultSelectionSetForModelWithIR = (
+	modelDefinition: SchemaModel,
+	modelIntrospection: ModelIntrospectionSchema,
+): Record<string, unknown> => {
+	const { fields } = modelDefinition;
+	const mappedFields = Object.values(fields)
+		.map(({ type, name }) => {
+			if (
+				typeof (type as { enum: string }).enum === 'string' ||
+				typeof type === 'string'
+			) {
+				return [name, FIELD_IR];
+			}
+
+			if (typeof (type as NonModelFieldType).nonModel === 'string') {
+				return [
+					name,
+					getDefaultSelectionSetForNonModelWithIR(
+						modelIntrospection.nonModels[(type as NonModelFieldType).nonModel],
+						modelIntrospection,
+					),
+				];
+			}
+
+			return undefined;
+		})
+		.filter(
+			(
+				pair: (string | Record<string, unknown>)[] | undefined,
+			): pair is (string | Record<string, unknown>)[] => pair !== undefined,
+		);
+
+	const ownerFields = resolveOwnerFields(modelDefinition).map(field => [
+		field,
+		FIELD_IR,
+	]);
+
+	return Object.fromEntries(mappedFields.concat(ownerFields));
+};
 
 function defaultSelectionSetForModel(modelDefinition: SchemaModel): string[] {
 	// fields that are explicitly part of the graphql schema; not
@@ -445,7 +485,7 @@ export function customSelectionSetToIR(
 
 			if (nested === SELECTION_SET_WILDCARD) {
 				result = {
-					[fieldName]: defaultSelectionSetForNonModelWithIR(
+					[fieldName]: getDefaultSelectionSetForNonModelWithIR(
 						relatedNonModelDefinition,
 						modelIntrospection,
 					),
@@ -471,8 +511,9 @@ export function customSelectionSetToIR(
 					modelIntrospection.models[relatedModel];
 
 				result = {
-					[fieldName]: modelsDefaultSelectionSetIR(
+					[fieldName]: getDefaultSelectionSetForModelWithIR(
 						nestedRelatedModelDefinition,
+						modelIntrospection,
 					),
 				};
 			} else {
@@ -529,23 +570,6 @@ export function customSelectionSetToIR(
 	);
 }
 
-const modelsDefaultSelectionSetIR = (relatedModelDefinition: SchemaModel) => {
-	const defaultSelectionSet = defaultSelectionSetForModel(
-		relatedModelDefinition,
-	);
-
-	const reduced = defaultSelectionSet.reduce(
-		(acc: Record<string, any>, curVal) => {
-			acc[curVal] = FIELD_IR;
-
-			return acc;
-		},
-		{},
-	);
-
-	return reduced;
-};
-
 /**
  * Stringifies selection set IR
  * * @example
@@ -697,6 +721,45 @@ export function generateGraphQLDocument(
 		selectionSet as ListArgs['selectionSet'],
 	);
 
+	// default PK args for get and list operations
+	// modified below for CPK
+	const getPkArgs = {
+		[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}!`,
+	};
+	const listPkArgs = {};
+
+	const generateSkArgs = (op: 'get' | 'list') => {
+		return sortKeyFieldNames.reduce(
+			(acc: Record<string, any>, fieldName: string) => {
+				const fieldType = fields[fieldName].type;
+
+				if (op === 'get') {
+					acc[fieldName] = `${fieldType}!`;
+				} else if (op === 'list') {
+					acc[fieldName] = `Model${fieldType}KeyConditionInput`;
+				}
+
+				return acc;
+			},
+			{},
+		);
+	};
+
+	if (isCustomPrimaryKey) {
+		Object.assign(getPkArgs, generateSkArgs('get'));
+
+		Object.assign(
+			listPkArgs,
+			{
+				// PK is only included in list query field args in the generated GQL
+				// when explicitly specifying PK with .identifier(['fieldName']) or @primaryKey in the schema definition
+				[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}`, // PK is always a nullable arg for list (no `!` after the type)
+				sortDirection: 'ModelSortDirection',
+			},
+			generateSkArgs('list'),
+		);
+	}
+
 	switch (modelOperation) {
 		case 'CREATE':
 		case 'UPDATE':
@@ -712,36 +775,15 @@ export function generateGraphQLDocument(
 		// TODO(Eslint): this this case clause correct without the break statement?
 		// eslint-disable-next-line no-fallthrough
 		case 'READ':
-			graphQLArguments ??
-				(graphQLArguments = isCustomPrimaryKey
-					? [primaryKeyFieldName, ...sortKeyFieldNames].reduce(
-							(acc: Record<string, any>, fieldName) => {
-								acc[fieldName] = `${fields[fieldName].type}!`;
-
-								return acc;
-							},
-							{},
-						)
-					: {
-							[primaryKeyFieldName]: `${fields[primaryKeyFieldName].type}!`,
-						});
+			graphQLArguments ?? (graphQLArguments = getPkArgs);
 			graphQLSelectionSet ?? (graphQLSelectionSet = selectionSetFields);
 		// TODO(Eslint): this this case clause correct without the break statement?
 		// eslint-disable-next-line no-fallthrough
 		case 'LIST':
 			graphQLArguments ??
 				(graphQLArguments = {
+					...listPkArgs,
 					filter: `Model${name}FilterInput`,
-					...(sortKeyFieldNames.length > 0
-						? [primaryKeyFieldName, ...sortKeyFieldNames].reduce(
-								(acc: Record<string, any>, fieldName) => {
-									acc[fieldName] = `${fields[fieldName].type}`;
-
-									return acc;
-								},
-								{ sortDirection: 'ModelSortDirection' },
-							)
-						: []),
 					limit: 'Int',
 					nextToken: 'String',
 				});

package/src/internals/operations/custom.ts

@@ -9,10 +9,10 @@ import { map } from 'rxjs';
 
 import {
 	authModeParams,
-	defaultSelectionSetForNonModelWithIR,
 	flattenItems,
 	generateSelectionSet,
 	getCustomHeaders,
+	getDefaultSelectionSetForNonModelWithIR,
 	initializeModel,
 	selectionSetIRToString,
 } from '../APIClient';
@@ -290,7 +290,7 @@ function operationSelectionSet(
 		const nonModel = modelIntrospection.nonModels[operation.type.nonModel];
 
 		return `{${selectionSetIRToString(
-			defaultSelectionSetForNonModelWithIR(nonModel, modelIntrospection),
+			getDefaultSelectionSetForNonModelWithIR(nonModel, modelIntrospection),
 		)}}`;
 	} else if (hasStringField(operation.type, 'model')) {
 		return `{${generateSelectionSet(modelIntrospection, operation.type.model)}}`;