Modify EC2 Get Instances action to allow partial results

packages/blocks/aws/src/lib/actions/ec2/ec2-get-instances-action.ts

@@ -10,6 +10,7 @@
   getAwsAccountsMultiSelectDropdown,
   getCredentialsListFromAuth,
   getEc2Instances,
+  getEc2InstancesWithPartialResults,
   groupARNsByRegion,
   parseArn,
 } from '@openops/common';
@@ -55,6 +56,13 @@
     }),
     ...filterTagsProperties(),
     dryRun: dryRunCheckBox(),
+    allowPartialResults: Property.Checkbox({
+      displayName: 'Allow partial results',
+      description:
+        'When enabled, the step returns { results, failedRegions } and continues when some regions fail.',
+      required: false,
+      defaultValue: false,
+    }),
   },
   async run(context) {
     try {
@@ -65,51 +73,57 @@
         tags,
         condition,
         dryRun,
+        allowPartialResults,
       } = context.propsValue;
       const filters: Filter[] = getFilters(context);
       const credentials = await getCredentialsListFromAuth(
         context.auth,
         accounts['accounts'],
       );
 
-      const promises: any[] = [];
-      if (filterByARNs) {
-        const arns = convertToStringArrayWithValidation(
-          filterProperty['arns'] as unknown as string[],
-          'Invalid input for ARNs: input should be a single string or an array of strings',
-        );
-        const groupedARNs = groupARNsByRegion(arns);
-
-        for (const region in groupedARNs) {
-          const arnsForRegion = groupedARNs[region];
-          const instanceIdFilter = {
-            Name: 'instance-id',
-            Values: arnsForRegion.map((arn) => parseArn(arn).resourceId),
-          };
-          promises.push(
-            ...credentials.map((credentials) =>
-              getEc2Instances(
-                credentials,
-                [region] as [string, ...string[]],
-                dryRun,
-                [...filters, instanceIdFilter],
-              ),
+      const partial = allowPartialResults === true;
+      const batches = buildEc2GetInstancesBatches(
+        filterByARNs,
+        filterProperty,
+        credentials,
+        filters,
+      );
+
+      if (partial) {
+        const partialOutcomes = await Promise.all(
+          batches.map((batch) =>
+            getEc2InstancesWithPartialResults(
+              batch.creds,
+              batch.regions,
+              dryRun,
+              batch.fetchFilters,
             ),
-          );
-        }
-      } else {
-        const regions = convertToStringArrayWithValidation(
-          filterProperty['regions'],
-          'Invalid input for regions: input should be a single string or an array of strings',
-        );
-        promises.push(
-          ...credentials.map((credentials) =>
-            getEc2Instances(credentials, regions, dryRun, filters),
           ),
         );
+        let instances = partialOutcomes.flatMap((o) => o.results);
+        const failedRegions = partialOutcomes.flatMap((o) => o.failedRegions);
+
+        if (tags?.length) {
+          instances = instances.filter((instance) =>
+            filterTags((instance.Tags as AwsTag[]) ?? [], tags, condition),
+          );
+        }
+
+        return { results: instances, failedRegions };
       }
 
-      const instances = (await Promise.all(promises)).flat();
+      const instances = (
+        await Promise.all(
+          batches.map((batch) =>
+            getEc2Instances(
+              batch.creds,
+              batch.regions,
+              dryRun,
+              batch.fetchFilters,
+            ),
+          ),
+        )
+      ).flat();
 
       if (tags?.length) {
         return instances.filter((instance) =>
@@ -126,6 +140,58 @@
   },
 });
 
+type Ec2GetInstancesBatch = {
+  creds: any;
+  regions: [string, ...string[]];
+  fetchFilters: Filter[];
+};
+
+function buildEc2GetInstancesBatches(
+  filterByARNs: boolean,
+  filterProperty: Record<string, unknown>,
+  credentials: any[],
+  filters: Filter[],
+): Ec2GetInstancesBatch[] {
+  const batches: Ec2GetInstancesBatch[] = [];
+
+  if (filterByARNs) {
+    const arns = convertToStringArrayWithValidation(
+      filterProperty['arns'] as unknown as string[],
+      'Invalid input for ARNs: input should be a single string or an array of strings',
+    );
+    const groupedARNs = groupARNsByRegion(arns);
+
+    for (const region in groupedARNs) {
+      const arnsForRegion = groupedARNs[region];
+      const instanceIdFilter: Filter = {
+        Name: 'instance-id',
+        Values: arnsForRegion.map((arn) => parseArn(arn).resourceId),
+      };
+      for (const creds of credentials) {
+        batches.push({
+          creds,
+          regions: [region] as [string, ...string[]],
+          fetchFilters: [...filters, instanceIdFilter],
+        });
+      }
+    }
+  } else {
+    const regions = convertToStringArrayWithValidation(
+      filterProperty['regions'],
+      'Invalid input for regions: input should be a single string or an array of strings',
+    );
+    for (const creds of credentials) {
+      batches.push({
+        creds,
+        regions: regions as [string, ...string[]],
+        fetchFilters: filters,
+      });
+    }
+  }
+
+  return batches;
+}
+
 function getFilters(context: any): Filter[] {
   const filters: Filter[] = [];
 

packages/blocks/aws/test/ec2/ec2-get-instances-action.test.ts

@@ -2,6 +2,7 @@ const openopsCommonMock = {
   ...jest.requireActual('@openops/common'),
   getCredentialsListFromAuth: jest.fn(),
   getEc2Instances: jest.fn(),
+  getEc2InstancesWithPartialResults: jest.fn(),
   dryRunCheckBox: jest.fn().mockReturnValue({
     required: false,
     defaultValue: false,
@@ -52,7 +53,7 @@ describe('ec2GetInstancesAction', () => {
   };
 
   test('should create action with input regions property', () => {
-    expect(Object.keys(ec2GetInstancesAction.props).length).toBe(8);
+    expect(Object.keys(ec2GetInstancesAction.props).length).toBe(9);
     expect(ec2GetInstancesAction.props).toMatchObject({
       accounts: {
         required: true,
@@ -79,6 +80,11 @@ describe('ec2GetInstancesAction', () => {
         defaultValue: false,
         type: 'CHECKBOX',
       },
+      allowPartialResults: {
+        required: false,
+        defaultValue: false,
+        type: 'CHECKBOX',
+      },
       filterByARNs: {
         type: 'CHECKBOX',
       },
@@ -514,4 +520,79 @@ describe('ec2GetInstancesAction', () => {
       [{ Name: 'instance-id', Values: ['i-2', 'i-4'] }],
     );
   });
+
+  test('when allowPartialResults, uses partial helper and returns object shape', async () => {
+    openopsCommonMock.getEc2InstancesWithPartialResults.mockResolvedValue({
+      results: [{ instance_id: 'i-1' }],
+      failedRegions: [
+        { region: 'eu-west-1', accountId: '111', error: 'AccessDenied' },
+      ],
+    });
+
+    const context = {
+      ...jest.requireActual('@openops/blocks-framework'),
+      auth: auth,
+      propsValue: {
+        filterProperty: { regions: ['us-east-1', 'eu-west-1'] },
+        dryRun: false,
+        accounts: {},
+        allowPartialResults: true,
+      },
+    };
+
+    const result = (await ec2GetInstancesAction.run(context)) as any;
+
+    expect(result).toEqual({
+      results: [{ instance_id: 'i-1' }],
+      failedRegions: [
+        { region: 'eu-west-1', accountId: '111', error: 'AccessDenied' },
+      ],
+    });
+    expect(
+      openopsCommonMock.getEc2InstancesWithPartialResults,
+    ).toHaveBeenCalledWith(
+      'credentials',
+      ['us-east-1', 'eu-west-1'],
+      false,
+      [],
+    );
+    expect(openopsCommonMock.getEc2Instances).not.toHaveBeenCalled();
+  });
+
+  test('when allowPartialResults, aggregates multiple credentials', async () => {
+    openopsCommonMock.getCredentialsListFromAuth.mockResolvedValue([
+      'cred-a',
+      'cred-b',
+    ]);
+    openopsCommonMock.getEc2InstancesWithPartialResults
+      .mockResolvedValueOnce({
+        results: [{ id: 'a' }],
+        failedRegions: [],
+      })
+      .mockResolvedValueOnce({
+        results: [{ id: 'b' }],
+        failedRegions: [{ region: 'us-west-2', accountId: '2', error: 'e' }],
+      });
+
+    const context = {
+      ...jest.requireActual('@openops/blocks-framework'),
+      auth: auth,
+      propsValue: {
+        filterProperty: { regions: ['us-east-1'] },
+        dryRun: true,
+        accounts: {},
+        allowPartialResults: true,
+      },
+    };
+
+    const result = (await ec2GetInstancesAction.run(context)) as any;
+
+    expect(result.results).toEqual([{ id: 'a' }, { id: 'b' }]);
+    expect(result.failedRegions).toEqual([
+      { region: 'us-west-2', accountId: '2', error: 'e' },
+    ]);
+    expect(
+      openopsCommonMock.getEc2InstancesWithPartialResults,
+    ).toHaveBeenCalledTimes(2);
+  });
 });