From 2d1a261754b1a03e3fc5eba18f9963b894c99bf8 Mon Sep 17 00:00:00 2001
From: David Lopez <lopezbnd@amazon.com>
Date: Thu, 10 Nov 2022 10:36:07 -0800
Subject: [PATCH] fix: stringify existing AWSJSON field value

---
 ...studio-ui-codegen-react-forms.test.ts.snap |  6 +-
 .../codegen-ui-react/lib/forms/form-state.ts  | 60 +++++++++++++++----
 2 files changed, 53 insertions(+), 13 deletions(-)

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 ab25a6830..c78acc861 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
@@ -3414,7 +3414,11 @@ export default function MyPostForm(props) {
     setUsername(cleanValues.username);
     setProfile_url(cleanValues.profile_url);
     setPost_url(cleanValues.post_url);
-    setMetadata(cleanValues.metadata);
+    setMetadata(
+      typeof cleanValues.metadata === \\"string\\"
+        ? cleanValues.metadata
+        : JSON.stringify(cleanValues.metadata)
+    );
     setErrors({});
   };
   const [postRecord, setPostRecord] = React.useState(post);
diff --git a/packages/codegen-ui-react/lib/forms/form-state.ts b/packages/codegen-ui-react/lib/forms/form-state.ts
index e6f4e8a8e..b0db5ad81 100644
--- a/packages/codegen-ui-react/lib/forms/form-state.ts
+++ b/packages/codegen-ui-react/lib/forms/form-state.ts
@@ -28,6 +28,9 @@ import {
   PropertyAccessChain,
   PropertyAssignment,
   PropertyName,
+  PropertyAccessExpression,
+  ElementAccessExpression,
+  ConditionalExpression,
 } from 'typescript';
 
 // used just to sanitize nested array field names
@@ -269,18 +272,24 @@ export const resetStateFunction = (fieldConfigs: Record<string, FieldConfigMetad
               factory.createStringLiteral(stateName),
             );
 
-        acc.push(
-          setStateExpression(
-            renderedName,
-            isArray && recordOrInitialValues === 'cleanValues'
-              ? factory.createBinaryExpression(
-                  accessExpression,
-                  factory.createToken(SyntaxKind.QuestionQuestionToken),
-                  factory.createArrayLiteralExpression([], false),
-                )
-              : accessExpression,
-          ),
-        );
+        // Initial values should have the correct values and not need a modifier
+        if (dataType === 'AWSJSON' && !isArray && recordOrInitialValues !== 'initialValues') {
+          const awsJSONAccessModifier = stringifyAWSJSONFieldValue(accessExpression);
+          acc.push(setStateExpression(renderedName, awsJSONAccessModifier));
+        } else {
+          acc.push(
+            setStateExpression(
+              renderedName,
+              isArray && recordOrInitialValues === 'cleanValues'
+                ? factory.createBinaryExpression(
+                    accessExpression,
+                    factory.createToken(SyntaxKind.QuestionQuestionToken),
+                    factory.createArrayLiteralExpression([], false),
+                  )
+                : accessExpression,
+            ),
+          );
+        }
         if (isArray) {
           acc.push(
             setStateExpression(
@@ -347,6 +356,33 @@ export const resetStateFunction = (fieldConfigs: Record<string, FieldConfigMetad
   );
 };
 
+/**
+ * Datastore allows JSON strings and normal JSON so check for a string
+ * before stringifying or else the string will return with escaped quotes
+ *
+ * Example output:
+ * typeof cleanValues.metadata === 'string' ? cleanValues.metadata : JSON.stringify(cleanValues.metadata)
+ */
+const stringifyAWSJSONFieldValue = (
+  value: PropertyAccessExpression | ElementAccessExpression,
+): ConditionalExpression => {
+  return factory.createConditionalExpression(
+    factory.createBinaryExpression(
+      factory.createTypeOfExpression(value),
+      factory.createToken(SyntaxKind.EqualsEqualsEqualsToken),
+      factory.createStringLiteral('string'),
+    ),
+    factory.createToken(SyntaxKind.QuestionToken),
+    value,
+    factory.createToken(SyntaxKind.ColonToken),
+    factory.createCallExpression(
+      factory.createPropertyAccessExpression(factory.createIdentifier('JSON'), factory.createIdentifier('stringify')),
+      undefined,
+      [value],
+    ),
+  );
+};
+
 /**
  * const [name, setName] = React.useState({default_expression});
  *