Add azure resource graph action (#1603)

Author: bigfluffycookie

packages/blocks/azure/src/index.ts

@@ -3,6 +3,7 @@ import { azureAuth } from '@openops/common';
 import { BlockCategory } from '@openops/shared';
 import { advisorAction } from './lib/actions/azure-advisor-action';
 import { azureCliAction } from './lib/actions/azure-cli-action';
+import { azureResourceGraphAction } from './lib/actions/azure-resource-graph-action';
 import { customAzureApiCallAction } from './lib/actions/custom-azure-api-action';
 
 export const azure = createBlock({
@@ -12,6 +13,11 @@ export const azure = createBlock({
   logoUrl: 'https://static.openops.com/blocks/azure.svg',
   authors: [],
   categories: [BlockCategory.CLOUD],
-  actions: [azureCliAction, advisorAction, customAzureApiCallAction],
+  actions: [
+    azureCliAction,
+    advisorAction,
+    azureResourceGraphAction,
+    customAzureApiCallAction,
+  ],
   triggers: [],
 });

packages/blocks/azure/src/lib/actions/azure-resource-graph-action.ts

@@ -0,0 +1,213 @@
+import { createAction, Property, Validators } from '@openops/blocks-framework';
+import {
+  azureAuth,
+  getUseHostSessionProperty,
+  makeHttpRequest,
+} from '@openops/common';
+import { AxiosHeaders } from 'axios';
+import { getAzureAccessToken } from '../auth/get-azure-access-token';
+import { createSubscriptionDynamicProperty } from '../common-properties';
+
+const RESOURCE_GRAPH_API_VERSION = '2024-04-01';
+const BATCH_SIZE = 1000;
+
+interface ResourceGraphResponse {
+  data: Record<string, unknown>[];
+  $skipToken?: string;
+}
+
+interface ResourceGraphRequestBody {
+  query: string;
+  options: {
+    $top: number;
+    $skipToken?: string;
+  };
+  subscriptions?: string[];
+}
+
+const buildResourceGraphUrl = (apiVersion?: string): string =>
+  `https://management.azure.com/providers/Microsoft.ResourceGraph/resources?api-version=${
+    apiVersion || RESOURCE_GRAPH_API_VERSION
+  }`;
+
+const buildRequestBody = (
+  query: string,
+  batch?: string[],
+  skipToken?: string,
+): ResourceGraphRequestBody => ({
+  query,
+  options: {
+    $top: BATCH_SIZE,
+    ...(skipToken && { $skipToken: skipToken }),
+  },
+  ...(batch?.length && { subscriptions: batch }),
+});
+
+const queryBatch = async (
+  url: string,
+  headers: AxiosHeaders,
+  query: string,
+  batch?: string[],
+  hardLimit?: number,
+): Promise<Record<string, unknown>[]> => {
+  const results: Record<string, unknown>[] = [];
+  let skipToken: string | undefined;
+
+  do {
+    const requestBody = buildRequestBody(query, batch, skipToken);
+    const response = await makeHttpRequest<ResourceGraphResponse>(
+      'POST',
+      url,
+      headers,
+      requestBody,
+    );
+
+    if (response.data?.length) {
+      if (hardLimit) {
+        const remaining = hardLimit - results.length;
+        results.push(...response.data.slice(0, remaining));
+        if (results.length >= hardLimit) {
+          break;
+        }
+      } else {
+        results.push(...response.data);
+      }
+    }
+
+    skipToken = response.$skipToken;
+  } while (skipToken);
+
+  return results;
+};
+
+export const azureResourceGraphAction = createAction({
+  auth: azureAuth,
+  name: 'resource_graph_query',
+  description:
+    'Query Azure Resource Graph using KQL to retrieve resources across multiple subscriptions',
+  displayName: 'Azure Resource Graph Query',
+  isWriteAction: false,
+  props: {
+    query: Property.LongText({
+      displayName: 'KQL Query',
+      description: 'The Kusto Query Language (KQL) query to execute.',
+      required: true,
+    }),
+    useHostSession: getUseHostSessionProperty('Azure', 'az login'),
+    querySubscriptions: createSubscriptionDynamicProperty(
+      {
+        displayName: 'Query Subscriptions',
+        description:
+          'Select Azure subscriptions to query for Azure Resource Graph.',
+        required: true,
+        multiSelect: true,
+        preselectAll: true,
+      },
+      'querySubscriptions',
+    ),
+    maxResults: Property.Number({
+      displayName: 'Maximum Results',
+      description:
+        'Maximum number of results to return. Leave empty for no limit.',
+      required: false,
+      validators: [Validators.minValue(1), Validators.integer],
+    }),
+    apiVersion: Property.ShortText({
+      displayName: 'API Version',
+      description:
+        'Azure Resource Graph API version. Leave empty to use the latest stable version (2022-10-01).',
+      required: false,
+      defaultValue: RESOURCE_GRAPH_API_VERSION,
+    }),
+  },
+  async run(context) {
+    const {
+      useHostSession,
+      querySubscriptions,
+      query,
+      maxResults,
+      apiVersion,
+    } = context.propsValue;
+
+    const kql = query?.trim();
+    if (!kql) {
+      throw new Error('KQL query is required.');
+    }
+
+    const normalizedMax =
+      typeof maxResults === 'number' ? maxResults : undefined;
+
+    const subscriptionList = querySubscriptions as string[] | undefined;
+
+    const token = await getAzureAccessToken(
+      context.auth,
+      !!useHostSession?.['useHostSessionCheckbox'],
+    );
+
+    const headers = new AxiosHeaders({
+      Authorization: `Bearer ${token}`,
+      'Content-Type': 'application/json',
+    });
+
+    const allResults = await getQueryResults(
+      splitSubscriptionListIntoBatches(subscriptionList),
+      buildResourceGraphUrl(apiVersion),
+      headers,
+      kql,
+      normalizedMax,
+    );
+
+    return {
+      totalRecords: allResults.length,
+      data: allResults,
+      query: kql,
+      querySubscriptions: subscriptionList,
+    };
+  },
+});
+
+function splitSubscriptionListIntoBatches(
+  subscriptionList?: string[],
+): (string[] | undefined)[] {
+  if (!subscriptionList?.length) {
+    return [undefined];
+  }
+
+  const batches: string[][] = [];
+  for (let i = 0; i < subscriptionList.length; i += BATCH_SIZE) {
+    const batch = subscriptionList.slice(i, i + BATCH_SIZE);
+    batches.push(batch);
+  }
+
+  return batches;
+}
+
+async function getQueryResults(
+  batches: (string[] | undefined)[],
+  url: string,
+  headers: AxiosHeaders,
+  query: string,
+  maxResults?: number,
+): Promise<Record<string, unknown>[]> {
+  const allResults: Record<string, unknown>[] = [];
+  let remaining = maxResults ?? Number.POSITIVE_INFINITY;
+
+  for (const batch of batches) {
+    if (remaining <= 0) {
+      break;
+    }
+
+    const batchResults = await queryBatch(
+      url,
+      headers,
+      query,
+      batch,
+      Number.isFinite(remaining) ? remaining : undefined,
+    );
+
+    allResults.push(...batchResults);
+    remaining = Math.max(0, remaining - batchResults.length);
+  }
+
+  return allResults;
+}