Use getNewExpressionCount for DynamoDB.DocumentClient (#526)

trivikr · web-flow · commit 00266aa7c94a · 2023-07-31T11:28:38.000-07:00
diff --git a/src/transforms/v2-to-v3/apis/getClientIdNamesFromNewExpr.ts b/src/transforms/v2-to-v3/apis/getClientIdNamesFromNewExpr.ts
@@ -1,7 +1,7 @@
 import { Collection, Identifier, JSCodeshift, NewExpression } from "jscodeshift";
 
-import { DYNAMODB } from "../config";
-import { getClientNewExpression, getDocClientNewExpression } from "../utils";
+import { DOCUMENT_CLIENT, DYNAMODB, DYNAMODB_DOCUMENT_CLIENT } from "../config";
+import { getClientNewExpression } from "../utils";
 
 export interface GetClientIdNamesFromNewExprOptions {
   v2ClientName: string;
@@ -50,7 +50,11 @@ export const getClientIdNamesFromNewExpr = (
       );
       if (v2ClientName === DYNAMODB) {
         namesFromGlobalModule.push(
-          ...getNames(j, source, getDocClientNewExpression({ v2GlobalName }))
+          ...getNames(
+            j,
+            source,
+            getClientNewExpression({ v2GlobalName, v2ClientName: DYNAMODB_DOCUMENT_CLIENT })
+          )
         );
       }
     }
@@ -59,7 +63,11 @@ export const getClientIdNamesFromNewExpr = (
     );
     if (v2ClientName === DYNAMODB) {
       namesFromServiceModule.push(
-        ...getNames(j, source, getDocClientNewExpression({ v2ClientLocalName }))
+        ...getNames(
+          j,
+          source,
+          getClientNewExpression({ v2ClientLocalName: `${v2ClientLocalName}.${DOCUMENT_CLIENT}` })
+        )
       );
     }
   }
diff --git a/src/transforms/v2-to-v3/client-instances/replaceDocClientCreation.ts b/src/transforms/v2-to-v3/client-instances/replaceDocClientCreation.ts
@@ -1,7 +1,7 @@
 import { Collection, JSCodeshift } from "jscodeshift";
 
-import { DYNAMODB } from "../config";
-import { getDocClientNewExpression } from "../utils";
+import { DOCUMENT_CLIENT, DYNAMODB, DYNAMODB_DOCUMENT_CLIENT } from "../config";
+import { getClientNewExpression } from "../utils";
 import { getDynamoDBForDocClient } from "./getDynamoDBForDocClient";
 
 export interface ReplaceDocClientCreationOptions {
@@ -19,7 +19,10 @@ export const replaceDocClientCreation = (
 
   if (v2GlobalName) {
     source
-      .find(j.NewExpression, getDocClientNewExpression({ v2GlobalName }))
+      .find(
+        j.NewExpression,
+        getClientNewExpression({ v2GlobalName, v2ClientName: DYNAMODB_DOCUMENT_CLIENT })
+      )
       .replaceWith((v2DocClientNewExpression) =>
         j.callExpression(
           j.memberExpression(j.identifier("DynamoDBDocument"), j.identifier("from")),
@@ -29,7 +32,10 @@ export const replaceDocClientCreation = (
   }
 
   source
-    .find(j.NewExpression, getDocClientNewExpression({ v2ClientLocalName }))
+    .find(
+      j.NewExpression,
+      getClientNewExpression({ v2ClientLocalName: `${v2ClientLocalName}.${DOCUMENT_CLIENT}` })
+    )
     .replaceWith((v2DocClientNewExpression) =>
       j.callExpression(j.memberExpression(j.identifier("DynamoDBDocument"), j.identifier("from")), [
         getDynamoDBForDocClient(j, v2DocClientNewExpression, { v2ClientLocalName }),
diff --git a/src/transforms/v2-to-v3/client-names/getNamesFromNewExpr.ts b/src/transforms/v2-to-v3/client-names/getNamesFromNewExpr.ts
@@ -1,6 +1,7 @@
 import { Collection, Identifier, JSCodeshift, MemberExpression } from "jscodeshift";
 
-import { getClientNewExpression, getDocClientNewExpression } from "../utils";
+import { DYNAMODB_DOCUMENT_CLIENT } from "../config";
+import { getClientNewExpression } from "../utils";
 
 export const getNamesFromNewExpr = (
   j: JSCodeshift,
@@ -14,7 +15,10 @@ export const getNamesFromNewExpr = (
       (newExpression) => ((newExpression.callee as MemberExpression).property as Identifier).name
     ),
   ...source
-    .find(j.NewExpression, getDocClientNewExpression({ v2GlobalName }))
+    .find(
+      j.NewExpression,
+      getClientNewExpression({ v2GlobalName, v2ClientName: DYNAMODB_DOCUMENT_CLIENT })
+    )
     .nodes()
     .map(
       (newExpression) =>
diff --git a/src/transforms/v2-to-v3/modules/addClientModules.ts b/src/transforms/v2-to-v3/modules/addClientModules.ts
@@ -8,9 +8,9 @@ import {
   isS3GetSignedUrlApiUsed,
   isS3UploadApiUsed,
 } from "../apis";
+import { DOCUMENT_CLIENT, DYNAMODB, DYNAMODB_DOCUMENT_CLIENT } from "../config";
 import { getV3ClientTypesCount } from "../ts-type";
 import { getClientTSTypeRefCount } from "./getClientTSTypeRefCount";
-import { getDocClientNewExpressionCount } from "./getDocClientNewExpressionCount";
 import { getNewExpressionCount } from "./getNewExpressionCount";
 import { hasImportEquals } from "./hasImportEquals";
 import { hasRequire } from "./hasRequire";
@@ -79,12 +79,22 @@ export const addClientModules = (
     }
   }
 
-  const docClientNewExpressionCount = getDocClientNewExpressionCount(j, source, options);
-  if (docClientNewExpressionCount > 0) {
-    addClientNamedModule(j, source, {
+  if (options.v2ClientName === DYNAMODB) {
+    const { v2ClientLocalName } = options;
+    const docClientOptions = {
       ...options,
-      importedName: "DynamoDBDocument",
-      v3ClientPackageName: "@aws-sdk/lib-dynamodb",
-    });
+      v2ClientName: DYNAMODB_DOCUMENT_CLIENT,
+      ...(v2ClientLocalName && {
+        v2ClientLocalName: `${v2ClientLocalName}.${DOCUMENT_CLIENT}`,
+      }),
+    };
+    const docClientNewExpressionCount = getNewExpressionCount(j, source, docClientOptions);
+    if (docClientNewExpressionCount > 0) {
+      addClientNamedModule(j, source, {
+        ...options,
+        importedName: "DynamoDBDocument",
+        v3ClientPackageName: "@aws-sdk/lib-dynamodb",
+      });
+    }
   }
 };
diff --git a/src/transforms/v2-to-v3/modules/getDocClientNewExpressionCount.ts b/src/transforms/v2-to-v3/modules/getDocClientNewExpressionCount.ts
diff --git a/src/transforms/v2-to-v3/modules/getNewExpressionCount.ts b/src/transforms/v2-to-v3/modules/getNewExpressionCount.ts
@@ -1,7 +1,7 @@
 import { Collection, JSCodeshift } from "jscodeshift";
 
+import { DOCUMENT_CLIENT, DYNAMODB, DYNAMODB_DOCUMENT_CLIENT } from "../config";
 import { getClientNewExpression } from "../utils";
-import { getDocClientNewExpressionCount } from "./getDocClientNewExpressionCount";
 import { ClientModulesOptions } from "./types";
 
 export const getNewExpressionCount = (
@@ -26,7 +26,15 @@ export const getNewExpressionCount = (
   );
   newExpressionCount += newExpressionsFromClientLocalName.length;
 
-  newExpressionCount += getDocClientNewExpressionCount(j, source, options);
+  if (v2ClientName === DYNAMODB) {
+    newExpressionCount += getNewExpressionCount(j, source, {
+      ...options,
+      v2ClientName: DYNAMODB_DOCUMENT_CLIENT,
+      ...(v2ClientLocalName && {
+        v2ClientLocalName: `${v2ClientLocalName}.${DOCUMENT_CLIENT}`,
+      }),
+    });
+  }
 
   return newExpressionCount;
 };
diff --git a/src/transforms/v2-to-v3/utils/getClientNewExpression.ts b/src/transforms/v2-to-v3/utils/getClientNewExpression.ts
@@ -24,17 +24,58 @@ export const getClientNewExpression = ({
   }
 
   if (v2GlobalName) {
+    if (v2ClientName) {
+      // Support for DynamoDB.DocumentClient
+      const [clientName, subClientName] = v2ClientName.split(".");
+
+      if (subClientName) {
+        return {
+          type: "NewExpression",
+          callee: {
+            type: "MemberExpression",
+            object: {
+              type: "MemberExpression",
+              object: { type: "Identifier", name: v2GlobalName },
+              property: { type: "Identifier", name: clientName },
+            },
+            property: { type: "Identifier", name: subClientName },
+          },
+        } as NewExpression;
+      }
+
+      return {
+        type: "NewExpression",
+        callee: {
+          object: { type: "Identifier", name: v2GlobalName },
+          property: { type: "Identifier", name: clientName },
+        },
+      } as NewExpression;
+    }
+
     return {
       type: "NewExpression",
       callee: {
         object: { type: "Identifier", name: v2GlobalName },
-        property: { type: "Identifier", ...(v2ClientName && { name: v2ClientName }) },
+        property: { type: "Identifier" },
+      },
+    } as NewExpression;
+  }
+
+  // Support for DynamoDB.DocumentClient
+  const [clientName, subClientName] = v2ClientLocalName!.split(".");
+
+  if (subClientName) {
+    return {
+      type: "NewExpression",
+      callee: {
+        object: { type: "Identifier", name: clientName },
+        property: { type: "Identifier", name: subClientName },
       },
     } as NewExpression;
   }
 
   return {
     type: "NewExpression",
-    callee: { type: "Identifier", name: v2ClientLocalName },
+    callee: { type: "Identifier", name: clientName },
   } as NewExpression;
 };
diff --git a/src/transforms/v2-to-v3/utils/getDocClientNewExpression.ts b/src/transforms/v2-to-v3/utils/getDocClientNewExpression.ts
diff --git a/src/transforms/v2-to-v3/utils/index.ts b/src/transforms/v2-to-v3/utils/index.ts
@@ -1,5 +1,4 @@
 export * from "./getClientDeepImportPath";
 export * from "./getClientNewExpression";
 export * from "./getDefaultLocalName";
-export * from "./getDocClientNewExpression";
 export * from "./isTypeScriptFile";
