Add commands for tuning options

Author: nasc17

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

@@ -414,6 +414,10 @@ 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)
 

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

@@ -51,6 +51,7 @@ def _postgres_parse_list_capability(result):
     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"]
+    index_tuning = [feature for feature in supported_features if feature.name == "IndexTuning"]
 
     if restricted == "Enabled":
         raise InvalidArgumentValueError("The location is restricted for provisioning of flexible servers. Please try using another region.")
@@ -60,6 +61,7 @@ def _postgres_parse_list_capability(result):
 
     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
+    index_tuning_supported = index_tuning[0].status == "Enabled" if index_tuning else False
 
     tiers = result[0].supported_server_editions
     tiers_dict = {}
@@ -102,7 +104,8 @@ def _postgres_parse_list_capability(result):
         'single_az': single_az,
         'geo_backup_supported': geo_backup_supported,
         'zones': zones,
-        'server_versions': versions
+        'server_versions': versions,
+        'index_tuning_supported': index_tuning_supported
     }
 
 

src/azure-cli/azure/cli/command_modules/rdbms/_helptext_pg.py

@@ -1170,3 +1170,32 @@
   - name: Update allowed mirrored databases.
     text: az postgres flexible-server fabric-mirroring update-databases -g testgroup -s testsvr --database-names testdb2 testdb3
 """
+
+helps['postgres flexible-server index-tuning'] = """
+type: group
+short-summary: Index tuning analyzes read queries captured in Query Store and recommends index changes to optimize these queries.
+"""
+
+helps['postgres flexible-server index-tuning update'] = """
+type: command
+short-summary: Update index tuning to be enabled/disabled on PostgreSQL flexible server.
+examples:
+  - name: Update index tuning to be enabled/disabled on PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning update -g testgroup -s testsvr --state Enabled
+"""
+
+helps['postgres flexible-server index-tuning list-tuning-options'] = """
+type: command
+short-summary: Get available tuning options associated with a PostgreSQL flexible server.
+examples:
+  - name: Get tuning options for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning list-tuning-options -g testgroup -s testsvr
+"""
+
+helps['postgres flexible-server index-tuning list-recommendations'] = """
+type: command
+short-summary: Get available tuning index recommendations associated with a PostgreSQL flexible server.
+examples:
+  - name: Get tuning index recommendations for a PostgreSQL flexible server. Filter by selected type.
+    text: az postgres flexible-server index-tuning list-recommendations -g testgroup -s testsvr --recommendation-type CreateIndex
+"""

src/azure-cli/azure/cli/command_modules/rdbms/_params.py

@@ -1034,6 +1034,26 @@ def _flexible_server_params(command_group):
                 c.argument('resource_group_name', arg_type=resource_group_name_type)
                 c.argument('server_name', options_list=['--server-name', '-s'], id_part='name', arg_type=server_name_arg_type, required=False)
 
+        # index tuning
+        if command_group == 'postgres':
+            for scope in ['update', 'list-tuning-options', 'list-recommendations']:
+                argument_context_string = '{} flexible-server index-tuning {}'.format(command_group, scope)
+                with self.argument_context(argument_context_string) as c:
+                    c.argument('server_name', options_list=['--server-name', '-s'], arg_type=server_name_arg_type)
+
+                with self.argument_context('{} flexible-server index-tuning update'.format(command_group)) as c:
+                    c.argument('state',
+                            options_list=['--state'],
+                            required=True,
+                            help='State of index tuning setting.',
+                            arg_type=get_enum_type(['Enabled', 'Disabled']))
+
+                with self.argument_context('{} flexible-server index-tuning list-recommendations'.format(command_group)) as c:
+                    c.argument('recommendation_type',
+                            options_list=['--recommendation-type', '-r'],
+                            help='Retrieve recommendations based on type.',
+                            arg_type=get_enum_type(['CreateIndex', 'DropIndex']))
+
         # GTID
         if command_group == 'mysql':
             with self.argument_context('{} flexible-server gtid reset'.format(command_group)) as c:

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_commands.py

@@ -23,7 +23,8 @@
     cf_postgres_flexible_private_link_resources,
     cf_postgres_flexible_virtual_endpoints,
     cf_postgres_flexible_server_threat_protection_settings,
-    cf_postgres_flexible_server_log_files)
+    cf_postgres_flexible_server_log_files,
+    cf_postgres_flexible_tuning_options)
 
 from ._transformers import (
     table_transform_output,
@@ -125,6 +126,11 @@ def load_flexibleserver_command_table(self, _):
         client_factory=cf_postgres_flexible_private_link_resources
     )
 
+    postgres_flexible_tuning_options_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.postgresqlflexibleservers.operations#TuningOptionsOperations.{}',
+        client_factory=cf_postgres_flexible_tuning_options
+    )
+
     # MERU COMMANDS
     flexible_server_custom_common = CliCommandType(
         operations_tmpl='azure.cli.command_modules.rdbms.flexible_server_custom_common#{}')
@@ -296,3 +302,9 @@ def load_flexibleserver_command_table(self, _):
         g.custom_command('start', 'flexible_server_fabric_mirroring_start')
         g.custom_command('stop', 'flexible_server_fabric_mirroring_stop')
         g.custom_command('update-databases', 'flexible_server_fabric_mirroring_update_databases')
+
+    with self.command_group('postgres flexible-server index-tuning', postgres_flexible_tuning_options_sdk,
+                            client_factory=cf_postgres_flexible_tuning_options) as g:
+        g.custom_command('update', 'index_tuning_update', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_command('list-tuning-options', 'tuning_options_list', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_command('list-recommendations', 'recommendations_list', custom_command_type=flexible_servers_custom_postgres)

src/azure-cli/azure/cli/command_modules/rdbms/flexible_server_custom_postgres.py

@@ -27,7 +27,8 @@
     cf_postgres_flexible_db, cf_postgres_check_resource_availability, cf_postgres_flexible_servers, \
     cf_postgres_check_resource_availability_with_location, \
     cf_postgres_flexible_private_dns_zone_suffix_operations, \
-    cf_postgres_flexible_private_endpoint_connections
+    cf_postgres_flexible_private_endpoint_connections, \
+    cf_postgres_flexible_config
 from ._flexible_server_util import generate_missing_parameters, resolve_poller, \
     generate_password, parse_maintenance_window, get_current_time, build_identity_and_data_encryption, \
     _is_resource_name, get_tenant_id, get_case_insensitive_key_value, get_enum_value_true_false
@@ -1596,6 +1597,54 @@ def _update_parameters(cmd, client, server_name, configuration_name, resource_gr
         client.begin_update(resource_group_name, server_name, configuration_name, parameters), cmd.cli_ctx, 'PostgreSQL Parameter update')
 
 
+def index_tuning_update(cmd, resource_group_name, server_name, state):
+    validate_resource_group(resource_group_name)
+    source = "user-override"
+    config_client = cf_postgres_flexible_config(cmd.cli_ctx, '_')
+
+    if state == 'Enabled':
+        subscription = get_subscription_id(cmd.cli_ctx)
+        postgres_source_client = get_postgresql_flexible_management_client(cmd.cli_ctx, subscription)
+        source_server_object = postgres_source_client.servers.get(resource_group_name, server_name)
+        location = ''.join(source_server_object.location.lower().split())
+        list_location_capability_info = get_postgres_location_capability_info(cmd, location)
+        index_tuning_supported = list_location_capability_info['index_tuning_supported']
+        if not index_tuning_supported:
+            raise CLIError("Index tuning is not supported for the server.")
+
+        configuration_name = "index_tuning.mode"
+        value = "report"
+        _update_parameters(cmd, config_client, server_name, configuration_name, resource_group_name, source, value)
+        configuration_name = "pg_qs.query_capture_mode"
+        value = "all"
+        return _update_parameters(cmd, config_client, server_name, configuration_name, resource_group_name, source, value)
+    else:
+        configuration_name = "index_tuning.mode"
+        value = "off"
+        return _update_parameters(cmd, config_client, server_name, configuration_name, resource_group_name, source, value)
+
+
+def tuning_options_list(client, resource_group_name, server_name):
+    validate_resource_group(resource_group_name)
+
+    return client.get(
+        resource_group_name=resource_group_name,
+        server_name=server_name,
+        tuning_option="index",
+    )
+
+
+def recommendations_list(client, resource_group_name, server_name, recommendation_type=None):
+    validate_resource_group(resource_group_name)
+
+    return client.list_recommendations(
+        resource_group_name=resource_group_name,
+        server_name=server_name,
+        tuning_option="index",
+        recommendation_type=recommendation_type
+    )
+
+
 def _update_private_endpoint_connection_status(cmd, client, resource_group_name, server_name,
                                                private_endpoint_connection_name, is_approved=True, description=None):  # pylint: disable=unused-argument
     validate_resource_group(resource_group_name)