{POSTGRES} Split postgres flexible server commands from RDBMS

src/azure-cli/azure/cli/command_modules/postgresql/__init__.py

@@ -0,0 +1,53 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core import AzCommandsLoader
+from azure.cli.core import ModExtensionSuppress
+from azure.cli.core.commands import CliCommandType
+from azure.cli.core.profiles import ResourceType
+from azure.cli.command_modules.postgresql._util import PostgreSQLArgumentContext
+from azure.cli.command_modules.postgresql.flexible_server_commands import load_flexibleserver_command_table
+from azure.cli.command_modules.postgresql._params import load_arguments
+
+
+# pylint: disable=import-outside-toplevel
+class RdbmsCommandsLoader(AzCommandsLoader):
+
+    def __init__(self, cli_ctx=None):
+
+        rdbms_custom = CliCommandType(operations_tmpl='azure.cli.command_modules.postgresql.custom#{}')
+        super().__init__(
+            cli_ctx=cli_ctx,
+            resource_type=ResourceType.MGMT_RDBMS,
+            custom_command_type=rdbms_custom,
+            argument_context_cls=PostgreSQLArgumentContext,
+            suppress_extension=ModExtensionSuppress(
+                __name__,
+                'rdbms-vnet',
+                '10.0.1',
+                reason='These commands are now in the CLI.',
+                recommend_remove=True))
+
+    def load_command_table(self, args):
+        from azure.cli.core.aaz import load_aaz_command_table
+        try:
+            from . import aaz
+        except ImportError:
+            aaz = None
+        if aaz:
+            load_aaz_command_table(
+                loader=self,
+                aaz_pkg_name=aaz.__name__,
+                args=args
+            )
+
+        load_flexibleserver_command_table(self, args)
+        return self.command_table
+
+    def load_arguments(self, command):
+        load_arguments(self, command)
+
+
+COMMAND_LOADER_CLS = RdbmsCommandsLoader

src/azure-cli/azure/cli/command_modules/postgresql/_breaking_change.py

@@ -0,0 +1,14 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.breaking_change import register_argument_deprecate, register_command_group_deprecate
+
+register_argument_deprecate('postgres flexible-server create', '--high-availability', redirect='--zonal-resiliency')
+register_argument_deprecate('postgres flexible-server update', '--high-availability', redirect='--zonal-resiliency')
+register_command_group_deprecate(command_group='postgres flexible-server index-tuning',
+                                 redirect='postgres flexible-server autonomous-tuning',
+                                 message='Index tuning feature has now expanded its capabilities to support '
+                                 'other automatically generated recommendations which are covered by the '
+                                 'new command.')

src/azure-cli/azure/cli/command_modules/postgresql/_client_factory.py

@@ -0,0 +1,136 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+from azure.cli.core.commands.client_factory import get_mgmt_service_client
+from azure.cli.core.profiles import ResourceType
+from azure.cli.core.auth.identity import get_environment_credential, AZURE_CLIENT_ID
+
+# pylint: disable=import-outside-toplevel
+
+RM_URI_OVERRIDE = 'AZURE_CLI_POSTGRESQL_FLEXIBLE_RM_URI'
+SUB_ID_OVERRIDE = 'AZURE_CLI_POSTGRESQL_FLEXIBLE_SUB_ID'
+
+
+def get_postgresql_flexible_management_client(cli_ctx, subscription_id=None, **_):
+    from os import getenv
+    from azure.mgmt.postgresqlflexibleservers import PostgreSQLManagementClient
+    # Allow overriding resource manager URI using environment variable
+    # for testing purposes. Subscription id is also determined by environment
+    # variable.
+    rm_uri_override = getenv(RM_URI_OVERRIDE)
+    subscription = subscription_id if subscription_id is not None else getenv(SUB_ID_OVERRIDE)
+    if rm_uri_override:
+        client_id = getenv(AZURE_CLIENT_ID)
+        if client_id:
+            credentials = get_environment_credential()
+        else:
+            from msrest.authentication import Authentication  # pylint: disable=import-error
+            credentials = Authentication()
+
+        return PostgreSQLManagementClient(
+            subscription_id=subscription,
+            base_url=rm_uri_override,
+            credential=credentials)
+    # Normal production scenario.
+    return get_mgmt_service_client(cli_ctx, PostgreSQLManagementClient, subscription_id=subscription)
+
+
+def cf_postgres_flexible_servers(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).servers
+
+
+def cf_postgres_flexible_firewall_rules(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).firewall_rules
+
+
+def cf_postgres_flexible_virtual_endpoints(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).virtual_endpoints
+
+
+def cf_postgres_flexible_config(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).configurations
+
+
+def cf_postgres_flexible_replica(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).replicas
+
+
+def cf_postgres_flexible_location_capabilities(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).capabilities_by_location
+
+
+def cf_postgres_flexible_server_capabilities(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).capabilities_by_server
+
+
+def cf_postgres_flexible_backups(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).backups_automatic_and_on_demand
+
+
+def cf_postgres_flexible_ltr_backups(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).backups_long_term_retention
+
+
+def cf_postgres_flexible_operations(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).operations
+
+
+def cf_postgres_flexible_admin(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).administrators_microsoft_entra
+
+
+def cf_postgres_flexible_migrations(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).migrations
+
+
+def cf_postgres_flexible_server_threat_protection_settings(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).server_threat_protection_settings
+
+
+def cf_postgres_flexible_advanced_threat_protection_settings(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).advanced_threat_protection_settings
+
+
+def cf_postgres_flexible_server_log_files(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).captured_logs
+
+
+def cf_postgres_check_resource_availability(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).name_availability
+
+
+def cf_postgres_flexible_db(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).databases
+
+
+def cf_postgres_flexible_private_dns_zone_suffix_operations(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).private_dns_zone_suffix
+
+
+def cf_postgres_flexible_private_endpoint_connections(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).private_endpoint_connections
+
+
+def cf_postgres_flexible_private_link_resources(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).private_link_resources
+
+
+def cf_postgres_flexible_tuning_options(cli_ctx, _):
+    return get_postgresql_flexible_management_client(cli_ctx).tuning_options
+
+
+def resource_client_factory(cli_ctx, subscription_id=None):
+    return get_mgmt_service_client(cli_ctx, ResourceType.MGMT_RESOURCE_RESOURCES, subscription_id=subscription_id)
+
+
+def private_dns_client_factory(cli_ctx, subscription_id=None):
+    from azure.mgmt.privatedns import PrivateDnsManagementClient
+    return get_mgmt_service_client(cli_ctx, PrivateDnsManagementClient, subscription_id=subscription_id).private_zones
+
+
+def private_dns_link_client_factory(cli_ctx, subscription_id=None):
+    from azure.mgmt.privatedns import PrivateDnsManagementClient
+    return get_mgmt_service_client(cli_ctx, PrivateDnsManagementClient,
+                                   subscription_id=subscription_id).virtual_network_links

src/azure-cli/azure/cli/command_modules/postgresql/_config_reader.py

@@ -0,0 +1,42 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+import json
+import os
+from azure.cli.core.azclierror import BadRequestError
+
+_config = None
+
+
+def get_config_json():
+    global _config  # pylint:disable=global-statement
+    if _config is not None:
+        return _config
+    script_dir = os.path.dirname(os.path.realpath(__file__))
+    with open(os.path.join(script_dir, "config.json"), "r") as f:
+        try:
+            _config = json.load(f)
+            return _config
+        except ValueError:
+            raise BadRequestError("Invalid json file. Make sure that the json file content is properly formatted.")
+
+
+def get_cloud(cmd):
+    config = get_config_json()
+    return config[cmd.cli_ctx.cloud.name]
+
+
+def get_cloud_cluster(cmd, location, subscription_id):
+    try:
+        cloud = get_cloud(cmd)
+        clusters = cloud[location]
+    except KeyError:
+        clusters = None
+    if clusters is not None:
+        for cluster in clusters:
+            if cloud[cluster] is not None:
+                if subscription_id in cloud[cluster]["subscriptions"]:
+                    return cloud[cluster]
+    return

src/azure-cli/azure/cli/command_modules/postgresql/_flexible_server_location_capabilities_util.py

@@ -0,0 +1,120 @@
+# --------------------------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for license information.
+# --------------------------------------------------------------------------------------------
+
+# pylint: disable=unused-argument, line-too-long, import-outside-toplevel, raise-missing-from
+from azure.cli.core.azclierror import InvalidArgumentValueError
+from azure.core.paging import ItemPaged
+from ._client_factory import cf_postgres_flexible_location_capabilities, cf_postgres_flexible_server_capabilities
+from collections import defaultdict
+
+
+def get_postgres_location_capability_info(cmd, location, is_offer_restriction_check_required=False):
+    list_location_capability_client = cf_postgres_flexible_location_capabilities(cmd.cli_ctx, '_')
+    list_location_capability_result = list_location_capability_client.list(location)
+    return _postgres_parse_list_capability(list_location_capability_result, is_offer_restriction_check_required)
+
+
+def get_postgres_server_capability_info(cmd, resource_group, server_name, is_offer_restriction_check_required=False):
+    list_server_capability_client = cf_postgres_flexible_server_capabilities(cmd.cli_ctx, '_')
+    list_server_capability_result = list_server_capability_client.list(resource_group_name=resource_group, server_name=server_name)
+    return _postgres_parse_list_capability(list_server_capability_result, is_offer_restriction_check_required)
+
+
+def get_performance_tiers_for_storage(storage_edition, storage_size):
+    performance_tiers = []
+    storage_size_mb = None if storage_size is None else storage_size * 1024
+    for storage_info in storage_edition.supported_storage_mb:
+        if storage_size_mb == storage_info.storage_size_mb:
+            for performance_tier in storage_info.supported_iops_tiers:
+                performance_tiers.append(performance_tier.name)
+    return performance_tiers
+
+
+def get_performance_tiers(storage_edition):
+    performance_tiers = []
+    for storage_info in storage_edition.supported_storage_mb:
+        for performance_tier in storage_info.supported_iops_tiers:
+            if performance_tier.name not in performance_tiers:
+                performance_tiers.append(performance_tier.name)
+    return performance_tiers
+
+
+# pylint: disable=too-many-locals
+def _postgres_parse_list_capability(result, is_offer_restriction_check_required=False):
+    result = _get_list_from_paged_response(result)
+
+    if not result:
+        raise InvalidArgumentValueError("No available SKUs in this location")
+
+    supported_features = result[0].supported_features if result[0].supported_features is not None else []
+    offer_restricted = [feature for feature in supported_features if feature.name == "OfferRestricted"]
+    restricted = offer_restricted[0].status if offer_restricted else None
+    zone_redundant = [feature for feature in supported_features if feature.name == "ZoneRedundantHa"]
+    geo_backup = [feature for feature in supported_features if feature.name == "GeoBackup"]
+    autonomous_tuning = [feature for feature in supported_features if feature.name == "IndexTuning"]
+
+    if restricted == "Enabled" and not is_offer_restriction_check_required:
+        raise InvalidArgumentValueError("The location is restricted for provisioning of flexible servers. Please try using another region.")
+
+    if restricted != "Disabled" and not is_offer_restriction_check_required:
+        raise InvalidArgumentValueError("No available SKUs in this location.")
+
+    single_az = zone_redundant[0].status != "Enabled" if zone_redundant else True
+    geo_backup_supported = geo_backup[0].status == "Enabled" if geo_backup else False
+    autonomous_tuning_supported = autonomous_tuning[0].status == "Enabled" if autonomous_tuning else False
+
+    tiers = result[0].supported_server_editions
+    tiers_dict = {}
+    for tier_info in tiers:
+        tier_name = tier_info.name
+        tier_dict = {}
+
+        skus = set()
+        zones = set()
+
+        for sku in tier_info.supported_server_skus:
+            skus.add(sku.name)
+            for zone in sku.supported_zones:
+                zones.add(zone)
+
+        storage_sizes = set()
+        for storage_edition in tier_info.supported_storage_editions:
+            if storage_edition.name == "ManagedDisk":
+                for storage_info in storage_edition.supported_storage_mb:
+                    storage_sizes.add(int(storage_info.storage_size_mb // 1024))
+                tier_dict["storage_edition"] = storage_edition
+            elif storage_edition.name == "ManagedDiskV2" and len(storage_edition.supported_storage_mb) > 0:
+                tier_dict["supported_storageV2_size"] = int(storage_edition.supported_storage_mb[0].storage_size_mb // 1024)
+                tier_dict["supported_storageV2_size_max"] = int(storage_edition.supported_storage_mb[0].maximum_storage_size_mb // 1024)
+                tier_dict["supported_storageV2_iops"] = storage_edition.supported_storage_mb[0].supported_iops
+                tier_dict["supported_storageV2_iops_max"] = storage_edition.supported_storage_mb[0].supported_maximum_iops
+                tier_dict["supported_storageV2_throughput"] = storage_edition.supported_storage_mb[0].supported_throughput
+                tier_dict["supported_storageV2_throughput_max"] = storage_edition.supported_storage_mb[0].supported_maximum_throughput
+
+        tier_dict["skus"] = skus
+        tier_dict["storage_sizes"] = storage_sizes
+        tiers_dict[tier_name] = tier_dict
+
+    versions = set()
+    for version in result[0].supported_server_versions:
+        versions.add(version.name)
+
+    supported_server_versions = defaultdict(list)
+    for version in result[0].supported_server_versions:
+        supported_server_versions[version.name] = version.supported_versions_to_upgrade
+
+    return {
+        'sku_info': tiers_dict,
+        'single_az': single_az,
+        'geo_backup_supported': geo_backup_supported,
+        'zones': zones,
+        'server_versions': versions,
+        'supported_server_versions': supported_server_versions,
+        'autonomous_tuning_supported': autonomous_tuning_supported
+    }
+
+
+def _get_list_from_paged_response(obj_list):
+    return list(obj_list) if isinstance(obj_list, ItemPaged) else obj_list