[front] Improve log with multiple memberships, use concurrent execution

front/admin/relocate_users.ts

@@ -1,7 +1,6 @@
-import { removeNulls } from "@dust-tt/types";
-import type { ApiResponse } from "auth0";
+import { concurrentExecutor, removeNulls } from "@dust-tt/types";
 
-import { getAuth0ManagemementClient } from "@app/lib/api/auth0";
+import { getAuth0ManagemementClient, throttleAuth0 } from "@app/lib/api/auth0";
 import { SUPPORTED_REGIONS } from "@app/lib/api/regions/config";
 import { Authenticator } from "@app/lib/auth";
 import { MembershipResource } from "@app/lib/resources/membership_resource";
@@ -34,32 +33,6 @@ makeScript(
     const auth = await Authenticator.internalAdminForWorkspace(workspaceId);
     const workspace = auth.getNonNullableWorkspace();
 
-    let remaining = 10;
-    let resetTime = Date.now();
-
-    const throttleAuth0 = async <T>(fn: () => Promise<ApiResponse<T>>) => {
-      if (remaining < rateLimitThreshold) {
-        const now = Date.now();
-        const waitTime = resetTime * 1000 - now;
-        logger.info({ waitTime }, "Waiting");
-        await new Promise((resolve) => setTimeout(resolve, waitTime));
-      }
-
-      const res = await fn();
-      if (res.status !== 200) {
-        logger.error({ res }, "When calling Auth0");
-        process.exit(1);
-      }
-
-      remaining = Number(res.headers.get("x-ratelimit-remaining"));
-      resetTime = Number(res.headers.get("x-ratelimit-reset"));
-
-      const limit = Number(res.headers.get("x-ratelimit-limit"));
-      logger.info({ limit, remaining, resetTime }, "Rate limit");
-
-      return res.data;
-    };
-
     const members = await MembershipResource.getMembershipsForWorkspace({
       workspace,
     });
@@ -70,8 +43,13 @@ makeScript(
       (m) => m.workspaceId !== workspace.id
     );
     if (externalMemberships.length > 0) {
+      const userIds = externalMemberships.map((m) => m.userId);
+      const users = await UserResource.fetchByModelIds(userIds);
       logger.error(
-        { users: externalMemberships.map((m) => m.user) },
+        {
+          externalMemberships,
+          users: users.map((u) => ({ sId: u.sId, email: u.email })),
+        },
         "Some users have mutiple memberships"
       );
       process.exit(1);
@@ -84,23 +62,31 @@ makeScript(
 
     let count = 0;
 
-    for (const auth0Id of auth0Ids) {
-      count++;
-      logger.info({ user: auth0Id, count }, "Setting region");
-      if (execute) {
-        await throttleAuth0(() =>
-          managementClient.users.update(
-            {
-              id: auth0Id,
-            },
-            {
-              app_metadata: {
-                region: destinationRegion,
-              },
-            }
-          )
-        );
+    await concurrentExecutor(
+      auth0Ids,
+      async (auth0Id) => {
+        count++;
+        logger.info({ user: auth0Id, count }, "Setting region");
+        if (execute) {
+          await throttleAuth0(
+            () =>
+              managementClient.users.update(
+                {
+                  id: auth0Id,
+                },
+                {
+                  app_metadata: {
+                    region: destinationRegion,
+                  },
+                }
+              ),
+            { rateLimitThreshold }
+          );
+        }
+      },
+      {
+        concurrency: 10,
       }
-    }
+    );
   }
 );

front/lib/api/auth0.ts

@@ -1,6 +1,7 @@
 import type { Session } from "@auth0/nextjs-auth0";
 import type { Result } from "@dust-tt/types";
 import { Err, Ok } from "@dust-tt/types";
+import type { ApiResponse } from "auth0";
 import { ManagementClient } from "auth0";
 import { isLeft } from "fp-ts/lib/Either";
 import * as t from "io-ts";
@@ -266,3 +267,33 @@ export async function getAuth0UsersFromEmail(emails: string[]) {
 
   return auth0Users;
 }
+
+const throttleAuth0State = {
+  remaining: -1,
+  resetTime: 0,
+};
+
+export const throttleAuth0 = async <T>(
+  fn: () => Promise<ApiResponse<T>>,
+  { rateLimitThreshold }: { rateLimitThreshold: number }
+) => {
+  const { remaining, resetTime } = throttleAuth0State;
+  if (remaining !== -1 && remaining < rateLimitThreshold) {
+    const now = Date.now();
+    const waitTime = resetTime * 1000 - now;
+    await new Promise((resolve) => setTimeout(resolve, waitTime));
+  }
+
+  const res = await fn();
+  if (res.status !== 200) {
+    logger.error({ res }, "When calling Auth0");
+    process.exit(1);
+  }
+
+  throttleAuth0State.remaining = Number(
+    res.headers.get("x-ratelimit-remaining")
+  );
+  throttleAuth0State.resetTime = Number(res.headers.get("x-ratelimit-reset"));
+
+  return res.data;
+};