From 5af5316ad75c1f69721784e30500a464185aeb25 Mon Sep 17 00:00:00 2001 From: Srinikhil Naravamakula Date: Tue, 9 May 2023 19:11:24 -0500 Subject: [PATCH 1/2] CosmosDB: Adds changes to handle enabling partition merge on CosmosDB account. --- .../cli/command_modules/cosmosdb/custom.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py b/src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py index b1574682051..ecbc3155b3b 100644 --- a/src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py +++ b/src/azure-cli/azure/cli/command_modules/cosmosdb/custom.py @@ -132,7 +132,8 @@ def cli_cosmosdb_create(cmd, tables_to_restore=None, is_restore_request=None, restore_source=None, - restore_timestamp=None): + restore_timestamp=None, + enable_partition_merge=None): """Create a new Azure Cosmos DB database account.""" from azure.cli.core.commands.client_factory import get_mgmt_service_client @@ -183,7 +184,8 @@ def cli_cosmosdb_create(cmd, databases_to_restore=databases_to_restore, gremlin_databases_to_restore=gremlin_databases_to_restore, tables_to_restore=tables_to_restore, - arm_location=resource_group_location) + arm_location=resource_group_location, + enable_partition_merge=enable_partition_merge) # pylint: disable=too-many-statements @@ -224,7 +226,8 @@ def _create_database_account(client, is_restore_request=None, restore_source=None, restore_timestamp=None, - arm_location=None): + arm_location=None, + enable_partition_merge=None): consistency_policy = None if default_consistency_level is not None: @@ -350,7 +353,8 @@ def _create_database_account(client, default_identity=default_identity, analytical_storage_configuration=analytical_storage_configuration, create_mode=create_mode, - restore_parameters=restore_parameters + restore_parameters=restore_parameters, + enable_partition_merge=enable_partition_merge ) async_docdb_create = client.begin_create_or_update(resource_group_name, account_name, params) @@ -386,7 +390,8 @@ def cli_cosmosdb_update(client, backup_redundancy=None, default_identity=None, analytical_storage_schema_type=None, - backup_policy_type=None): + backup_policy_type=None, + enable_partition_merge=None): """Update an existing Azure Cosmos DB database account. """ existing = client.get(resource_group_name, account_name) @@ -468,7 +473,8 @@ def cli_cosmosdb_update(client, api_properties=api_properties, backup_policy=backup_policy, default_identity=default_identity, - analytical_storage_configuration=analytical_storage_configuration) + analytical_storage_configuration=analytical_storage_configuration, + enable_partition_merge=enable_partition_merge) async_docdb_update = client.begin_update(resource_group_name, account_name, params) docdb_account = async_docdb_update.result() From 5704fea59a9e4eee2bd60846ddc76110a2f405de Mon Sep 17 00:00:00 2001 From: Srinikhil Naravamakula Date: Tue, 9 May 2023 19:35:45 -0500 Subject: [PATCH 2/2] CosmosDB: Adds changes to handle enabling partition merge on CosmosDB account. --- src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py b/src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py index 72780d124d9..69911aef425 100644 --- a/src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py +++ b/src/azure-cli/azure/cli/command_modules/cosmosdb/_params.py @@ -87,6 +87,7 @@ def load_arguments(self, _): c.argument('databases_to_restore', nargs='+', action=CreateDatabaseRestoreResource, arg_group='Restore') c.argument('gremlin_databases_to_restore', nargs='+', action=CreateGremlinDatabaseRestoreResource, arg_group='Restore') c.argument('tables_to_restore', nargs='+', action=CreateTableRestoreResource, arg_group='Restore') + c.argument('enable_partition_merge', arg_type=get_three_state_flag(), help="Flag to enable partition merge on the account.") for scope in ['cosmosdb create', 'cosmosdb update']: with self.argument_context(scope) as c: @@ -116,7 +117,8 @@ def load_arguments(self, _): c.argument('default_identity', help="The primary identity to access key vault in CMK related features. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more. User-assigned identities are specified in format 'UserAssignedIdentity='.", is_preview=True) c.argument('analytical_storage_schema_type', options_list=['--analytical-storage-schema-type', '--as-schema'], arg_type=get_enum_type(AnalyticalStorageSchemaType), help="Schema type for analytical storage.", arg_group='Analytical Storage Configuration') c.argument('backup_policy_type', arg_type=get_enum_type(BackupPolicyType), help="The type of backup policy of the account to create", arg_group='Backup Policy') - + c.argument('enable_partition_merge', arg_type=get_three_state_flag(), help="Flag to enable partition merge on the account.") + for scope in ['cosmosdb regenerate-key', 'cosmosdb keys regenerate']: with self.argument_context(scope) as c: c.argument('key_kind', arg_type=get_enum_type(KeyKind), help="The access key to regenerate.")