Skip to content

Commit eee5ac5

Browse files
MarceloRGoncMarceloRGonc
authored
Merge branch 'mg/OPS-3003-6' into mg/OPS-3003 
2 parents 9ca3b61 + a04499b commit eee5ac5

12 files changed

Lines changed: 733 additions & 140 deletions

File tree

‎packages/blocks/azure/src/index.ts‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import { azureAuth } from '@openops/common';
33
import { BlockCategory } from '@openops/shared';
44
import { advisorAction } from './lib/actions/azure-advisor-action';
55
import { azureCliAction } from './lib/actions/azure-cli-action';
6+
import { azureResourceGraphAction } from './lib/actions/azure-resource-graph-action';
67
import { customAzureApiCallAction } from './lib/actions/custom-azure-api-action';
78

89
export const azure = createBlock({
@@ -12,6 +13,11 @@ export const azure = createBlock({
1213
logoUrl: 'https://static.openops.com/blocks/azure.svg',
1314
authors: [],
1415
categories: [BlockCategory.CLOUD],
15-
actions: [azureCliAction, advisorAction, customAzureApiCallAction],
16+
actions: [
17+
azureCliAction,
18+
advisorAction,
19+
azureResourceGraphAction,
20+
customAzureApiCallAction,
21+
],
1622
triggers: [],
1723
});

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

Lines changed: 213 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,213 @@
1+
import { createAction, Property, Validators } from '@openops/blocks-framework';
2+
import {
3+
azureAuth,
4+
getUseHostSessionProperty,
5+
makeHttpRequest,
6+
} from '@openops/common';
7+
import { AxiosHeaders } from 'axios';
8+
import { getAzureAccessToken } from '../auth/get-azure-access-token';
9+
import { createSubscriptionDynamicProperty } from '../common-properties';
10+
11+
const RESOURCE_GRAPH_API_VERSION = '2024-04-01';
12+
const BATCH_SIZE = 1000;
13+
14+
interface ResourceGraphResponse {
15+
data: Record<string, unknown>[];
16+
$skipToken?: string;
17+
}
18+
19+
interface ResourceGraphRequestBody {
20+
query: string;
21+
options: {
22+
$top: number;
23+
$skipToken?: string;
24+
};
25+
subscriptions?: string[];
26+
}
27+
28+
const buildResourceGraphUrl = (apiVersion?: string): string =>
29+
`https://management.azure.com/providers/Microsoft.ResourceGraph/resources?api-version=${
30+
apiVersion || RESOURCE_GRAPH_API_VERSION
31+
}`;
32+
33+
const buildRequestBody = (
34+
query: string,
35+
batch?: string[],
36+
skipToken?: string,
37+
): ResourceGraphRequestBody => ({
38+
query,
39+
options: {
40+
$top: BATCH_SIZE,
41+
...(skipToken && { $skipToken: skipToken }),
42+
},
43+
...(batch?.length && { subscriptions: batch }),
44+
});
45+
46+
const queryBatch = async (
47+
url: string,
48+
headers: AxiosHeaders,
49+
query: string,
50+
batch?: string[],
51+
hardLimit?: number,
52+
): Promise<Record<string, unknown>[]> => {
53+
const results: Record<string, unknown>[] = [];
54+
let skipToken: string | undefined;
55+
56+
do {
57+
const requestBody = buildRequestBody(query, batch, skipToken);
58+
const response = await makeHttpRequest<ResourceGraphResponse>(
59+
'POST',
60+
url,
61+
headers,
62+
requestBody,
63+
);
64+
65+
if (response.data?.length) {
66+
if (hardLimit) {
67+
const remaining = hardLimit - results.length;
68+
results.push(...response.data.slice(0, remaining));
69+
if (results.length >= hardLimit) {
70+
break;
71+
}
72+
} else {
73+
results.push(...response.data);
74+
}
75+
}
76+
77+
skipToken = response.$skipToken;
78+
} while (skipToken);
79+
80+
return results;
81+
};
82+
83+
export const azureResourceGraphAction = createAction({
84+
auth: azureAuth,
85+
name: 'resource_graph_query',
86+
description:
87+
'Query Azure Resource Graph using KQL to retrieve resources across multiple subscriptions',
88+
displayName: 'Azure Resource Graph Query',
89+
isWriteAction: false,
90+
props: {
91+
query: Property.LongText({
92+
displayName: 'KQL Query',
93+
description: 'The Kusto Query Language (KQL) query to execute.',
94+
required: true,
95+
}),
96+
useHostSession: getUseHostSessionProperty('Azure', 'az login'),
97+
querySubscriptions: createSubscriptionDynamicProperty(
98+
{
99+
displayName: 'Query Subscriptions',
100+
description:
101+
'Select Azure subscriptions to query for Azure Resource Graph.',
102+
required: true,
103+
multiSelect: true,
104+
preselectAll: true,
105+
},
106+
'querySubscriptions',
107+
),
108+
maxResults: Property.Number({
109+
displayName: 'Maximum Results',
110+
description:
111+
'Maximum number of results to return. Leave empty for no limit.',
112+
required: false,
113+
validators: [Validators.minValue(1), Validators.integer],
114+
}),
115+
apiVersion: Property.ShortText({
116+
displayName: 'API Version',
117+
description:
118+
'Azure Resource Graph API version. Leave empty to use the latest stable version (2022-10-01).',
119+
required: false,
120+
defaultValue: RESOURCE_GRAPH_API_VERSION,
121+
}),
122+
},
123+
async run(context) {
124+
const {
125+
useHostSession,
126+
querySubscriptions,
127+
query,
128+
maxResults,
129+
apiVersion,
130+
} = context.propsValue;
131+
132+
const kql = query?.trim();
133+
if (!kql) {
134+
throw new Error('KQL query is required.');
135+
}
136+
137+
const normalizedMax =
138+
typeof maxResults === 'number' ? maxResults : undefined;
139+
140+
const subscriptionList = querySubscriptions as string[] | undefined;
141+
142+
const token = await getAzureAccessToken(
143+
context.auth,
144+
!!useHostSession?.['useHostSessionCheckbox'],
145+
);
146+
147+
const headers = new AxiosHeaders({
148+
Authorization: `Bearer ${token}`,
149+
'Content-Type': 'application/json',
150+
});
151+
152+
const allResults = await getQueryResults(
153+
splitSubscriptionListIntoBatches(subscriptionList),
154+
buildResourceGraphUrl(apiVersion),
155+
headers,
156+
kql,
157+
normalizedMax,
158+
);
159+
160+
return {
161+
totalRecords: allResults.length,
162+
data: allResults,
163+
query: kql,
164+
querySubscriptions: subscriptionList,
165+
};
166+
},
167+
});
168+
169+
function splitSubscriptionListIntoBatches(
170+
subscriptionList?: string[],
171+
): (string[] | undefined)[] {
172+
if (!subscriptionList?.length) {
173+
return [undefined];
174+
}
175+
176+
const batches: string[][] = [];
177+
for (let i = 0; i < subscriptionList.length; i += BATCH_SIZE) {
178+
const batch = subscriptionList.slice(i, i + BATCH_SIZE);
179+
batches.push(batch);
180+
}
181+
182+
return batches;
183+
}
184+
185+
async function getQueryResults(
186+
batches: (string[] | undefined)[],
187+
url: string,
188+
headers: AxiosHeaders,
189+
query: string,
190+
maxResults?: number,
191+
): Promise<Record<string, unknown>[]> {
192+
const allResults: Record<string, unknown>[] = [];
193+
let remaining = maxResults ?? Number.POSITIVE_INFINITY;
194+
195+
for (const batch of batches) {
196+
if (remaining <= 0) {
197+
break;
198+
}
199+
200+
const batchResults = await queryBatch(
201+
url,
202+
headers,
203+
query,
204+
batch,
205+
Number.isFinite(remaining) ? remaining : undefined,
206+
);
207+
208+
allResults.push(...batchResults);
209+
remaining = Math.max(0, remaining - batchResults.length);
210+
}
211+
212+
return allResults;
213+
}

‎packages/blocks/azure/src/lib/actions/custom-azure-api-action.ts‎

Lines changed: 7 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,9 @@
11
import { createCustomApiCallAction } from '@openops/blocks-common';
22
import { Property } from '@openops/blocks-framework';
3-
import {
4-
authenticateUserWithAzure,
5-
azureAuth,
6-
getUseHostSessionProperty,
7-
} from '@openops/common';
8-
import { runCommand } from '../azure-cli';
3+
import { azureAuth, getUseHostSessionProperty } from '@openops/common';
4+
import { getAzureAccessToken } from '../auth/get-azure-access-token';
95
import { getSubscriptionsDropdownForHostSession } from '../common-properties';
106

11-
const getHostAccessToken = async (
12-
auth: unknown,
13-
subscription: string,
14-
): Promise<string> => {
15-
const output = await runCommand(
16-
'account get-access-token --resource https://management.azure.com --output json',
17-
auth,
18-
true,
19-
subscription,
20-
);
21-
const parsed = JSON.parse(output ?? '{}');
22-
const token = parsed?.accessToken;
23-
if (!token) {
24-
throw new Error('Failed to obtain Azure access token');
25-
}
26-
return token as string;
27-
};
28-
297
export const customAzureApiCallAction = createCustomApiCallAction({
308
auth: azureAuth,
319
name: 'custom_azure_api_call',
@@ -65,9 +43,11 @@ export const customAzureApiCallAction = createCustomApiCallAction({
6543
const selectedSubscription =
6644
context.propsValue?.subscriptions?.['subDropdown'];
6745

68-
const token = shouldUseHostCredentials
69-
? await getHostAccessToken(context.auth, selectedSubscription)
70-
: (await authenticateUserWithAzure(context.auth)).access_token;
46+
const token = await getAzureAccessToken(
47+
context.auth,
48+
!!shouldUseHostCredentials,
49+
selectedSubscription,
50+
);
7151

7252
return {
7353
Authorization: `Bearer ${token}`,

‎packages/blocks/azure/src/lib/auth/get-azure-access-token.ts‎

Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
import { authenticateUserWithAzure } from '@openops/common';
2+
import { runCommand } from '../azure-cli';
3+
4+
function parseAzAccessToken(raw?: string | null): string {
5+
const parsed = JSON.parse(raw ?? '{}');
6+
const token = parsed?.accessToken;
7+
if (!token) {
8+
throw new Error('Failed to obtain Azure access token');
9+
}
10+
return token;
11+
}
12+
13+
export const getHostAccessToken = async (
14+
auth: unknown,
15+
subscription?: string,
16+
): Promise<string> => {
17+
const output = await runCommand(
18+
'account get-access-token --resource https://management.azure.com --output json',
19+
auth,
20+
true,
21+
subscription,
22+
);
23+
return parseAzAccessToken(output);
24+
};
25+
26+
export const getAzureAccessToken = async (
27+
auth: unknown,
28+
useHost: boolean,
29+
subscription?: string,
30+
): Promise<string> => {
31+
if (useHost) {
32+
return getHostAccessToken(auth, subscription);
33+
}
34+
const { access_token } = await authenticateUserWithAzure(auth);
35+
return access_token;
36+
};

0 commit comments

Comments
 (0)