Abstract the Postgres management client to allow both single and flexible server

package.json

@@ -1100,7 +1100,8 @@
     },
     "dependencies": {
         "@azure/arm-cosmosdb": "^9.1.0",
-        "@azure/arm-postgresql": "^4.4.0",
+        "@azure/arm-postgresql": "^5.0.0",
+        "@azure/arm-postgresql-flexible": "^1.0.0",
         "@azure/cosmos": "^3.6.3",
         "antlr4ts": "^0.4.1-alpha.0",
         "bson": "^1.1.3",

src/AzureDBExperiences.ts

@@ -12,7 +12,8 @@ export enum API {
     Graph = 'Graph',
     Table = 'Table',
     Core = 'Core',
-    Postgres = 'Postgres'
+    PostgresSingle = 'PostgresSingle',
+    PostgresFlexible = 'PostgresFlexible'
 }
 
 export enum DBAccountKind {
@@ -109,8 +110,9 @@ const CoreExperience: Experience = { api: API.Core, longName: "Core", descriptio
 export const MongoExperience: Experience = { api: API.MongoDB, longName: "Azure Cosmos DB for MongoDB API", shortName: "MongoDB", kind: DBAccountKind.MongoDB, tag: "Azure Cosmos DB for MongoDB API" };
 const TableExperience: Experience = { api: API.Table, longName: "Azure Table", shortName: "Table", kind: DBAccountKind.GlobalDocumentDB, capability: 'EnableTable', tag: "Azure Table" };
 const GremlinExperience: Experience = { api: API.Graph, longName: "Gremlin", description: "(graph)", shortName: "Gremlin", kind: DBAccountKind.GlobalDocumentDB, capability: 'EnableGremlin', tag: "Gremlin (graph)" };
-const PostgresExperience: Experience = { api: API.Postgres, longName: "PostgreSQL", shortName: "PostgreSQL" };
+const PostgresSingleExperience: Experience = { api: API.PostgresSingle, longName: "PostgreSQL Single Server", shortName: "PostgreSQLSingle"};
+const PostgresFlexibleExperience: Experience = { api: API.PostgresFlexible, longName: "PostgreSQL Flexible Server", shortName: "PostgreSQLFlexible"};
 
 const cosmosExperiencesArray: Experience[] = [CoreExperience, MongoExperience, TableExperience, GremlinExperience];
-const experiencesArray: Experience[] = [...cosmosExperiencesArray, PostgresExperience];
+const experiencesArray: Experience[] = [...cosmosExperiencesArray, PostgresSingleExperience, PostgresFlexibleExperience];
 const experiencesMap = new Map<API, Experience>(experiencesArray.map((info: Experience): [API, Experience] => [info.api, info]));

src/commands/api/DatabaseAccountTreeItemInternal.ts

@@ -94,7 +94,7 @@ export class DatabaseAccountTreeItemInternal implements DatabaseAccountTreeItem
             if (this._parsedCS instanceof ParsedMongoConnectionString) {
                 apiType = API.MongoDB;
             } else if (this._parsedCS instanceof ParsedPostgresConnectionString) {
-                apiType = API.Postgres;
+                apiType = API.PostgresSingle;
             } else {
                 apiType = API.Core;
             }

src/postgres/abstract/AbstractPostgresClient.ts

@@ -0,0 +1,22 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { PostgreSQLManagementClient } from "@azure/arm-postgresql";
+import { PostgreSQLFlexibleManagementClient } from "@azure/arm-postgresql-flexible";
+import { createAzureClient, ISubscriptionContext } from "vscode-azureextensionui";
+import { PostgresServerType } from "./models";
+
+export type AbstractPostgresClient = PostgreSQLFlexibleManagementClient | PostgreSQLManagementClient;
+
+export function createAbstractPostgresClient(serverType: PostgresServerType, clientInfo: ISubscriptionContext): AbstractPostgresClient {
+    switch (serverType) {
+        case PostgresServerType.Flexible:
+            return createAzureClient(clientInfo, PostgreSQLFlexibleManagementClient);
+        case PostgresServerType.Single:
+            return createAzureClient(clientInfo, PostgreSQLManagementClient);
+        default:
+            throw new Error("Service not implemented.");
+    }
+}

src/postgres/abstract/models.ts

@@ -0,0 +1,60 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import * as FlexibleModels from "@azure/arm-postgresql-flexible/esm/models";
+import * as SingleModels from "@azure/arm-postgresql/esm/models";
+
+export enum PostgresServerType {
+    Flexible,
+    Single
+}
+
+export type PostgresAbstractServer = (SingleModels.Server | FlexibleModels.Server) & { serverType?: PostgresServerType; }
+
+export type PostgresAbstractServerList = Array<PostgresAbstractServer>;
+
+export type PostgresAbstractDatabase = SingleModels.Database | FlexibleModels.Database;
+
+export type PostgresAbstractDatabaseList = Array<PostgresAbstractDatabase>;
+
+/**
+ * Billing information related properties of a server.
+ */
+ export interface AbstractSku {
+    /**
+     * The name of the sku, typically, tier + family + cores, e.g. B_Gen4_1, GP_Gen5_8.
+     */
+    name: string;
+    /**
+     * The tier of the particular SKU, e.g. Basic. Possible values include: 'Basic',
+     * 'GeneralPurpose', 'MemoryOptimized'
+     */
+    tier?: SingleModels.SkuTier | FlexibleModels.SkuTier;
+    /**
+     * The scale up/out capacity, representing server's compute units.
+     */
+    capacity?: number;
+    /**
+     * The size code, to be interpreted by resource as appropriate.
+     */
+    size?: string;
+    /**
+     * The family of hardware.
+     */
+    family?: string;
+}
+
+export interface AbstractServerCreate {
+    location: string;
+    sku: AbstractSku;
+    administratorLogin: string;
+    administratorLoginPassword: string;
+    version: string;
+    storageMB: number;
+}
+
+export type AbstractNameAvailability = SingleModels.NameAvailability | FlexibleModels.NameAvailability;
+
+export type AbstractFirewallRule = SingleModels.FirewallRule | FlexibleModels.FirewallRule;

src/postgres/commands/configurePostgresFirewall.ts

@@ -3,14 +3,14 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { PostgreSQLManagementClient } from '@azure/arm-postgresql';
-import { FirewallRule } from '@azure/arm-postgresql/src/models';
 import * as publicIp from 'public-ip';
 import * as vscode from 'vscode';
-import { createAzureClient, DialogResponses, IActionContext } from "vscode-azureextensionui";
+import { DialogResponses, IActionContext } from "vscode-azureextensionui";
 import { ext } from "../../extensionVariables";
 import { localize } from "../../utils/localize";
 import { nonNullProp } from '../../utils/nonNull';
+import { createAbstractPostgresClient } from '../abstract/AbstractPostgresClient';
+import { AbstractFirewallRule, PostgresServerType } from '../abstract/models';
 import { PostgresServerTreeItem } from "../tree/PostgresServerTreeItem";
 
 export async function configurePostgresFirewall(context: IActionContext, treeItem?: PostgresServerTreeItem): Promise<void> {
@@ -30,12 +30,14 @@ export async function configurePostgresFirewall(context: IActionContext, treeIte
 
 export async function setFirewallRule(context: IActionContext, treeItem: PostgresServerTreeItem, ip: string): Promise<void> {
 
-    const client: PostgreSQLManagementClient = createAzureClient(treeItem.root, PostgreSQLManagementClient);
+    const serverType: PostgresServerType = nonNullProp(treeItem, 'serverType');
+    const client = createAbstractPostgresClient(serverType, treeItem.root);
     const resourceGroup: string = nonNullProp(treeItem, 'resourceGroup');
     const serverName: string = nonNullProp(treeItem, 'azureName');
+
     const firewallRuleName: string = "azureDatabasesForVSCode-publicIp";
 
-    const newFirewallRule: FirewallRule = {
+    const newFirewallRule: AbstractFirewallRule = {
         startIpAddress: ip,
         endIpAddress: ip
     };

src/postgres/commands/createPostgresServer/IPostgresServerWizardContext.ts

@@ -3,8 +3,8 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Server, Sku } from "@azure/arm-postgresql/src/models";
 import { IAzureDBWizardContext } from "../../../tree/IAzureDBWizardContext";
+import { AbstractSku, PostgresAbstractServer, PostgresServerType } from "../../abstract/models";
 
 export interface IPostgresServerWizardContext extends IAzureDBWizardContext {
     /**
@@ -17,6 +17,7 @@ export interface IPostgresServerWizardContext extends IAzureDBWizardContext {
     longUserName?: string;
     adminPassword?: string;
 
-    server?: Server;
-    sku?: Sku;
+    server?: PostgresAbstractServer;
+    sku?: AbstractSku;
+    serverType?: PostgresServerType;
 }