[RDBMS] az postgres flexible-server index-tuning: Support tuning options operations (Azure#30851)

* Add commands for tuning options

* New commands and upates for index-tuning group

* Lint fix

* Testing

* Updates from PR

* check casing

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

@@ -681,7 +681,7 @@
 type: command
 short-summary: Get the parameter for a flexible server."
 examples:
-  - name: Get the parameter for a server.W
+  - name: Get the parameter for a server.
     text: az postgres flexible-server parameter show --resource-group testGroup --server-name servername --name parameterName
 """
 
@@ -1170,3 +1170,56 @@
   - 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 for a PostgreSQL flexible server.
+examples:
+  - name: Update index tuning to be enabled/disabled for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning update -g testgroup -s testsvr --enabled True
+"""
+
+helps['postgres flexible-server index-tuning show'] = """
+type: command
+short-summary: Show state of index tuning for a PostgreSQL flexible server.
+examples:
+  - name: Show state of index tuning for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning show -g testgroup -s testsvr
+"""
+
+helps['postgres flexible-server index-tuning list-settings'] = """
+type: command
+short-summary: Get tuning settings associated for a PostgreSQL flexible server.
+examples:
+  - name: Get tuning settings for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning list-settings -g testgroup -s testsvr
+"""
+
+helps['postgres flexible-server index-tuning show-settings'] = """
+type: command
+short-summary: Get a tuning setting for a PostgreSQL flexible server.
+examples:
+  - name: Get a tuning setting for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning show-settings -g testgroup -s testsvr --name setting-name
+"""
+
+helps['postgres flexible-server index-tuning set-settings'] = """
+type: command
+short-summary: Update a tuning setting for a PostgreSQL flexible server.
+examples:
+  - name: Update a tuning setting for a PostgreSQL flexible server.
+    text: az postgres flexible-server index-tuning set-settings -g testgroup -s testsvr --name setting-name --value setting-value
+"""
+
+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

@@ -23,6 +23,7 @@
 from .randomname.generate import generate_username
 from ._flexible_server_util import get_current_time
 from argcomplete.completers import FilesCompleter
+from ._util import get_index_tuning_settings_map
 
 
 def load_arguments(self, _):    # pylint: disable=too-many-statements, too-many-locals
@@ -1034,6 +1035,37 @@ 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', 'show', 'list-settings', 'show-settings', 'set-settings', '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('index_tuning_enabled',
+                           options_list=['--enabled'],
+                           required=True,
+                           help='Enable or disable index tuning feature.',
+                           arg_type=get_enum_type(['True', 'False']))
+
+            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']))
+
+            for scope in ['show-settings', 'set-settings']:
+                argument_context_string = '{} flexible-server index-tuning {}'.format(command_group, scope)
+                with self.argument_context(argument_context_string) as c:
+                    c.argument('setting_name', options_list=['--name', '-n'], required=True,
+                               arg_type=get_enum_type(get_index_tuning_settings_map().keys()),
+                               help='The name of the tuning setting.')
+
+            with self.argument_context('{} flexible-server index-tuning set-settings'.format(command_group)) as c:
+                c.argument('value', options_list=['--value', '-v'],
+                           help='Value of the tuning setting.')
+
         # 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/_util.py

@@ -103,3 +103,20 @@ def call(*args, **kwargs):
                 sleep(interval_sec)
         return call
     return decorate
+
+
+def get_index_tuning_settings_map():
+    return {
+        'analysis_interval': 'index_tuning.analysis_interval',
+        'max_columns_per_index': 'index_tuning.max_columns_per_index',
+        'max_index_count': 'index_tuning.max_index_count',
+        'max_indexes_per_table': 'index_tuning.max_indexes_per_table',
+        'max_queries_per_database': 'index_tuning.max_queries_per_database',
+        'max_regression_factor': 'index_tuning.max_regression_factor',
+        'max_total_size_factor': 'index_tuning.max_total_size_factor',
+        'min_improvement_factor': 'index_tuning.min_improvement_factor',
+        'mode': 'index_tuning.mode',
+        'unused_dml_per_table': 'index_tuning.unused_dml_per_table',
+        'unused_min_period': 'index_tuning.unused_min_period',
+        'unused_reads_per_table': 'index_tuning.unused_reads_per_table'
+    }

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

@@ -296,3 +296,12 @@ 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_config_sdk,
+                            client_factory=cf_postgres_flexible_config) as g:
+        g.custom_command('update', 'index_tuning_update', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_show_command('show', 'index_tuning_show', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_command('list-settings', 'index_tuning_settings_list', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_command('show-settings', 'index_tuning_settings_get', custom_command_type=flexible_servers_custom_postgres)
+        g.custom_command('set-settings', 'index_tuning_settings_set', 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,11 +27,13 @@
     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_tuning_options
 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
 from ._flexible_server_location_capabilities_util import get_postgres_location_capability_info
+from ._util import get_index_tuning_settings_map
 from .flexible_server_custom_common import create_firewall_rule
 from .flexible_server_virtual_network import prepare_private_network, prepare_private_dns_zone, prepare_public_network
 from .validators import pg_arguments_validator, validate_server_name, validate_and_format_restore_point_in_time, \
@@ -1606,6 +1608,91 @@ 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, client, resource_group_name, server_name, index_tuning_enabled):
+    validate_resource_group(resource_group_name)
+    source = "user-override"
+
+    if index_tuning_enabled == "True":
+        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.")
+
+        logger.warning("Enabling index tuning for the server.")
+        configuration_name = "index_tuning.mode"
+        value = "report"
+        _update_parameters(cmd, client, server_name, configuration_name, resource_group_name, source, value)
+        configuration_name = "pg_qs.query_capture_mode"
+        query_capture_mode_configuration = client.get(resource_group_name, server_name, configuration_name)
+
+        if query_capture_mode_configuration.value.lower() == "none":
+            value = "all"
+            _update_parameters(cmd, client, server_name, configuration_name, resource_group_name, source, value)
+        logger.warning("Index tuning is enabled for the server.")
+    else:
+        logger.warning("Disabling index tuning for the server.")
+        configuration_name = "index_tuning.mode"
+        value = "off"
+        _update_parameters(cmd, client, server_name, configuration_name, resource_group_name, source, value)
+        logger.warning("Index tuning is disabled for the server.")
+
+
+def index_tuning_show(client, resource_group_name, server_name):
+    validate_resource_group(resource_group_name)
+    index_tuning_configuration = client.get(resource_group_name, server_name, "index_tuning.mode")
+    query_capture_mode_configuration = client.get(resource_group_name, server_name, "pg_qs.query_capture_mode")
+
+    if index_tuning_configuration.value.lower() == "report" and query_capture_mode_configuration.value.lower() != "none":
+        logger.warning("Index tuning is enabled for the server.")
+    else:
+        logger.warning("Index tuning is disabled for the server.")
+
+
+def index_tuning_settings_list(cmd, client, resource_group_name, server_name):
+    validate_resource_group(resource_group_name)
+    index_tuning_configurations_map_values = get_index_tuning_settings_map().values()
+    configurations_list = client.list_by_server(resource_group_name, server_name)
+
+    # Filter the list based on the values in the dictionary
+    index_tuning_settings = [setting for setting in configurations_list if setting.name in index_tuning_configurations_map_values]
+
+    return index_tuning_settings
+
+
+def index_tuning_settings_get(cmd, client, resource_group_name, server_name, setting_name):
+    validate_resource_group(resource_group_name)
+    index_tuning_configurations_map = get_index_tuning_settings_map()
+    index_tuning_configuration_name = index_tuning_configurations_map[setting_name]
+
+    return client.get(
+        resource_group_name=resource_group_name,
+        server_name=server_name,
+        configuration_name=index_tuning_configuration_name)
+
+
+def index_tuning_settings_set(client, resource_group_name, server_name, setting_name, value=None):
+    source = "user-override" if value else None
+    tuning_settings = get_index_tuning_settings_map()
+    configuration_name = tuning_settings[setting_name]
+    return flexible_parameter_update(client, server_name, configuration_name, resource_group_name, source, value)
+
+
+def recommendations_list(cmd, resource_group_name, server_name, recommendation_type=None):
+    validate_resource_group(resource_group_name)
+    tuning_options_client = cf_postgres_flexible_tuning_options(cmd.cli_ctx, None)
+
+    return tuning_options_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)