Merge pull request #265 from JupiterOne/INT-10529-improv12

use concurrent queue to fetch group users

package.json

@@ -34,6 +34,7 @@
   },
   "dependencies": {
     "@jupiterone/integration-sdk-http-client": "^12.4.0",
+    "bottleneck": "^2.19.5",
     "lodash": "^4.17.21",
     "lodash.startcase": "^4.4.0",
     "p-map": "^4",

src/client.ts

@@ -3,7 +3,6 @@ import {
   IntegrationProviderAuthenticationError,
 } from '@jupiterone/integration-sdk-core';
 import { IntegrationConfig } from './config';
-export type ResourceIteratee<T> = (each: T) => Promise<void> | void;
 import {
   ApplicationGroupAssignment,
   Group,
@@ -24,6 +23,9 @@ import {
   BaseAPIClient,
   RetryOptions,
 } from '@jupiterone/integration-sdk-http-client';
+import Bottleneck from 'bottleneck';
+
+export type ResourceIteratee<T> = (each: T) => Promise<void> | void;
 
 const NINETY_DAYS_AGO = 90 * 24 * 60 * 60 * 1000;
 const DEFAULT_RATE_LIMIT_THRESHOLD = 0.5;
@@ -281,28 +283,78 @@ export class APIClient extends BaseAPIClient {
     } while (nextUrl);
   }
 
+  public async getGroupUsersLimit(
+    groupId: string,
+  ): Promise<number | undefined> {
+    const response = await this.retryableRequest(
+      `/api/v1/groups/${groupId}/users?limit=1`,
+    );
+    await response.text(); // Consume body to avoid memory leaks
+    if (!response.headers.has('x-rate-limit-limit')) {
+      return;
+    }
+    const limitHeader = response.headers.get('x-rate-limit-limit');
+    return parseInt(limitHeader as string, 10);
+  }
+
   /**
    * Iterates each user resource assigned to a given group.
    *
    * @param iteratee receives each resource to produce relationships
    */
-  public async iterateUsersForGroup(
+  public iterateUsersForGroup(
     groupId: string,
     iteratee: ResourceIteratee<OktaUser>,
-  ): Promise<void> {
-    try {
-      for await (const user of this.iteratePages<OktaUser>(
-        `/api/v1/groups/${groupId}/users?limit=1000`,
-      )) {
-        await iteratee(user);
+    limiter: Bottleneck,
+    tasksState: { error: any },
+  ): void {
+    const initialUrl = `/api/v1/groups/${groupId}/users?limit=1000`;
+    const iteratePages = async (url: string) => {
+      if (tasksState.error) {
+        // Stop processing if an error has occurred in previous tasks
+        // This happens when this task has been queued before the error occurred
+        return;
       }
-    } catch (err) {
-      if (err.status === 404) {
-        //ignore it. It's probably a group that got deleted between steps
-      } else {
-        throw err;
+      let nextUrl: string | undefined;
+      try {
+        nextUrl = await this.requestPage(url, iteratee);
+      } catch (err) {
+        if (err.status === 404) {
+          //ignore it. It's probably a group that got deleted between steps
+        } else {
+          err.groupId = groupId;
+          throw err;
+        }
+      }
+      if (nextUrl) {
+        // Queue another task to process the next page
+        void limiter.schedule(() => iteratePages(nextUrl as string));
       }
+    };
+    void limiter.schedule(() => iteratePages(initialUrl));
+  }
+
+  private async requestPage<T>(
+    url: string,
+    iteratee: ResourceIteratee<T>,
+  ): Promise<string | undefined> {
+    const response = await this.retryableRequest(url);
+    const data = await response.json();
+    for (const item of data) {
+      await iteratee(item);
+    }
+
+    const link = response.headers.get('link') as string | undefined;
+    if (!link) {
+      return;
+    }
+
+    const parsedLink = parse(link);
+    if (!parsedLink?.next?.url) {
+      return;
     }
+
+    return parsedLink.next.url;
   }
 
   /**

src/converters/group.test.ts

@@ -1,5 +1,5 @@
 import { Group } from '@okta/okta-sdk-nodejs';
-import { OktaIntegrationConfig, StandardizedOktaUserGroup } from '../types';
+import { OktaIntegrationConfig } from '../types';
 import { createGroupUserRelationship, createUserGroupEntity } from './group';
 
 const config: OktaIntegrationConfig = {
@@ -105,47 +105,14 @@ describe('creating group entity', () => {
 
 describe('creating group entity differently', () => {
   test('with APP_GROUP type', () => {
-    const group: StandardizedOktaUserGroup = {
-      _class: ['UserGroup'],
-      _key: 'id',
-      _rawData: [
-        {
-          name: 'default',
-          rawData: {
-            created: '2019-04-22T21:43:53.000Z',
-            id: 'id',
-            lastMembershipUpdated: '2019-04-22T21:43:53.000Z',
-            lastUpdated: '2019-04-22T21:43:53.000Z',
-            profile: {
-              description: 'description',
-              name: 'name',
-            },
-            type: 'APP_GROUP',
-          },
-        },
-      ],
-      _type: 'okta_app_user_group',
-      created: 1555969433000,
-      createdOn: 1555969433000,
-      displayName: 'name',
-      id: 'id',
-      lastMembershipUpdated: 1555969433000,
-      lastMembershipUpdatedOn: 1555969433000,
-      lastUpdated: 1555969433000,
-      lastUpdatedOn: 1555969433000,
-      description: 'description',
-      name: 'name',
-      type: 'APP_GROUP',
-      webLink: '/admin/group/id',
-    };
-    expect(createGroupUserRelationship(group, 'id')).toEqual({
+    expect(createGroupUserRelationship('group_id', 'id')).toEqual({
       _class: 'HAS',
-      _fromEntityKey: 'id',
-      _key: 'id|has_user|id',
+      _fromEntityKey: 'group_id',
+      _key: 'group_id|has_user|id',
       _toEntityKey: 'id',
       _type: 'okta_group_has_user',
       displayName: 'HAS',
-      groupId: 'id',
+      groupId: 'group_id',
       userId: 'id',
     });
   });

src/converters/group.ts

@@ -3,7 +3,6 @@ import * as url from 'url';
 import {
   createDirectRelationship,
   createIntegrationEntity,
-  Entity,
   parseTimePropertyValue,
   Relationship,
   RelationshipClass,
@@ -74,20 +73,20 @@ export function createUserGroupEntity(
 }
 
 export function createGroupUserRelationship(
-  group: Entity,
+  groupKey: string,
   userKey: string,
 ): Relationship {
   return createDirectRelationship({
     _class: RelationshipClass.HAS,
     fromType: Entities.USER_GROUP._type,
-    fromKey: group._key,
+    fromKey: groupKey,
     toType: Entities.USER._type,
     toKey: userKey,
     properties: {
-      _key: `${group._key}|has_user|${userKey}`,
+      _key: `${groupKey}|has_user|${userKey}`,
       _type: Relationships.USER_GROUP_HAS_USER._type,
       userId: userKey,
-      groupId: group.id as string,
+      groupId: groupKey,
     },
   });
 }

src/steps/constants.ts

@@ -5,6 +5,9 @@ import {
 } from '@jupiterone/integration-sdk-core';
 
 export const DATA_ACCOUNT_ENTITY = 'DATA_ACCOUNT_ENTITY';
+export const USER_GROUP_IDS = 'USER_GROUP_IDS';
+export const APP_USER_GROUP_IDS = 'APP_USER_GROUP_IDS';
+export const EVERYONE_GROUP_KEY = 'EVERYONE_GROUP_KEY';
 
 export const Steps = {
   ACCOUNT: 'fetch-account',