From 6e370e62e9dd01e7effb70f61542044b89c8edfa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gustaf=20R=C3=A4ntil=C3=A4?= <g.rantila@gmail.com>
Date: Fri, 26 Apr 2024 09:04:19 +0200
Subject: [PATCH] fix: Handle reading huge amounts of users without crashing
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Gustaf Räntilä <g.rantila@gmail.com>
---
 .changeset/silent-wombats-hang.md             |  5 ++
 .../src/microsoftGraph/read.test.ts           | 46 +++++++++++++++++++
 .../src/microsoftGraph/read.ts                |  6 +--
 3 files changed, 54 insertions(+), 3 deletions(-)
 create mode 100644 .changeset/silent-wombats-hang.md

diff --git a/.changeset/silent-wombats-hang.md b/.changeset/silent-wombats-hang.md
new file mode 100644
index 0000000000..85bda32007
--- /dev/null
+++ b/.changeset/silent-wombats-hang.md
@@ -0,0 +1,5 @@
+---
+'@backstage/plugin-catalog-backend-module-msgraph': patch
+---
+
+Handle fetching huge amounts of users from Azure without crashing
diff --git a/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.test.ts b/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.test.ts
index a610df2c29..9f4b9ee08c 100644
--- a/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.test.ts
+++ b/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.test.ts
@@ -1064,5 +1064,51 @@ describe('read microsoft graph', () => {
       );
       expect(client.getUserPhotoWithSizeLimit).toHaveBeenCalledTimes(1);
     });
+
+    it('should handle loading huge amounts of users', async () => {
+      client.getOrganization.mockResolvedValue(getExampleOrg());
+
+      const userCount = 200_000;
+
+      async function* getHugeAmountsOfExampleUsers() {
+        for (let i = 0; i < userCount; ++i) {
+          yield {
+            id: `userid-${i}`,
+            displayName: 'User Name',
+            mail: 'user.name@example.com',
+          };
+        }
+      }
+
+      client.getUsers.mockImplementation(getHugeAmountsOfExampleUsers);
+
+      client.getGroups.mockImplementation(getExampleGroups);
+      client.getGroupMembers.mockImplementation(getExampleGroupMembers);
+      client.getGroupPhotoWithSizeLimit.mockResolvedValue(
+        'data:image/jpeg;base64,...',
+      );
+
+      const { users } = await readMicrosoftGraphOrg(client, 'tenantid', {
+        logger: getVoidLogger(),
+        loadUserPhotos: false,
+      });
+
+      expect(users.length).toBe(userCount);
+
+      expect(client.getUsers).toHaveBeenCalledTimes(1);
+      expect(client.getUsers).toHaveBeenCalledWith(
+        {
+          top: 999,
+        },
+        undefined,
+      );
+      expect(client.getGroups).toHaveBeenCalledTimes(1);
+      expect(client.getGroups).toHaveBeenCalledWith(
+        {
+          top: 999,
+        },
+        undefined,
+      );
+    });
   });
 });
diff --git a/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.ts b/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.ts
index 0f3d81c224..2ca21fca9c 100644
--- a/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.ts
+++ b/plugins/catalog-backend-module-msgraph/src/microsoftGraph/read.ts
@@ -384,7 +384,7 @@ export async function readMicrosoftGraphOrg(
     logger: LoggerService;
   },
 ): Promise<{ users: UserEntity[]; groups: GroupEntity[] }> {
-  const users: UserEntity[] = [];
+  let users: UserEntity[] = [];
 
   if (options.userGroupMemberFilter || options.userGroupMemberSearch) {
     const { users: usersInGroups } = await readMicrosoftGraphUsersInGroups(
@@ -401,7 +401,7 @@ export async function readMicrosoftGraphOrg(
         logger: options.logger,
       },
     );
-    users.push(...usersInGroups);
+    users = usersInGroups;
   } else {
     const { users: usersWithFilter } = await readMicrosoftGraphUsers(client, {
       queryMode: options.queryMode,
@@ -412,7 +412,7 @@ export async function readMicrosoftGraphOrg(
       transformer: options.userTransformer,
       logger: options.logger,
     });
-    users.push(...usersWithFilter);
+    users = usersWithFilter;
   }
   const { groups, rootGroup, groupMember, groupMemberOf } =
     await readMicrosoftGraphGroups(client, tenantId, {