@aws-amplify/data-schema 1.2.7 → 1.2.9

package/src/ModelSchema.ts

@@ -117,6 +117,10 @@ type RDSModelSchemaFunctions =
   | 'renameModelFields'
   | 'renameModels';
 
+type OmitFromEach<Models, Modifier extends string> = {
+  [ModelName in keyof Models]: Omit<Models[ModelName], Modifier>;
+};
+
 export type RDSModelSchema<
   T extends RDSModelSchemaParamShape,
   UsedMethods extends RDSModelSchemaFunctions = never,
@@ -173,7 +177,7 @@ export type RDSModelSchema<
     >;
     setAuthorization: (
       callback: (
-        models: BaseSchema<T, true>['models'],
+        models: OmitFromEach<BaseSchema<T, true>['models'], 'secondaryIndexes'>,
         schema: RDSModelSchema<T, UsedMethods | 'setAuthorization'>,
       ) => void,
     ) => RDSModelSchema<T>;
@@ -182,7 +186,12 @@ export type RDSModelSchema<
         Partial<Record<keyof T['types'], RelationshipTemplate>>
       >,
     >(
-      callback: (models: BaseSchema<T, true>['models']) => Relationships,
+      callback: (
+        models: OmitFromEach<
+          BaseSchema<T, true>['models'],
+          'authorization' | 'fields' | 'secondaryIndexes'
+        >,
+      ) => Relationships,
     ) => RDSModelSchema<
       UnionToIntersection<Relationships[number]> extends infer RelationshipsDefs
         ? RelationshipsDefs extends Record<string, RelationshipTemplate>

package/src/SchemaProcessor.ts

@@ -1245,13 +1245,29 @@ const schemaPreprocessor = (
           getRefType,
         );
 
+        topLevelTypes.push(...implicitTypes);
+
         mergeCustomTypeAuthRules(
           customTypeInheritedAuthRules,
           customTypeAuthRules,
         );
-        Object.assign(lambdaFunctions, lambdaFunctionDefinition);
 
-        topLevelTypes.push(...implicitTypes);
+        if (customTypeAuthRules) {
+          const nestedCustomTypeNames = extractNestedCustomTypeNames(
+            customTypeAuthRules,
+            topLevelTypes,
+            getRefType,
+          );
+
+          for (const nestedCustomType of nestedCustomTypeNames) {
+            mergeCustomTypeAuthRules(customTypeInheritedAuthRules, {
+              typeName: nestedCustomType,
+              authRules: customTypeAuthRules.authRules, // apply the same auth rules as the top-level custom type
+            });
+          }
+        }
+
+        Object.assign(lambdaFunctions, lambdaFunctionDefinition);
 
         if (jsFunctionForField) {
           jsFunctions.push(jsFunctionForField);
@@ -1666,6 +1682,55 @@ function generateCustomOperationTypes({
   return types;
 }
 
+function extractNestedCustomTypeNames(
+  customTypeAuthRules: CustomTypeAuthRules,
+  topLevelTypes: [string, any][],
+  getRefType: ReturnType<typeof getRefTypeForSchema>,
+): string[] {
+  if (!customTypeAuthRules) {
+    return [];
+  }
+
+  const [_, customTypeDef] = topLevelTypes.find(
+    ([topLevelTypeName]) => customTypeAuthRules.typeName === topLevelTypeName,
+  )!;
+
+  // traverse the custom type's fields and extract any nested custom type names.
+  // Those nested custom types also inherit the custom op's auth configuration.
+  // Supports both inline custom types and refs to custom types
+  const traverseCustomTypeFields = (
+    name: string,
+    typeDef: any,
+    namesList: string[] = [],
+  ) => {
+    const fields = typeDef.data.fields as Record<string, any>;
+
+    for (const [fieldName, fieldDef] of Object.entries(fields)) {
+      if (isCustomType(fieldDef)) {
+        const customTypeName = `${capitalize(name)}${capitalize(fieldName)}`;
+        namesList.push(customTypeName);
+        traverseCustomTypeFields(customTypeName, fieldDef, namesList);
+      } else if (isRefField(fieldDef)) {
+        const refType = getRefType(fieldDef.data.link, name);
+
+        if (refType.type === 'CustomType') {
+          namesList.push(fieldDef.data.link);
+          traverseCustomTypeFields(fieldDef.data.link, refType.def, namesList);
+        }
+      }
+    }
+
+    return namesList;
+  };
+
+  const nestedCustomTypeNames = traverseCustomTypeFields(
+    customTypeAuthRules.typeName,
+    customTypeDef,
+  );
+
+  return nestedCustomTypeNames;
+}
+
 /**
  * Returns API definition from ModelSchema or string schema
  * @param arg - { schema }