Merge pull request #23 from owenr88/release/1.2.8

1.2.8

Author: owenr88

package-lock.json

@@ -1,7 +1,7 @@
 {
   "name": "cognito-export",
   "description": "A simple CLI tool and JavaScript package for exporting users from an AWS Cognito user pool.",
-  "version": "1.2.7",
+  "version": "1.2.8",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "bin": "dist/cli.js",

packages/cognito-export/package.json

@@ -1,5 +1,6 @@
 import { CognitoIdentityProvider } from "@aws-sdk/client-cognito-identity-provider";
 import chalk from "chalk";
+import { z } from "zod";
 
 export type CognitoExportProps = {
   userPoolId: string;
@@ -11,6 +12,7 @@ class CognitoBase {
   protected verbose: boolean;
   protected userPoolId: string;
   protected cognito: CognitoIdentityProvider | undefined;
+  protected customAttributes: Record<string, z.ZodType<string | number>> = {}
 
   constructor(options: CognitoExportProps) {
     this.verbose = !!options.verbose;
@@ -78,6 +80,16 @@ class CognitoBase {
         );
       }
 
+      // Set the custom attributes
+      userPool.UserPool?.SchemaAttributes?.forEach((attr) => {
+        if (attr.Name && attr.Name.startsWith("custom:")) {
+          this.customAttributes[attr.Name] = attr.AttributeDataType === "String" ? z.string() : z.number();
+        }
+      });
+      if(Object.keys(this.customAttributes).length) {
+        this.log('Extracted custom attributes: ' + Object.keys(this.customAttributes).join(','))
+      }
+
       this.log("Found user pool");
     } catch (e: any) {
       this.log(e.message, "error");

packages/cognito-export/src/CognitoBase.ts

@@ -5,8 +5,7 @@ import { z } from "zod";
 import {
   ImportRecordSchema,
   ImportRecordSchemaType,
-  extractCustomAttributeKeys,
-  parseUserAttributes,
+  parseUserAttributes
 } from "zod-cognito";
 import CognitoBase from "./CognitoBase";
 
@@ -28,21 +27,17 @@ export class CognitoExport extends CognitoBase {
         // Get the users
         const res = await this.cognito?.listUsers(params);
 
+          // Add these to the pool custom attributes
+        const customAttr = {
+          ...this.customAttributes,
+          sub: z.string(),
+        }
+        const customAttributes = z.object(customAttr);
+
         // Loop through the users and parse them into the array
         res?.Users?.forEach((user) => {
           if (users.length >= limit) return false;
 
-          // Convert the customAttributeKeys to an object
-          const customAttributeKeys = extractCustomAttributeKeys(
-            user.Attributes ?? []
-          );
-          const customAttr: Record<string, z.ZodType<string | number>> = {};
-          customAttributeKeys?.forEach((key) => {
-            customAttr[key] = z.union([z.string(), z.number()]);
-          });
-          customAttr.sub = z.string(); // Add this to exported data
-          const customAttributes = z.object(customAttr);
-
           // Map the attributes to a key-value object
           const attrbs = parseUserAttributes<typeof customAttr>(
             user.Attributes,

packages/cognito-export/src/exporter.ts

@@ -1,7 +1,7 @@
 {
   "name": "cognito-import",
   "description": "A simple CLI tool and JavaScript package for importing users into an AWS Cognito user pool.",
-  "version": "1.2.7",
+  "version": "1.2.8",
   "main": "dist/index.js",
   "bin": "dist/cli.js",
   "repository": {
@@ -31,6 +31,7 @@
     "commander": "^12.0.0",
     "node-fetch": "^2.7.0",
     "papaparse": "^5.4.1",
+    "zod": "^3.23.8",
     "zod-cognito": "^1.2.6"
   },
   "devDependencies": {

packages/cognito-import/package.json

@@ -1,6 +1,7 @@
 import { CognitoIdentityProvider } from "@aws-sdk/client-cognito-identity-provider";
 import { IAM } from "@aws-sdk/client-iam";
 import chalk from "chalk";
+import { z } from "zod";
 
 type CognitoExportProps = {
   userPoolId: string;
@@ -14,6 +15,7 @@ class CognitoBase {
   protected iamArn: string | undefined;
   protected userPoolId: string;
   protected cognito: CognitoIdentityProvider | undefined;
+  protected customAttributes: Record<string, z.ZodType<string | number>> = {}
 
   constructor(options: CognitoExportProps) {
     this.verbose = !!options.verbose;
@@ -83,6 +85,16 @@ class CognitoBase {
       }
       this.log("Found user pool " + userPool.UserPool?.Name);
 
+      // Set the custom attributes
+      userPool.UserPool?.SchemaAttributes?.forEach((attr) => {
+        if (attr.Name && attr.Name.startsWith("custom:")) {
+          this.customAttributes[attr.Name] = attr.AttributeDataType === "String" ? z.string() : z.number();
+        }
+      });
+      if(Object.keys(this.customAttributes).length) {
+        this.log('Extracted custom attributes: ' + Object.keys(this.customAttributes).join(','))
+      }
+
       // Create the iamArn if it doesn't exist
       if (!this.iamArn && !!userPool.UserPool?.Arn) {
         // arn:aws:cognito-idp:eu-west-1:000000000000:userpool/eu-west-1_XXXXXXX

packages/cognito-import/src/CognitoBase.ts

@@ -1,7 +1,7 @@
 {
   "name": "zod-cognito",
   "description": "Zod schemas for AWS Cognito User Pools",
-  "version": "1.2.7",
+  "version": "1.2.8",
   "main": "dist/index.js",
   "types": "dist/index.d.ts",
   "repository": {