From 7c5ab39ddc3e33236c84c70085049660b8b8d03c Mon Sep 17 00:00:00 2001
From: Justin Shih <36183898+Jshhhh@users.noreply.github.com>
Date: Tue, 10 Jan 2023 16:40:37 -0800
Subject: [PATCH] fix: parse and stringify nonmodel fields (#882)

Co-authored-by: Justin Shih <jushih@amazon.com>
---
 ...studio-ui-codegen-react-forms.test.ts.snap |  25 +++-
 .../forms/form-renderer-helper/cta-props.ts   |  63 ++------
 .../forms/form-renderer-helper/form-state.ts  |  68 ++++++++-
 .../form-renderer-helper/parse-fields.ts      | 134 ++++++++++++++++++
 4 files changed, 233 insertions(+), 57 deletions(-)
 create mode 100644 packages/codegen-ui-react/lib/forms/form-renderer-helper/parse-fields.ts

diff --git a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
index 85c3bb069..67422a608 100644
--- a/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
+++ b/packages/codegen-ui-react/lib/__tests__/__snapshots__/studio-ui-codegen-react-forms.test.ts.snap
@@ -2872,6 +2872,9 @@ export default function MyPostForm(props) {
               nonModelFieldArray: modelFields.nonModelFieldArray.map((s) =>
                 JSON.parse(s)
               ),
+              nonModelField: modelFields.nonModelField
+                ? JSON.parse(modelFields.nonModelField)
+                : modelFields.nonModelField,
             })
           );
           if (onSuccess) {
@@ -6876,6 +6879,9 @@ export default function MyPostForm(props) {
               nonModelFieldArray: modelFields.nonModelFieldArray.map((s) =>
                 JSON.parse(s)
               ),
+              nonModelField: modelFields.nonModelField
+                ? JSON.parse(modelFields.nonModelField)
+                : modelFields.nonModelField,
             })
           );
           if (onSuccess) {
@@ -10603,7 +10609,11 @@ export default function MyPostForm(props) {
         ? cleanValues.nonModelField
         : JSON.stringify(cleanValues.nonModelField)
     );
-    setNonModelFieldArray(cleanValues.nonModelFieldArray ?? []);
+    setNonModelFieldArray(
+      cleanValues.nonModelFieldArray?.map((item) =>
+        typeof item === \\"string\\" ? item : JSON.stringify(item)
+      ) ?? []
+    );
     setCurrentNonModelFieldArrayValue(\\"\\");
     setErrors({});
   };
@@ -10693,7 +10703,15 @@ export default function MyPostForm(props) {
           });
           await DataStore.save(
             Post.copyOf(postRecord, (updated) => {
-              Object.assign(updated, modelFields);
+              Object.assign(updated, {
+                ...modelFields,
+                nonModelFieldArray: modelFields.nonModelFieldArray.map((s) =>
+                  JSON.parse(s)
+                ),
+                nonModelField: modelFields.nonModelField
+                  ? JSON.parse(modelFields.nonModelField)
+                  : modelFields.nonModelField,
+              });
             })
           );
           if (onSuccess) {
@@ -16580,6 +16598,9 @@ export default function PostCreateFormRow(props) {
               nonModelFieldArray: modelFields.nonModelFieldArray.map((s) =>
                 JSON.parse(s)
               ),
+              nonModelField: modelFields.nonModelField
+                ? JSON.parse(modelFields.nonModelField)
+                : modelFields.nonModelField,
             })
           );
           if (onSuccess) {
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper/cta-props.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper/cta-props.ts
index 0ac22fb94..fa3f798dd 100644
--- a/packages/codegen-ui-react/lib/forms/form-renderer-helper/cta-props.ts
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper/cta-props.ts
@@ -13,7 +13,7 @@
   See the License for the specific language governing permissions and
   limitations under the License.
  */
-import { FieldConfigMetadata, GenericDataSchema } from '@aws-amplify/codegen-ui';
+import { FieldConfigMetadata, GenericDataSchema, isNonModelDataType } from '@aws-amplify/codegen-ui';
 import {
   factory,
   NodeFlags,
@@ -34,6 +34,7 @@ import {
 import { isManyToManyRelationship } from './map-from-fieldConfigs';
 import { ImportCollection } from '../../imports';
 import { getBiDirectionalRelationshipStatements } from './bidirectional-relationship';
+import { generateParsePropertyAssignments, generateUpdateModelObject } from './parse-fields';
 
 const getRecordUpdateDataStoreCallExpression = ({
   modelName,
@@ -91,7 +92,10 @@ const getRecordUpdateDataStoreCallExpression = ({
                       factory.createIdentifier('assign'),
                     ),
                     undefined,
-                    [factory.createIdentifier(updatedObjectName), factory.createIdentifier(modelFieldsObjectName)],
+                    [
+                      factory.createIdentifier(updatedObjectName),
+                      generateUpdateModelObject(fieldConfigs, modelFieldsObjectName),
+                    ],
                   ),
                 ),
                 ...relationshipBasedUpdates,
@@ -205,12 +209,17 @@ export const buildDataStoreExpression = (
   const hasManyRelationshipFields: string[] = [];
   const nonModelArrayFields: string[] = [];
   const savedRecordName = lowerCaseFirst(modelName);
+  const nonModelFields: string[] = [];
 
   Object.entries(fieldConfigs).forEach((fieldConfig) => {
     const [fieldName, fieldConfigMetaData] = fieldConfig;
     const { dataType, isArray } = fieldConfigMetaData;
-    if (isArray && dataType && typeof dataType === 'object' && 'nonModel' in dataType) {
-      nonModelArrayFields.push(fieldName);
+    if (isNonModelDataType(dataType)) {
+      if (isArray) {
+        nonModelArrayFields.push(fieldName);
+      } else {
+        nonModelFields.push(fieldName);
+      }
     }
     relationshipsPromisesAccessStatements.push(
       ...getBiDirectionalRelationshipStatements({
@@ -279,51 +288,7 @@ export const buildDataStoreExpression = (
     );
   });
 
-  nonModelArrayFields.forEach((field) => {
-    // nonModelFieldArray: modelFields.nonModelFieldArray.map(s => JSON.parse(s))
-    propertyAssignments.push(
-      factory.createPropertyAssignment(
-        factory.createIdentifier(field),
-        factory.createCallExpression(
-          factory.createPropertyAccessExpression(
-            factory.createPropertyAccessExpression(
-              factory.createIdentifier('modelFields'),
-              factory.createIdentifier(field),
-            ),
-            factory.createIdentifier('map'),
-          ),
-          undefined,
-          [
-            factory.createArrowFunction(
-              undefined,
-              undefined,
-              [
-                factory.createParameterDeclaration(
-                  undefined,
-                  undefined,
-                  undefined,
-                  factory.createIdentifier('s'),
-                  undefined,
-                  undefined,
-                  undefined,
-                ),
-              ],
-              undefined,
-              factory.createToken(SyntaxKind.EqualsGreaterThanToken),
-              factory.createCallExpression(
-                factory.createPropertyAccessExpression(
-                  factory.createIdentifier('JSON'),
-                  factory.createIdentifier('parse'),
-                ),
-                undefined,
-                [factory.createIdentifier('s')],
-              ),
-            ),
-          ],
-        ),
-      ),
-    );
-  });
+  propertyAssignments.push(...generateParsePropertyAssignments(nonModelArrayFields, nonModelFields));
 
   const modelFieldsObject = propertyAssignments.length
     ? factory.createObjectLiteralExpression(
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper/form-state.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper/form-state.ts
index 8c15c39af..f3ed7cd21 100644
--- a/packages/codegen-ui-react/lib/forms/form-renderer-helper/form-state.ts
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper/form-state.ts
@@ -37,6 +37,7 @@ import {
   PropertyAccessExpression,
   ElementAccessExpression,
   ConditionalExpression,
+  CallChain,
 } from 'typescript';
 import { capitalizeFirstLetter, lowerCaseFirst, getSetNameIdentifier, buildUseStateExpression } from '../../helpers';
 import { getElementAccessExpression } from './invalid-variable-helpers';
@@ -259,22 +260,22 @@ export const resetStateFunction = (fieldConfigs: Record<string, FieldConfigMetad
     const renderedName = sanitizedFieldName || stateName;
     if (!stateNames.has(stateName)) {
       const accessExpression = getElementAccessExpression(recordOrInitialValues, stateName);
+      const isNonModelField = isNonModelDataType(dataType);
 
       // Initial values should have the correct values and not need a modifier
-      if (
-        (dataType === 'AWSJSON' || isNonModelDataType(dataType)) &&
-        !isArray &&
-        recordOrInitialValues !== 'initialValues'
-      ) {
+      if ((dataType === 'AWSJSON' || isNonModelField) && !isArray && recordOrInitialValues !== 'initialValues') {
         const awsJSONAccessModifier = stringifyAWSJSONFieldValue(accessExpression);
         acc.push(setStateExpression(renderedName, awsJSONAccessModifier));
       } else {
+        const stringifiedOrAccessExpression = isNonModelField
+          ? stringifyAWSJSONFieldArrayValues(accessExpression)
+          : accessExpression;
         acc.push(
           setStateExpression(
             renderedName,
             isArray && recordOrInitialValues === 'cleanValues'
               ? factory.createBinaryExpression(
-                  accessExpression,
+                  stringifiedOrAccessExpression,
                   factory.createToken(SyntaxKind.QuestionQuestionToken),
                   factory.createArrayLiteralExpression([], false),
                 )
@@ -408,6 +409,61 @@ const stringifyAWSJSONFieldValue = (
   );
 };
 
+/**
+ * Datastore allows JSON strings and normal JSON so make sure items in array are string type
+ *
+ * Example output:
+ * cleanValues.nonModelFieldArray?.map(item => typeof item === "string" ? item : JSON.stringify(item))
+ */
+const stringifyAWSJSONFieldArrayValues = (value: PropertyAccessExpression | ElementAccessExpression): CallChain => {
+  return factory.createCallChain(
+    factory.createPropertyAccessChain(
+      value,
+      factory.createToken(SyntaxKind.QuestionDotToken),
+      factory.createIdentifier('map'),
+    ),
+    undefined,
+    undefined,
+    [
+      factory.createArrowFunction(
+        undefined,
+        undefined,
+        [
+          factory.createParameterDeclaration(
+            undefined,
+            undefined,
+            undefined,
+            factory.createIdentifier('item'),
+            undefined,
+            undefined,
+            undefined,
+          ),
+        ],
+        undefined,
+        factory.createToken(SyntaxKind.EqualsGreaterThanToken),
+        factory.createConditionalExpression(
+          factory.createBinaryExpression(
+            factory.createTypeOfExpression(factory.createIdentifier('item')),
+            factory.createToken(SyntaxKind.EqualsEqualsEqualsToken),
+            factory.createStringLiteral('string'),
+          ),
+          factory.createToken(SyntaxKind.QuestionToken),
+          factory.createIdentifier('item'),
+          factory.createToken(SyntaxKind.ColonToken),
+          factory.createCallExpression(
+            factory.createPropertyAccessExpression(
+              factory.createIdentifier('JSON'),
+              factory.createIdentifier('stringify'),
+            ),
+            undefined,
+            [factory.createIdentifier('item')],
+          ),
+        ),
+      ),
+    ],
+  );
+};
+
 /**
  * turns ['myNestedObject', 'value', 'nestedValue', 'leaf']
  *
diff --git a/packages/codegen-ui-react/lib/forms/form-renderer-helper/parse-fields.ts b/packages/codegen-ui-react/lib/forms/form-renderer-helper/parse-fields.ts
new file mode 100644
index 000000000..5fe159181
--- /dev/null
+++ b/packages/codegen-ui-react/lib/forms/form-renderer-helper/parse-fields.ts
@@ -0,0 +1,134 @@
+/*
+  Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+
+  Licensed under the Apache License, Version 2.0 (the "License").
+  You may not use this file except in compliance with the License.
+  You may obtain a copy of the License at
+
+      http://www.apache.org/licenses/LICENSE-2.0
+
+  Unless required by applicable law or agreed to in writing, software
+  distributed under the License is distributed on an "AS IS" BASIS,
+  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+  See the License for the specific language governing permissions and
+  limitations under the License.
+ */
+import { isNonModelDataType, FieldConfigMetadata } from '@aws-amplify/codegen-ui';
+import {
+  PropertyAccessExpression,
+  Identifier,
+  factory,
+  SyntaxKind,
+  Expression,
+  ObjectLiteralExpression,
+} from 'typescript';
+/**
+ * JSON.parse(s)
+ */
+export const parseValue = (expression: Expression) =>
+  factory.createCallExpression(
+    factory.createPropertyAccessExpression(factory.createIdentifier('JSON'), factory.createIdentifier('parse')),
+    undefined,
+    [expression],
+  );
+
+/**
+ * modelFields.nonModelFieldArray.map(s => JSON.parse(item))
+ */
+export const parseArrayValues = (accessName: PropertyAccessExpression | Identifier) => {
+  return factory.createCallExpression(
+    factory.createPropertyAccessExpression(accessName, factory.createIdentifier('map')),
+    undefined,
+    [
+      factory.createArrowFunction(
+        undefined,
+        undefined,
+        [
+          factory.createParameterDeclaration(
+            undefined,
+            undefined,
+            undefined,
+            factory.createIdentifier('s'),
+            undefined,
+            undefined,
+          ),
+        ],
+        undefined,
+        factory.createToken(SyntaxKind.EqualsGreaterThanToken),
+        parseValue(factory.createIdentifier('s')),
+      ),
+    ],
+  );
+};
+
+/**
+ * arrayFields = nonModelFieldArray: modelFields.nonModelFieldArray.map(s => JSON.parse(item))
+ * singleFields =
+ *  nonModelField: modelFields.nonModelField ? JSON.parse(modelFields.nonModelField) : modelFields.nonModelField
+ */
+export const generateParsePropertyAssignments = (arrayFields: string[], nonArrayFields: string[]) => {
+  const parseArrayFields = arrayFields.map((field) =>
+    factory.createPropertyAssignment(
+      factory.createIdentifier(field),
+      parseArrayValues(
+        factory.createPropertyAccessExpression(
+          factory.createIdentifier('modelFields'),
+          factory.createIdentifier(field),
+        ),
+      ),
+    ),
+  );
+  const parseFields = nonArrayFields.map((field) =>
+    factory.createPropertyAssignment(
+      factory.createIdentifier(field),
+      factory.createConditionalExpression(
+        factory.createPropertyAccessExpression(
+          factory.createIdentifier('modelFields'),
+          factory.createIdentifier(field),
+        ),
+        factory.createToken(SyntaxKind.QuestionToken),
+        parseValue(
+          factory.createPropertyAccessExpression(
+            factory.createIdentifier('modelFields'),
+            factory.createIdentifier(field),
+          ),
+        ),
+        factory.createToken(SyntaxKind.ColonToken),
+        factory.createPropertyAccessExpression(
+          factory.createIdentifier('modelFields'),
+          factory.createIdentifier(field),
+        ),
+      ),
+    ),
+  );
+  return [...parseArrayFields, ...parseFields];
+};
+
+//
+export const generateUpdateModelObject = (
+  fieldConfigs: Record<string, FieldConfigMetadata>,
+  modelFieldsObjectName: string,
+) => {
+  const nonModelFields: string[] = [];
+  const nonModelArrayFields: string[] = [];
+
+  Object.entries(fieldConfigs).forEach(([name, { dataType, sanitizedFieldName, isArray }]) => {
+    if (isNonModelDataType(dataType)) {
+      const renderedFieldName = sanitizedFieldName || name;
+      if (!isArray) {
+        nonModelFields.push(renderedFieldName);
+      } else {
+        nonModelArrayFields.push(renderedFieldName);
+      }
+    }
+  });
+  const parsePropertyAssignments = generateParsePropertyAssignments(nonModelArrayFields, nonModelFields);
+  let updateModelObject: ObjectLiteralExpression | Identifier = factory.createIdentifier(modelFieldsObjectName);
+  if (parsePropertyAssignments.length) {
+    updateModelObject = factory.createObjectLiteralExpression(
+      [factory.createSpreadAssignment(factory.createIdentifier(modelFieldsObjectName)), ...parsePropertyAssignments],
+      true,
+    );
+  }
+  return updateModelObject;
+};