fix(graphql-relational-transformer): support non string type in sort …

…key(aws-amplify#2985)
tcydig · Nov 11, 2024 · fcb81bd · fcb81bd
1 parent 048349d
commit fcb81bd
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/.../src/__tests__/__snapshots__/amplify-graphql-has-many-references-transformer.test.ts.snap b/.../src/__tests__/__snapshots__/amplify-graphql-has-many-references-transformer.test.ts.snap
@@ -1367,7 +1367,7 @@ exports[`has many references with partition key + sort key 1`] = `
   },
   \\"expressionValues\\": {
       \\":partitionKey\\": $util.dynamodb.toDynamoDB($partitionKeyValue),
-      \\":sortKey\\": $util.dynamodb.toDynamoDB(\\"\${sortKeyValue0}\\")
+      \\":sortKey\\": $util.dynamodb.toDynamoDB($sortKeyValue0)
   }
 } )
   #set( $args = $util.defaultIfNull($ctx.stash.transformedArgs, $ctx.args) )

diff --git a/packages/amplify-graphql-relational-transformer/src/resolver/ddb-generator.ts b/packages/amplify-graphql-relational-transformer/src/resolver/ddb-generator.ts
@@ -51,14 +51,14 @@ const PARTITION_KEY_VALUE = 'partitionKeyValue';
 export class DDBRelationalResolverGenerator extends RelationalResolverGenerator {
   makeExpression = (keySchema: any[], connectionAttributes: string[]): ObjectNode => {
     if (keySchema[1] && connectionAttributes[1]) {
-      let condensedSortKeyValue;
+      let condensedSortKeyValue = `$${SORT_KEY_VALUE}0`;
 
       if (connectionAttributes.length > 2) {
         const rangeKeyFields = connectionAttributes.slice(1);
 
-        condensedSortKeyValue = rangeKeyFields
+        condensedSortKeyValue = `"${rangeKeyFields
           .map((keyField, idx) => `\${${SORT_KEY_VALUE}${idx}}`)
-          .join(ModelResourceIDs.ModelCompositeKeySeparator());
+          .join(ModelResourceIDs.ModelCompositeKeySeparator())}"`;
       }
 
       return obj({
@@ -69,7 +69,7 @@ export class DDBRelationalResolverGenerator extends RelationalResolverGenerator
         }),
         expressionValues: obj({
           ':partitionKey': ref(`util.dynamodb.toDynamoDB($${PARTITION_KEY_VALUE})`),
-          ':sortKey': ref(`util.dynamodb.toDynamoDB(${condensedSortKeyValue ? `"${condensedSortKeyValue}"` : `$${SORT_KEY_VALUE}0`})`),
+          ':sortKey': ref(`util.dynamodb.toDynamoDB(${condensedSortKeyValue})`),
         }),
       });
     }

diff --git a/packages/amplify-graphql-relational-transformer/src/resolver/ddb-references-generator.ts b/packages/amplify-graphql-relational-transformer/src/resolver/ddb-references-generator.ts
@@ -57,9 +57,10 @@ export class DDBRelationalReferencesResolverGenerator extends DDBRelationalResol
       // #set( $sortKeyValue1 = <pk sk 1>
       // These variables are then used in the query object as
       // ":sortKey": $util.dynamodb.toDynamoDB("${sortKeyValue0}#${sortKeyValue1}"")
-      const condensedSortKeyValue = rangeKeyFields
-        .map((key, idx) => `\${${SORT_KEY_VALUE}${idx}}`)
-        .join(ModelResourceIDs.ModelCompositeKeySeparator());
+      const condensedSortKeyValue =
+        rangeKeyFields.length === 1
+          ? `$${SORT_KEY_VALUE}0`
+          : `"${rangeKeyFields.map((key, idx) => `\${${SORT_KEY_VALUE}${idx}}`).join(ModelResourceIDs.ModelCompositeKeySeparator())}"`;
 
       return obj({
         expression: str('#partitionKey = :partitionKey AND #sortKey = :sortKey'),
@@ -69,7 +70,7 @@ export class DDBRelationalReferencesResolverGenerator extends DDBRelationalResol
         }),
         expressionValues: obj({
           ':partitionKey': ref(`util.dynamodb.toDynamoDB($${PARTITION_KEY_VALUE})`),
-          ':sortKey': ref(`util.dynamodb.toDynamoDB("${condensedSortKeyValue}")`),
+          ':sortKey': ref(`util.dynamodb.toDynamoDB(${condensedSortKeyValue})`),
         }),
       });
     }