Skip to content

Commit c1c1169

Browse files
BenjinBenjin
authored
Improving logging for Azure account sign-in for connection (#19465)
* improving logging for Azure account sign-in * Error handling around stale credentials * another case * revert styling * adding tests
1 parent 6d04381 commit c1c1169

File tree

4 files changed

+190
-25
lines changed

4 files changed

+190
-25
lines changed

‎src/connectionconfig/azureHelpers.ts

Lines changed: 36 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@ import { sendErrorEvent } from "../telemetry/telemetry";
2222
import { getErrorMessage, listAllIterator } from "../utils/utils";
2323
import { MssqlVSCodeAzureSubscriptionProvider } from "../azure/MssqlVSCodeAzureSubscriptionProvider";
2424
import { configSelectedAzureSubscriptions } from "../constants/constants";
25+
import { Logger } from "../models/logger";
2526

2627
//#region VS Code integration
2728

@@ -59,6 +60,7 @@ export async function confirmVscodeAzureSignin(): Promise<
5960
*/
6061
export async function promptForAzureSubscriptionFilter(
6162
state: ConnectionDialogWebviewState,
63+
logger: Logger,
6264
): Promise<boolean> {
6365
try {
6466
const auth = await confirmVscodeAzureSignin();
@@ -90,7 +92,7 @@ export async function promptForAzureSubscriptionFilter(
9092
return true;
9193
} catch (error) {
9294
state.formError = l10n.t("Error loading Azure subscriptions.");
93-
console.error(state.formError + "\n" + getErrorMessage(error));
95+
logger.error(state.formError + "\n" + getErrorMessage(error));
9496
return false;
9597
}
9698
}
@@ -176,6 +178,7 @@ export async function fetchServersFromAzure(sub: AzureSubscription): Promise<Azu
176178

177179
export async function getAccounts(
178180
azureAccountService: AzureAccountService,
181+
logger: Logger,
179182
): Promise<FormItemOptions[]> {
180183
let accounts: IAccount[] = [];
181184
try {
@@ -187,7 +190,7 @@ export async function getAccounts(
187190
};
188191
});
189192
} catch (error) {
190-
console.error(`Error loading Azure accounts: ${getErrorMessage(error)}`);
193+
logger.error(`Error loading Azure accounts: ${getErrorMessage(error)}`);
191194

192195
sendErrorEvent(
193196
TelemetryViews.ConnectionDialog,
@@ -213,27 +216,45 @@ export async function getAccounts(
213216
export async function getTenants(
214217
azureAccountService: AzureAccountService,
215218
accountId: string,
219+
logger: Logger,
216220
): Promise<FormItemOptions[]> {
217221
let tenants: ITenant[] = [];
218222
try {
219223
const account = (await azureAccountService.getAccounts()).find(
220224
(account) => account.displayInfo?.userId === accountId,
221225
);
222-
if (!account) {
226+
227+
if (!account?.properties?.tenants) {
228+
const missingProp = !account
229+
? "account"
230+
: !account.properties
231+
? "properties"
232+
: "tenants";
233+
const message = `Unable to retrieve tenants for the selected account due to undefined ${missingProp}`;
234+
logger.error(message);
235+
236+
sendErrorEvent(
237+
TelemetryViews.ConnectionDialog,
238+
TelemetryActions.LoadAzureTenantsForEntraAuth,
239+
new Error(message),
240+
true, // includeErrorMessage
241+
undefined, // errorCode
242+
`missing_${missingProp}`, // errorType
243+
);
244+
223245
return [];
224246
}
247+
225248
tenants = account.properties.tenants;
226-
if (!tenants) {
227-
return [];
228-
}
249+
229250
return tenants.map((tenant) => {
230251
return {
231252
displayName: tenant.displayName,
232253
value: tenant.id,
233254
};
234255
});
235256
} catch (error) {
236-
console.error(`Error loading Azure tenants: ${getErrorMessage(error)}`);
257+
logger.error(`Error loading Azure tenants: ${getErrorMessage(error)}`);
237258

238259
sendErrorEvent(
239260
TelemetryViews.ConnectionDialog,
@@ -307,9 +328,9 @@ export async function constructAzureAccountForTenant(
307328

308329
//#endregion
309330

310-
//#region Miscellaneous Auzre helpers
331+
//#region Miscellaneous Azure helpers
311332

312-
function extractFromResourceId(resourceId: string, property: string): string | undefined {
333+
export function extractFromResourceId(resourceId: string, property: string): string | undefined {
313334
if (!property.endsWith("/")) {
314335
property += "/";
315336
}
@@ -322,7 +343,12 @@ function extractFromResourceId(resourceId: string, property: string): string | u
322343
startIndex += property.length;
323344
}
324345

325-
return resourceId.substring(startIndex, resourceId.indexOf("/", startIndex));
346+
let endIndex = resourceId.indexOf("/", startIndex);
347+
if (endIndex === -1) {
348+
endIndex = undefined;
349+
}
350+
351+
return resourceId.substring(startIndex, endIndex);
326352
}
327353

328354
//#endregion

‎src/connectionconfig/connectionDialogWebviewController.ts

Lines changed: 38 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
155155
// Load connection form components
156156
this.state.formComponents = await generateConnectionComponents(
157157
this._mainController.connectionManager,
158-
getAccounts(this._mainController.azureAccountService),
158+
getAccounts(this._mainController.azureAccountService, this.logger),
159159
this.getAzureActionButtons(),
160160
);
161161

@@ -299,7 +299,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
299299

300300
this.registerReducer("filterAzureSubscriptions", async (state) => {
301301
try {
302-
if (await promptForAzureSubscriptionFilter(state)) {
302+
if (await promptForAzureSubscriptionFilter(state, this.logger)) {
303303
await this.loadAllAzureServers(state);
304304
}
305305
} catch (err) {
@@ -514,6 +514,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
514514
const tenants = await getTenants(
515515
this._mainController.azureAccountService,
516516
this.state.connectionProfile.accountId,
517+
this.logger,
517518
);
518519
if (tenants.length === 1) {
519520
hiddenProperties.push("tenantId");
@@ -935,6 +936,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
935936

936937
accountsComponent.options = await getAccounts(
937938
this._mainController.azureAccountService,
939+
this.logger,
938940
);
939941

940942
this.state.connectionProfile.accountId = account.key.id;
@@ -956,15 +958,28 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
956958
(account) => account.displayInfo.userId === this.state.connectionProfile.accountId,
957959
);
958960
if (account) {
959-
const session =
960-
await this._mainController.azureAccountService.getAccountSecurityToken(
961-
account,
962-
undefined,
961+
let isTokenExpired = false;
962+
try {
963+
const session =
964+
await this._mainController.azureAccountService.getAccountSecurityToken(
965+
account,
966+
undefined,
967+
);
968+
isTokenExpired = !AzureController.isTokenValid(
969+
session.token,
970+
session.expiresOn,
963971
);
964-
const isTokenExpired = !AzureController.isTokenValid(
965-
session.token,
966-
session.expiresOn,
967-
);
972+
} catch (err) {
973+
this.logger.verbose(
974+
`Error getting token or checking validity; prompting for refresh. Error: ${getErrorMessage(err)}`,
975+
);
976+
977+
this.vscodeWrapper.showErrorMessage(
978+
"Error validating Entra authentication token; you may need to refresh your token.",
979+
);
980+
981+
isTokenExpired = true;
982+
}
968983

969984
if (isTokenExpired) {
970985
actionButtons.push({
@@ -979,12 +994,18 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
979994
this.state.connectionProfile.accountId,
980995
);
981996
if (account) {
982-
const session =
983-
await this._mainController.azureAccountService.getAccountSecurityToken(
984-
account,
985-
undefined,
997+
try {
998+
const session =
999+
await this._mainController.azureAccountService.getAccountSecurityToken(
1000+
account,
1001+
undefined,
1002+
);
1003+
this.logger.log("Token refreshed", session.expiresOn);
1004+
} catch (err) {
1005+
this.logger.error(
1006+
`Error refreshing token: ${getErrorMessage(err)}`,
9861007
);
987-
this.logger.log("Token refreshed", session.expiresOn);
1008+
}
9881009
}
9891010
},
9901011
});
@@ -1016,6 +1037,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
10161037
tenants = await getTenants(
10171038
this._mainController.azureAccountService,
10181039
this.state.connectionProfile.accountId,
1040+
this.logger,
10191041
);
10201042
if (tenantComponent) {
10211043
tenantComponent.options = tenants;
@@ -1040,6 +1062,7 @@ export class ConnectionDialogWebviewController extends FormWebviewController<
10401062
tenants = await getTenants(
10411063
this._mainController.azureAccountService,
10421064
this.state.connectionProfile.accountId,
1065+
this.logger,
10431066
);
10441067
if (tenantComponent) {
10451068
tenantComponent.options = tenants;

‎test/unit/azureHelpers.test.ts

Lines changed: 86 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,86 @@
1+
/*---------------------------------------------------------------------------------------------
2+
* Copyright (c) Microsoft Corporation. All rights reserved.
3+
* Licensed under the MIT License. See License.txt in the project root for license information.
4+
*--------------------------------------------------------------------------------------------*/
5+
6+
import { expect } from "chai";
7+
import { AzureAccountService } from "../../src/services/azureAccountService";
8+
import * as sinon from "sinon";
9+
import * as azureHelpers from "../../src/connectionconfig/azureHelpers";
10+
import { Logger } from "../../src/models/logger";
11+
import { IAccount } from "vscode-mssql";
12+
13+
suite("Azure Helpers", () => {
14+
let sandbox: sinon.SinonSandbox;
15+
let mockAzureAccountService: AzureAccountService;
16+
let mockLogger: Logger;
17+
18+
setup(() => {
19+
sandbox = sinon.createSandbox();
20+
mockAzureAccountService = sandbox.createStubInstance(AzureAccountService);
21+
mockLogger = sandbox.createStubInstance(Logger);
22+
});
23+
24+
teardown(() => {
25+
sandbox.restore();
26+
});
27+
test("getTenants handles error cases", async () => {
28+
const getAccountsStub = mockAzureAccountService.getAccounts as sinon.SinonStub;
29+
// undefined tenants
30+
getAccountsStub.resolves([
31+
{
32+
displayInfo: {
33+
userId: "test-user-id",
34+
},
35+
properties: {
36+
tenants: undefined,
37+
},
38+
} as IAccount,
39+
]);
40+
41+
let result = await azureHelpers.getTenants(
42+
mockAzureAccountService,
43+
"test-user-id",
44+
mockLogger,
45+
);
46+
expect(result).to.be.an("array").that.is.empty;
47+
expect(
48+
(mockLogger.error as sinon.SinonStub).calledWithMatch("undefined tenants"),
49+
"logger should have been called with 'undefined tenants'",
50+
).to.be.true;
51+
52+
// reset mocks for next case
53+
getAccountsStub.reset();
54+
(mockLogger.error as sinon.SinonStub).resetHistory();
55+
56+
// undefined properties
57+
getAccountsStub.resolves([
58+
{
59+
displayInfo: {
60+
userId: "test-user-id",
61+
},
62+
properties: undefined,
63+
} as IAccount,
64+
]);
65+
66+
result = await azureHelpers.getTenants(mockAzureAccountService, "test-user-id", mockLogger);
67+
expect(result).to.be.an("array").that.is.empty;
68+
expect(
69+
(mockLogger.error as sinon.SinonStub).calledWithMatch("undefined properties"),
70+
"logger should have been called with 'undefined properties'",
71+
).to.be.true;
72+
});
73+
74+
test("extractFromResourceId", () => {
75+
const resourceId =
76+
"subscriptions/test-subscription/resourceGroups/test-resource-group/providers/Microsoft.Sql/servers/test-server/databases/test-database";
77+
let result = azureHelpers.extractFromResourceId(resourceId, "servers");
78+
expect(result).to.equal("test-server");
79+
80+
result = azureHelpers.extractFromResourceId(resourceId, "databases");
81+
expect(result).to.equal("test-database");
82+
83+
result = azureHelpers.extractFromResourceId(resourceId, "fakeProperty");
84+
expect(result).to.be.undefined;
85+
});
86+
});

‎test/unit/connectionDialogWebviewController.test.ts

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ import {
3838
import { CreateSessionResponse } from "../../src/models/contracts/objectExplorer/createSessionRequest";
3939
import { TreeNodeInfo } from "../../src/objectExplorer/nodes/treeNodeInfo";
4040
import { mockGetCapabilitiesRequest } from "./mocks";
41+
import { AzureController } from "../../src/azure/azureController";
4142

4243
suite("ConnectionDialogWebviewController Tests", () => {
4344
let sandbox: sinon.SinonSandbox;
@@ -491,4 +492,33 @@ suite("ConnectionDialogWebviewController Tests", () => {
491492
});
492493
});
493494
});
495+
496+
test("getAzureActionButtons", async () => {
497+
controller.state.connectionProfile.authenticationType = AuthenticationType.AzureMFA;
498+
controller.state.connectionProfile.accountId = "TestEntraAccountId";
499+
500+
const actionButtons = await controller["getAzureActionButtons"]();
501+
expect(actionButtons.length).to.equal(1, "Should always have the Sign In button");
502+
expect(actionButtons[0].id).to.equal("azureSignIn");
503+
504+
controller.state.connectionProfile.authenticationType = AuthenticationType.AzureMFA;
505+
controller.state.connectionProfile.accountId = "TestUserId";
506+
507+
const isTokenValidStub = sandbox.stub(AzureController, "isTokenValid").returns(false);
508+
509+
// When there's no error, we should have refreshToken button
510+
let buttons = await controller["getAzureActionButtons"]();
511+
expect(buttons.length).to.equal(2);
512+
expect(buttons[1].id).to.equal("refreshToken");
513+
514+
// Test error handling when getAccountSecurityToken throws
515+
isTokenValidStub.restore();
516+
sandbox
517+
.stub(mainController.azureAccountService, "getAccountSecurityToken")
518+
.throws(new Error("Test error"));
519+
520+
buttons = await controller["getAzureActionButtons"]();
521+
expect(buttons.length).to.equal(2);
522+
expect(buttons[1].id).to.equal("refreshToken");
523+
});
494524
});

0 commit comments

Comments
 (0)