Support for MongoDB clusters / vCore (#22512)

* Initial support for MongoDB clusters / vCore

* Get cluster connection string from arm

* Preserve Azure account for any auth type

When a service has been selected through the Azure
browser, we want to preserve the Azure account
information even if a different authentication has
been selected. This allows doing ARM operations
using the signed in Azure account for any resources
including those with a different login types such
as clusters.

extensions/azurecore/src/azureResource/constants.ts

@@ -20,7 +20,8 @@ export enum AzureResourceItemType {
 	message = 'azure.resource.itemType.message',
 	azureMonitor = 'azure.resource.itemType.azureMonitor',
 	azureMonitorContainer = 'azure.resource.itemType.azureMonitorContainer',
-	cosmosDBMongoAccount = 'azure.resource.itemType.cosmosDBMongoAccount'
+	cosmosDBMongoAccount = 'azure.resource.itemType.cosmosDBMongoAccount',
+	cosmosDBMongoCluster = 'azure.resource.itemType.cosmosDBMongoCluster'
 }
 
 export enum AzureResourceServiceNames {

extensions/azurecore/src/azureResource/interfaces.ts

@@ -50,6 +50,7 @@ export interface DbServerGraphData extends GraphData {
 	properties: {
 		fullyQualifiedDomainName: string;
 		administratorLogin: string;
+		connectionString: string;
 	};
 }
 

extensions/azurecore/src/azureResource/providers/cosmosdb/mongo/cosmosDbMongoService.ts

@@ -10,15 +10,26 @@ import { cosmosMongoDbQuery } from '../../queryStringConstants';
 import { DbServerGraphData } from '../../../interfaces';
 import { COSMOSDB_MONGO_PROVIDER_ID } from '../../../../constants';
 
+export interface AzureResourceMongoDatabaseServer extends azureResource.AzureResourceDatabaseServer {
+	isServer: boolean;
+}
+
 export class CosmosDbMongoService extends ResourceServiceBase<DbServerGraphData> {
 	public override queryFilter: string = cosmosMongoDbQuery;
 
-	public convertServerResource(resource: DbServerGraphData): azureResource.AzureResourceDatabaseServer | undefined {
+	public convertServerResource(resource: DbServerGraphData): AzureResourceMongoDatabaseServer | undefined {
+		let host = resource.name;
+		const isServer = resource.type === azureResource.AzureResourceType.cosmosDbCluster;
+		if (isServer) {
+			const url = new URL(resource.properties.connectionString);
+			host = url.hostname;
+		}
 		return {
 			id: resource.id,
 			name: resource.name,
 			provider: COSMOSDB_MONGO_PROVIDER_ID,
-			fullName: resource.properties.fullyQualifiedDomainName,
+			isServer: isServer,
+			fullName: host,
 			loginName: resource.properties.administratorLogin,
 			defaultDatabaseName: '',
 			tenant: resource.tenantId,

extensions/azurecore/src/azureResource/providers/cosmosdb/mongo/cosmosDbMongoTreeDataProvider.ts

@@ -8,6 +8,7 @@ import * as nls from 'vscode-nls';
 const localize = nls.loadMessageBundle();
 
 import { AzureResourceItemType, AzureResourcePrefixes, cosmosDBProvider } from '../../../constants';
+import { AzureResourceMongoDatabaseServer } from './cosmosDbMongoService';
 import { generateGuid } from '../../../utils';
 import { DbServerGraphData, GraphData } from '../../../interfaces';
 import { ResourceTreeDataProviderBase } from '../../resourceTreeDataProviderBase';
@@ -25,7 +26,7 @@ export class CosmosDbMongoTreeDataProvider extends ResourceTreeDataProviderBase<
 		super(databaseServerService);
 	}
 
-	public getTreeItemForResource(databaseServer: azureResource.AzureResourceDatabaseServer, account: azdata.Account): azdata.TreeItem {
+	public getTreeItemForResource(databaseServer: AzureResourceMongoDatabaseServer, account: azdata.Account): azdata.TreeItem {
 		return {
 			id: `${AzureResourcePrefixes.cosmosdb}${account.key.accountId}${databaseServer.id ?? databaseServer.name}`,
 			label: `${databaseServer.name} (CosmosDB Mongo API)`,
@@ -38,16 +39,18 @@ export class CosmosDbMongoTreeDataProvider extends ResourceTreeDataProviderBase<
 			payload: {
 				id: generateGuid(),
 				connectionName: databaseServer.name,
-				serverName: databaseServer.name,
+				serverName: databaseServer.fullName,
 				userName: databaseServer.loginName,
 				password: '',
-				authenticationType: azdata.connection.AuthenticationType.AzureMFA,
+				authenticationType: databaseServer.isServer ? azdata.connection.AuthenticationType.SqlLogin : azdata.connection.AuthenticationType.AzureMFA,
 				savePassword: true,
 				groupFullName: '',
 				groupId: '',
 				providerName: cosmosDBProvider,
 				saveProfile: false,
-				options: {},
+				options: {
+					isServer: databaseServer.isServer,
+				},
 				azureAccount: account.key.accountId,
 				azureTenantId: databaseServer.tenant,
 				azureResourceId: databaseServer.id,

extensions/azurecore/src/azureResource/providers/queryStringConstants.ts

@@ -71,7 +71,7 @@ export const kustoClusterQuery = `type == "${azureResource.AzureResourceType.kus
 /**
  * Lists all Cosmos DB for MongoDB accounts
  */
-export const cosmosMongoDbQuery = `type == "${azureResource.AzureResourceType.cosmosDbAccount}" and kind == "${Constants.mongoDbKind}"`;
+export const cosmosMongoDbQuery = `(type == "${azureResource.AzureResourceType.cosmosDbAccount}" and kind == "${Constants.mongoDbKind}") or type == "${azureResource.AzureResourceType.cosmosDbCluster}"`;
 
 /**
  * Lists all Log Analytics workspaces

extensions/azurecore/src/azureResource/providers/universal/universalService.ts

@@ -79,7 +79,7 @@ export class AzureResourceUniversalService implements azureResource.IAzureResour
 
 	public getProviderFromResourceType(type: string, kind?: string):
 		[provider: azureResource.IAzureResourceTreeDataProvider, category: ResourceCategory] {
-		if (type === azureResource.AzureResourceType.cosmosDbAccount && kind === mongoDbKind) {
+		if ((type === azureResource.AzureResourceType.cosmosDbAccount && kind === mongoDbKind) || type === azureResource.AzureResourceType.cosmosDbCluster) {
 			return [this.getRegisteredTreeDataProviderInstance(COSMOSDB_MONGO_PROVIDER_ID), ResourceCategory.Server];
 		} else if (type === azureResource.AzureResourceType.sqlDatabase || type === azureResource.AzureResourceType.sqlSynapseSqlDatabase) {
 			return [this.getRegisteredTreeDataProviderInstance(DATABASE_PROVIDER_ID), ResourceCategory.Database];

extensions/azurecore/src/azurecore.d.ts

@@ -383,6 +383,7 @@ declare module 'azurecore' {
 			storageAccount = 'microsoft.storage/storageaccounts',
 			logAnalytics = 'microsoft.operationalinsights/workspaces',
 			cosmosDbAccount = 'microsoft.documentdb/databaseaccounts',
+			cosmosDbCluster = 'microsoft.documentdb/mongoclusters',
 			mysqlFlexibleServer = 'microsoft.dbformysql/flexibleservers'
 		}
 

src/sql/workbench/services/connection/browser/connectionWidget.ts

@@ -974,9 +974,9 @@ export class ConnectionWidget extends lifecycle.Disposable {
 				});
 			}
 
-			if (this.authType === AuthenticationType.AzureMFA || this.authType === AuthenticationType.AzureMFAAndUser) {
+			if (this.authType === AuthenticationType.AzureMFA || this.authType === AuthenticationType.AzureMFAAndUser || connectionInfo.azureAccount !== null) {
 				this.fillInAzureAccountOptions().then(async () => {
-					let accountName = (this.authType === AuthenticationType.AzureMFA)
+					let accountName = ((this.authType === AuthenticationType.AzureMFA) || connectionInfo.azureAccount !== null)
 						? connectionInfo.azureAccount : connectionInfo.userName;
 					let account: azdata.Account;
 					if (accountName) {
@@ -1262,7 +1262,11 @@ export class ConnectionWidget extends lifecycle.Disposable {
 				model.userName = this.userName;
 				model.password = this.password;
 				model.authenticationType = this.authenticationType;
-				model.azureAccount = this.authToken;
+				const azureAccount = this.authToken;
+				if (azureAccount) {
+					// set the azureAccount only if one has been selected, otherwise preserve the initial model value
+					model.azureAccount = azureAccount;
+				}
 				model.savePassword = this._rememberPasswordCheckBox.checked;
 				model.databaseName = this.databaseName;
 				if (this._customOptionWidgets) {