From 5d936093900b6815147377470fd8795cabd09d00 Mon Sep 17 00:00:00 2001
From: Megan Mott <59709511+motm32@users.noreply.github.com>
Date: Mon, 18 Sep 2023 08:46:03 -0700
Subject: [PATCH] serviceconnector: Add database subtypes (#1581)

* Add database subtypes

* Change switch case to group

* Requested change
---
 serviceconnector/constants.ts                         |  5 +++++
 .../src/createLinker/CosmosDBAccountListStep.ts       | 10 +++++++++-
 .../src/createLinker/ICreateLinkerContext.ts          |  3 ++-
 serviceconnector/src/createLinker/LinkerCreateStep.ts |  2 +-
 .../src/createLinker/TargetServiceListStep.ts         | 11 ++++++++---
 5 files changed, 25 insertions(+), 6 deletions(-)

diff --git a/serviceconnector/constants.ts b/serviceconnector/constants.ts
index b3ce216baa..3e403c0187 100644
--- a/serviceconnector/constants.ts
+++ b/serviceconnector/constants.ts
@@ -11,6 +11,11 @@ export type TargetServiceType = {
 
 export enum TargetServiceTypeName {
     Storage = 'Storage',
+    MongoDB = 'MongoDB',
+    Cassandra = 'Cassandra',
+    Gremlin = 'Gremlin, Sql',
+    NoSQL = 'Sql',
+    Table = 'Table, Sql',
     CosmosDB = 'Cosmos DB',
     KeyVault = 'Key Vault',
 }
diff --git a/serviceconnector/src/createLinker/CosmosDBAccountListStep.ts b/serviceconnector/src/createLinker/CosmosDBAccountListStep.ts
index 759ba5eaa3..20935d3e07 100644
--- a/serviceconnector/src/createLinker/CosmosDBAccountListStep.ts
+++ b/serviceconnector/src/createLinker/CosmosDBAccountListStep.ts
@@ -16,13 +16,17 @@ export interface DatabaseAccountJsonResponse {
     id: string;
     name: string;
     kind: string;
+    properties: {
+        EnabledApiTypes: string;
+    }
 }
 
 export class CosmosDBAccountListStep extends AzureWizardPromptStep<ICreateLinkerContext>{
     public async prompt(context: ICreateLinkerContext): Promise<void> {
         const placeHolder: string = vscode.l10n.t('Select a database account');
         const accounts = await getCosmosDBDatabaseAccounts(context);
-        context.databaseAccount = (await context.ui.showQuickPick(this.getPicks(accounts.value), { placeHolder })).data;
+        const filteredAccounts = filterDatabaseAccounts(accounts.value, context.targetServiceType?.type);
+        context.databaseAccount = (await context.ui.showQuickPick(this.getPicks(filteredAccounts), { placeHolder })).data;
     }
 
     public shouldPrompt(context: ICreateLinkerContext): boolean {
@@ -40,3 +44,7 @@ async function getCosmosDBDatabaseAccounts(context: ICreateLinkerContext): Promi
     const url = `${context.environment.resourceManagerEndpointUrl}subscriptions/${context.subscriptionId}/providers/Microsoft.DocumentDB/databaseAccounts?api-version=2023-03-15`
     return (<AzExtPipelineResponse>await sendRequestWithTimeout(context, { url, method: 'GET' }, 5000, context)).parsedBody as DatabaseJsonResponse;
 }
+
+function filterDatabaseAccounts(accounts: DatabaseAccountJsonResponse[], kind: string | undefined): DatabaseAccountJsonResponse[] {
+    return accounts.filter(a => a.properties.EnabledApiTypes === kind);
+}
diff --git a/serviceconnector/src/createLinker/ICreateLinkerContext.ts b/serviceconnector/src/createLinker/ICreateLinkerContext.ts
index c8b48a1b28..eef4f9da25 100644
--- a/serviceconnector/src/createLinker/ICreateLinkerContext.ts
+++ b/serviceconnector/src/createLinker/ICreateLinkerContext.ts
@@ -4,7 +4,7 @@
 *--------------------------------------------------------------------------------------------*/
 import { AuthInfoBase, KnownClientType, LinkerResource } from "@azure/arm-servicelinker";
 import { IStorageAccountWizardContext } from "@microsoft/vscode-azext-azureutils";
-import { ExecuteActivityContext, ISubscriptionActionContext } from "@microsoft/vscode-azext-utils";
+import { ExecuteActivityContext, IAzureQuickPickItem, ISubscriptionActionContext } from "@microsoft/vscode-azext-utils";
 import { TargetServiceType } from "../../constants";
 import { DatabaseAccountJsonResponse } from "./CosmosDBAccountListStep";
 import { KeyVaultAccountJsonResponse } from "./KeyVaultListStep";
@@ -20,6 +20,7 @@ export interface ICreateLinkerContext extends ISubscriptionActionContext, IStora
 
     //Target service
     targetServiceType?: TargetServiceType;
+    targetService?: IAzureQuickPickItem<TargetServiceType>;
     databaseAccount?: DatabaseAccountJsonResponse;
     keyVaultAccount?: KeyVaultAccountJsonResponse;
 
diff --git a/serviceconnector/src/createLinker/LinkerCreateStep.ts b/serviceconnector/src/createLinker/LinkerCreateStep.ts
index a88f81d386..d2919bd4b1 100644
--- a/serviceconnector/src/createLinker/LinkerCreateStep.ts
+++ b/serviceconnector/src/createLinker/LinkerCreateStep.ts
@@ -43,7 +43,7 @@ export class LinkerCreateStep extends AzureWizardExecuteStep<ICreateLinkerContex
     }
 
     private getSourceResourceId(context: ICreateLinkerContext): string {
-        switch (context.targetServiceType?.type) {
+        switch (context.targetService?.group) {
             case TargetServiceTypeName.Storage:
                 return nonNullValue(context.storageAccount?.id);
             case TargetServiceTypeName.CosmosDB:
diff --git a/serviceconnector/src/createLinker/TargetServiceListStep.ts b/serviceconnector/src/createLinker/TargetServiceListStep.ts
index cd8d8ff6ac..fd39dc031b 100644
--- a/serviceconnector/src/createLinker/TargetServiceListStep.ts
+++ b/serviceconnector/src/createLinker/TargetServiceListStep.ts
@@ -19,11 +19,16 @@ export class TargetServiceListStep extends AzureWizardPromptStep<ICreateLinkerCo
             { label: vscode.l10n.t('Queue'), data: { name: "storageQueue", type: TargetServiceTypeName.Storage, id: '/queueServices/default' }, group: TargetServiceTypeName.Storage },
             { label: vscode.l10n.t('Table'), data: { name: "storageTable", type: TargetServiceTypeName.Storage, id: '/tableServices/default' }, group: TargetServiceTypeName.Storage },
             { label: vscode.l10n.t('File'), data: { name: "storageFile", type: TargetServiceTypeName.Storage, id: '/fileServices/default' }, group: TargetServiceTypeName.Storage },
-            { label: vscode.l10n.t('CosmosDB'), data: { name: "cosmosDB", type: TargetServiceTypeName.CosmosDB, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
+            { label: vscode.l10n.t('MongoDB'), data: { name: "mongodb", type: TargetServiceTypeName.MongoDB, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
+            { label: vscode.l10n.t('Cassandra'), data: { name: "cassandra", type: TargetServiceTypeName.Cassandra, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
+            { label: vscode.l10n.t('Apache Gremlin'), data: { name: "gremlin", type: TargetServiceTypeName.Gremlin, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
+            { label: vscode.l10n.t('NoSQL'), data: { name: "sql", type: TargetServiceTypeName.NoSQL, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
+            { label: vscode.l10n.t('Table'), data: { name: "table", type: TargetServiceTypeName.Table, id: '/databases' }, group: TargetServiceTypeName.CosmosDB },
             { label: vscode.l10n.t('Key Vault'), data: { name: "keyVault", type: TargetServiceTypeName.KeyVault, id: '/keyVault' }, group: TargetServiceTypeName.KeyVault },
         ];
 
-        context.targetServiceType = (await context.ui.showQuickPick(picks, { placeHolder, enableGrouping: true })).data
+        context.targetService = (await context.ui.showQuickPick(picks, { placeHolder, enableGrouping: true }));
+        context.targetServiceType = context.targetService.data;
     }
 
     public shouldPrompt(context: ICreateLinkerContext): boolean {
@@ -39,7 +44,7 @@ export class TargetServiceListStep extends AzureWizardPromptStep<ICreateLinkerCo
             replication: StorageAccountReplication.LRS
         };
 
-        switch (context.targetServiceType?.type) {
+        switch (context.targetService?.group) {
             case TargetServiceTypeName.Storage:
                 promptSteps.push(new StorageAccountListStep(storageAccountCreateOptions));
                 break;