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": "^2.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 (Preview)", 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 { PostgreSQLManagementClient as 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,55 @@
+/*---------------------------------------------------------------------------------------------
+ *  Copyright (c) Microsoft Corporation. All rights reserved.
+ *  Licensed under the MIT License. See License.txt in the project root for license information.
+ *--------------------------------------------------------------------------------------------*/
+
+import { PostgreSQLManagementModels as SingleModels } from "@azure/arm-postgresql";
+import { PostgreSQLManagementModels as FlexibleModels } from "@azure/arm-postgresql-flexible";
+
+export enum PostgresServerType {
+    Flexible,
+    Single
+}
+
+export type PostgresAbstractServer = (SingleModels.Server | FlexibleModels.Server) & { serverType?: PostgresServerType; }
+
+export type PostgresAbstractDatabase = SingleModels.Database | FlexibleModels.Database;
+
+/**
+ * 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;
+    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;
 }

src/postgres/commands/createPostgresServer/steps/PostgresServerCreateStep.ts

@@ -2,13 +2,14 @@
  *  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 { ServerForCreate } from '@azure/arm-postgresql/src/models';
+import { PostgreSQLManagementClient, PostgreSQLManagementModels as SingleModels } from "@azure/arm-postgresql";
+import { PostgreSQLManagementClient as PostgreSQLFlexibleManagementClient, PostgreSQLManagementModels as FlexibleModels } from "@azure/arm-postgresql-flexible";
 import { Progress } from 'vscode';
 import { AzureWizardExecuteStep, callWithMaskHandling, createAzureClient, LocationListStep } from 'vscode-azureextensionui';
 import { ext } from '../../../../extensionVariables';
 import { localize } from '../../../../utils/localize';
 import { nonNullProp } from '../../../../utils/nonNull';
+import { AbstractServerCreate, PostgresServerType } from '../../../abstract/models';
 import { IPostgresServerWizardContext } from '../IPostgresServerWizardContext';
 
 export class PostgresServerCreateStep extends AzureWizardExecuteStep<IPostgresServerWizardContext> {
@@ -24,31 +25,76 @@ export class PostgresServerCreateStep extends AzureWizardExecuteStep<IPostgresSe
 
         return await callWithMaskHandling(
             async () => {
-                const client: PostgreSQLManagementClient = createAzureClient(context, PostgreSQLManagementClient);
+                const serverType = nonNullProp(context, 'serverType');
                 const createMessage: string = localize('creatingPostgresServer', 'Creating PostgreSQL Server "{0}"... It should be ready in several minutes.', context.newServerName);
+
                 ext.outputChannel.appendLog(createMessage);
                 progress.report({ message: createMessage });
-                const options: ServerForCreate = {
+                const options: AbstractServerCreate = {
                     location: locationName,
                     sku: nonNullProp(context, 'sku'),
-                    properties: {
-                        administratorLogin: nonNullProp(context, 'shortUserName'),
-                        administratorLoginPassword: password,
-                        sslEnforcement: "Enabled",
-                        createMode: "Default",
-                        version: "10",
-                        storageProfile: {
-                            storageMB: parseInt(storageMB)
-                        }
-                    },
+                    administratorLogin: nonNullProp(context, 'shortUserName'),
+                    administratorLoginPassword: password,
+                    storageMB: parseInt(storageMB)
                 };
 
-                context.server = await client.servers.create(rgName, newServerName, options);
+                switch (serverType){
+                    case PostgresServerType.Single:
+                        const singleClient = createAzureClient(context, PostgreSQLManagementClient);
+                        context.server = await singleClient.servers.create(rgName, newServerName, this.asSingleParameters(options));
+                        break;
+                    case PostgresServerType.Flexible:
+                        const flexiClient = createAzureClient(context, PostgreSQLFlexibleManagementClient);
+                        context.server = await flexiClient.servers.create(rgName, newServerName, this.asFlexibleParameters(options));
+                        break;
+                }
+                context.server.serverType = serverType;
             },
             password);
     }
 
     public shouldExecute(context: IPostgresServerWizardContext): boolean {
         return !context.server;
     }
+
+
+    private asFlexibleParameters(parameters: AbstractServerCreate) : FlexibleModels.Server {
+        return {
+            location: parameters.location,
+            version: "12",
+            administratorLogin: parameters.administratorLogin,
+            administratorLoginPassword: parameters.administratorLoginPassword,
+            storageProfile: {
+                storageMB: parameters.storageMB
+            },
+            sku: {
+                name: parameters.sku.name,
+                tier: parameters.sku.tier as FlexibleModels.SkuTier
+            },
+        }
+    }
+
+    private asSingleParameters(parameters: AbstractServerCreate) : SingleModels.ServerForCreate {
+        return {
+            location: parameters.location,
+            sku: {
+                name: parameters.sku.name,
+                capacity: parameters.sku.capacity,
+                size: parameters.sku.size,
+                family: parameters.sku.family,
+                tier: parameters.sku.tier as SingleModels.SkuTier
+            },
+            properties: {
+                administratorLogin: parameters.administratorLogin,
+                administratorLoginPassword: parameters.administratorLoginPassword,
+                sslEnforcement: "Enabled",
+                createMode: "Default",
+                version: "11",
+                storageProfile: {
+                    storageMB: parameters.storageMB
+                }
+            }
+        }
+    }
+
 }

src/postgres/commands/createPostgresServer/steps/PostgresServerNameStep.ts

@@ -3,20 +3,21 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { PostgreSQLManagementClient } from '@azure/arm-postgresql';
-import { NameAvailability, NameAvailabilityRequest } from '@azure/arm-postgresql/src/models';
-import { AzureNameStep, createAzureClient, ResourceGroupListStep, resourceGroupNamingRules } from 'vscode-azureextensionui';
+import { AzureNameStep, ResourceGroupListStep, resourceGroupNamingRules } from 'vscode-azureextensionui';
 import { localize } from '../../../../utils/localize';
+import { nonNullProp } from '../../../../utils/nonNull';
+import { AbstractPostgresClient, createAbstractPostgresClient } from '../../../abstract/AbstractPostgresClient';
+import { AbstractNameAvailability, PostgresServerType } from '../../../abstract/models';
 import { IPostgresServerWizardContext } from '../IPostgresServerWizardContext';
 
 export class PostgresServerNameStep extends AzureNameStep<IPostgresServerWizardContext> {
 
     public async prompt(context: IPostgresServerWizardContext): Promise<void> {
-        const client: PostgreSQLManagementClient = createAzureClient(context, PostgreSQLManagementClient);
+        const client = createAbstractPostgresClient(nonNullProp(context, "serverType"), context);
         context.newServerName = (await context.ui.showInputBox({
             placeHolder: localize('serverNamePlaceholder', 'Server name'),
             prompt: localize('enterServerNamePrompt', 'Provide a name for the PostgreSQL Server.'),
-            validateInput: (name: string) => validatePostgresServerName(name, client)
+            validateInput: (name: string) => validatePostgresServerName(name, client, nonNullProp(context, "serverType"))
         })).trim();
         context.valuesToMask.push(context.newServerName);
         context.relatedNameTask = this.generateRelatedName(context, context.newServerName, resourceGroupNamingRules);
@@ -31,7 +32,7 @@ export class PostgresServerNameStep extends AzureNameStep<IPostgresServerWizardC
     }
 }
 
-async function validatePostgresServerName(name: string, client: PostgreSQLManagementClient): Promise<string | undefined> {
+async function validatePostgresServerName(name: string, client: AbstractPostgresClient, serverType: PostgresServerType): Promise<string | undefined> {
     name = name ? name.trim() : '';
 
     const min = 3;
@@ -44,14 +45,13 @@ async function validatePostgresServerName(name: string, client: PostgreSQLManage
     } else if (name.startsWith('-') || name.endsWith('-')) {
         return localize('serverNamePrefixSuffixCheck', 'Server name must not start or end in a hyphen.');
     }
-
-    const availabilityRequest: NameAvailabilityRequest = { name: name, type: "Microsoft.DBforPostgreSQL" };
-    const availability: NameAvailability = (await client.checkNameAvailability.execute(availabilityRequest));
+    const resourceType = serverType === PostgresServerType.Single ? "Microsoft.DBforPostgreSQL" : "Microsoft.DBforPostgreSQL/flexibleServers";
+    const availability: AbstractNameAvailability = await client.checkNameAvailability.execute({name: name, type: resourceType});
 
     if (!availability.nameAvailable) {
-        if (availability.reason === 'AlreadyExists') {
-            return localize('serverNameAvailabilityCheck', 'A server named "{0}" already exists.', name);
-        }
+        return availability.message ?
+            availability.message :
+            localize('serverNameAvailabilityCheck', 'Server name "{0}" is not available.', name);
     }
 
     return undefined;

src/postgres/commands/createPostgresServer/steps/PostgresServerSetCredentialsStep.ts

@@ -3,13 +3,13 @@
  *  Licensed under the MIT License. See License.txt in the project root for license information.
  *--------------------------------------------------------------------------------------------*/
 
-import { Server } from '@azure/arm-postgresql/src/models';
 import * as vscode from 'vscode';
 import { Progress } from 'vscode';
 import { AzureWizardExecuteStep } from 'vscode-azureextensionui';
 import { ext } from '../../../../extensionVariables';
 import { localize } from '../../../../utils/localize';
 import { nonNullProp } from '../../../../utils/nonNull';
+import { PostgresAbstractServer } from '../../../abstract/models';
 import { setPostgresCredentials } from '../../setPostgresCredentials';
 import { IPostgresServerWizardContext } from '../IPostgresServerWizardContext';
 
@@ -25,7 +25,7 @@ export class PostgresServerSetCredentialsStep extends AzureWizardExecuteStep<IPo
         progress.report({ message: setupMessage });
         ext.outputChannel.appendLog(setupMessage);
         const password: string = nonNullProp(context, 'adminPassword');
-        const server: Server = nonNullProp(context, 'server');
+        const server: PostgresAbstractServer = nonNullProp(context, 'server');
 
         await setPostgresCredentials(user, password, nonNullProp(server, 'id'));
         const completedMessage: string = localize('addedCredentialsMessage', 'Successfully setup credentials for server "{0}".', newServerName);