Nasc/index tuning init1112

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

@@ -1126,3 +1126,24 @@
   - name: Get the private link resource for a flexible server using --ids parameter.
     text: az postgres flexible-server private-link-resource show --ids /subscriptions/xxxxxxxx-xxxx-xxxx-xxxx-xxxxxxxxxxxx/resourceGroups/testGroup/providers/Microsoft.DBforPostgreSQL/flexibleServers/testServer
 """
+
+helps['postgres flexible-server update-index-tuning'] = """
+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 update-index-tuning -g testgroup -s testsvr --state Enabled
+"""
+
+helps['postgres flexible-server tuning-options'] = """
+type: group
+short-summary: Get tuning options for a PostgreSQL flexible server.
+"""
+
+helps['postgres flexible-server tuning-options show'] = """
+type: command
+short-summary: Get a tuning option associated with a PostgreSQL flexible server.
+examples:
+  - name: Get tuning option associated with a PostgreSQL flexible server.
+    text: az postgres flexible-server tuning-options show -g testgroup -s testsvr
+"""

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

@@ -987,6 +987,17 @@ 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':
+            with self.argument_context('{} flexible-server update-index-tuning'.format(command_group)) as c:
+                c.argument('server_name', options_list=['--server-name', '-s'], arg_type=server_name_arg_type)
+                c.argument('state', options_list=['--state'], arg_type=get_enum_type(['Enabled', 'Disabled']),
+                           help='Whether index tuning is enabled.')
+                c.argument('yes', arg_type=yes_arg_type)
+
+            with self.argument_context('{} flexible-server tuning-options show'.format(command_group)) as c:
+                c.argument('server_name', options_list=['--server-name', '-s'], arg_type=server_name_arg_type)
+
         # 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,
@@ -61,7 +62,7 @@ def load_flexibleserver_command_table(self, _):
     )
 
     postgres_flexible_location_capabilities_sdk = CliCommandType(
-        operations_tmpl='azure.mgmt.rdbms..postgresql_flexibleservers.operations#LocationBasedCapabilitiesOperations.{}',
+        operations_tmpl='azure.mgmt.postgresqlflexibleservers.operations#LocationBasedCapabilitiesOperations.{}',
         client_factory=cf_postgres_flexible_location_capabilities
     )
 
@@ -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#{}')
@@ -286,3 +292,12 @@ def load_flexibleserver_command_table(self, _):
                             client_factory=cf_postgres_flexible_private_link_resources) as g:
         g.command('list', 'list_by_server')
         g.custom_show_command('show', 'flexible_server_private_link_resource_get', custom_command_type=flexible_servers_custom_postgres)
+
+    with self.command_group('postgres flexible-server', postgres_flexible_config_sdk,
+                            custom_command_type=flexible_servers_custom_postgres,
+                            client_factory=cf_postgres_flexible_config, table_transformer=table_transform_output_parameters) as g:
+        g.custom_command('update-index-tuning', 'index_tuning_update', custom_command_type=flexible_servers_custom_postgres)
+
+    with self.command_group('postgres flexible-server tuning-options', postgres_flexible_tuning_options_sdk,
+                            client_factory=cf_postgres_flexible_tuning_options) as g:
+        g.custom_command('show', 'tuning_options_get', custom_command_type=flexible_servers_custom_postgres)

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

@@ -22,7 +22,7 @@
 from azure.core.exceptions import HttpResponseError, ResourceNotFoundError
 from azure.mgmt.core.tools import resource_id, is_valid_resource_id, parse_resource_id
 from azure.cli.core.azclierror import BadRequestError, FileOperationError, MutuallyExclusiveArgumentError, RequiredArgumentMissingError, ArgumentUsageError, InvalidArgumentValueError
-from azure.mgmt.rdbms import postgresql_flexibleservers
+from azure.mgmt import postgresqlflexibleservers as postgresql_flexibleservers
 from ._client_factory import cf_postgres_flexible_firewall_rules, get_postgresql_flexible_management_client, \
     cf_postgres_flexible_db, cf_postgres_check_resource_availability, cf_postgres_flexible_servers, \
     cf_postgres_check_resource_availability_with_location, \
@@ -484,14 +484,52 @@ def flexible_parameter_update(client, server_name, configuration_name, resource_
         source = "user-override"
 
     parameters = postgresql_flexibleservers.models.Configuration(
-        configuration_name=configuration_name,
         value=value,
         source=source
     )
 
     return client.begin_update(resource_group_name, server_name, configuration_name, parameters)
 
 
+def index_tuning_update(cmd, client, resource_group_name, server_name, state):
+    validate_resource_group(resource_group_name)
+    source = "user-override"
+
+    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"
+        _index_tuning_parameter_update(cmd, client, server_name, configuration_name, resource_group_name, source, value)
+        configuration_name = "pg_qs.query_capture_mode"
+        value = "all"
+        _index_tuning_parameter_update(cmd, client, server_name, configuration_name, resource_group_name, source, value)
+    else:
+        configuration_name = "index_tuning.mode"
+        parameters = postgresql_flexibleservers.models.Configuration(
+            value="off",
+            source=source
+        )
+        return client.begin_update(resource_group_name, server_name, configuration_name, parameters)
+
+
+def _index_tuning_parameter_update(cmd, client, server_name, configuration_name, resource_group_name, source, value):
+    parameters = postgresql_flexibleservers.models.Configuration(
+        value=value,
+        source=source
+    )
+
+    return resolve_poller(
+        client.begin_update(resource_group_name, server_name, configuration_name, parameters), cmd.cli_ctx, 'PostgreSQL Parameter update')
+
+
 def flexible_list_skus(cmd, client, location):
     result = client.execute(location)
     logger.warning('For prices please refer to https://aka.ms/postgres-pricing')
@@ -1401,6 +1439,16 @@ def flexible_server_private_link_resource_get(
         group_name="postgresqlServer")
 
 
+def tuning_options_get(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 _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)

src/azure-cli/azure/cli/command_modules/rdbms/tests/latest/test_rdbms_flexible_commands.py

@@ -1018,7 +1018,6 @@ def main_tests(geo_redundant_backup):
         main_tests(True)
 
 
-
 class FlexibleServerProxyResourceMgmtScenarioTest(ScenarioTest):
 
     postgres_location = 'eastus'
@@ -2541,7 +2540,6 @@ def _test_server_logs_mgmt(self, database_engine, resource_group):
         self.cmd('{} flexible-server delete -g {} -n {} --yes'.format(database_engine, resource_group, server_name))
 
 
-
 class FlexibleServerPrivateEndpointsMgmtScenarioTest(ScenarioTest):
 
     postgres_location = 'eastus'
@@ -2743,3 +2741,46 @@ def _test_private_link_resource(self, database_engine, resource_group, server, g
         result = self.cmd('{} flexible-server private-link-resource show -g {} -s {}'
                           .format(database_engine, resource_group, server)).get_output_in_json()
         self.assertEqual(result['groupId'], group_id)
+
+
+class FlexibleServerTuningOptionsResourceMgmtScenarioTest(ScenarioTest):
+
+    postgres_location = 'eastus2euap'
+
+    @AllowLargeResponse()
+    @ResourceGroupPreparer(location=postgres_location)
+    def test_postgres_flexible_server_tuning_options_resource(self, resource_group):
+        self._test_tuning_options_mgmt('postgres', resource_group)
+
+    def _test_tuning_options_mgmt(self, database_engine, resource_group):
+
+        # Create server with at least 4 vCores and running PostgreSQL major version of 13 or later
+        location = self.postgres_location
+        server_name = self.create_random_name(SERVER_NAME_PREFIX, SERVER_NAME_MAX_LENGTH)
+        version = '16'
+        storage_size = 128
+        sku_name = 'Standard_D4s_v3'
+        tier = 'GeneralPurpose'
+
+        self.cmd('{} flexible-server create -g {} -n {} --sku-name {} --tier {} --storage-size {} --version {} -l {} --public-access none --yes'.format(
+                 database_engine, resource_group, server_name, sku_name, tier, storage_size, version, location))
+
+        # Get tuning options for server
+        self.cmd('{} flexible-server tuning-options show -g {} -s {}'.format(database_engine, resource_group, server_name),
+                 checks=[JMESPathCheck('name', 'index')])
+
+        # Enable index tuning for server
+        self.cmd('{} flexible-server update-index-tuning -g {} -s {} --state Enabled'.format(database_engine, resource_group, server_name))
+
+        # Check for updated parameters
+        self.cmd('{} flexible-server parameter show --name {} -g {} -s {}'.format(database_engine, 'index_tuning.mode', resource_group, server_name),
+                 checks=[JMESPathCheck('value', 'report')])
+        self.cmd('{} flexible-server parameter show --name {} -g {} -s {}'.format(database_engine, 'pg_qs.query_capture_mode', resource_group, server_name),
+                 checks=[JMESPathCheck('value', 'all')])
+
+        # Disable index tuning for server
+        self.cmd('{} flexible-server update-index-tuning -g {} -s {} --state Disabled'.format(database_engine, resource_group, server_name))
+
+        # Check for updated parameters
+        self.cmd('{} flexible-server parameter show --name {} -g {} -s {}'.format(database_engine, 'index_tuning.mode', resource_group, server_name),
+                 checks=[JMESPathCheck('value', 'off')])