Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[CosmosDB] Adding support to retrieve and redistribute physical partition throughput #4864

Merged
merged 6 commits into from
Jul 22, 2022
Merged
Show file tree
Hide file tree
Changes from 4 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions src/cosmosdb-preview/HISTORY.rst
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,10 @@
Release History
===============

0.18.0
++++++
* Add support for retrieving and redistributing throughput at physical partition level.

0.17.0
++++++
* Add support for new Continuous 7 Days backup mode
Expand Down
1 change: 1 addition & 0 deletions src/cosmosdb-preview/README.md
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@ This package provides commands to
- Update an Azure Cosmos DB database account to enable materialized views
- Create/Delete a cosmosdb materialized views builder service resource
- Provision and update database account with Continuous 7 days backup mode
- Retrieve and redistribute throughput at physical partition level.

## How to use ##

Expand Down
48 changes: 48 additions & 0 deletions src/cosmosdb-preview/azext_cosmosdb_preview/_help.py
Original file line number Diff line number Diff line change
Expand Up @@ -644,3 +644,51 @@
text: |-
az cosmosdb mongodb collection merge -g my-resource-group -a my-account -d my-db --name my-mongodb-collection
"""

helps['cosmosdb sql container retrieve-partition-throughput'] = """
type: command
short-summary: "Retrieve the partition throughput of a sql container."
examples:
- name: Retrieve container container_name's throughput for specific physical partitions
text: |-
az cosmosdb sql container retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --physical-partition-ids 8 9"
- name: Retrieve container container_name's throughput for all physical partitions
text: |-
az cosmosdb sql container retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --all-partitions
"""

helps['cosmosdb sql container redistribute-partition-throughput'] = """
type: command
short-summary: "Redistributes the partition throughput of a sql container."
examples:
- name: Evenly distributes the partition throughput for a sql container among all physical partitions
text: |-
az cosmosdb sql container redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --evenly-distribute
- name: Redistributes the partition throughput for a sql container from source partitions to target partitions
text: |-
az cosmosdb sql container redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --target-partition-info 8=1200 6=1200]' --source-partition-info 9]'
"""

helps['cosmosdb mongodb collection retrieve-partition-throughput'] = """
type: command
short-summary: "Retrieve the partition throughput of a mongodb collection."
examples:
- name: Retrieve container container_name's throughput for specific physical partitions
text: |-
az cosmosdb mongodb collection retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --physical-partition-ids 8 9
- name: Retrieve container container_name's throughput for all physical partitions
text: |-
az cosmosdb mongodb collection retrieve-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --all-partitions
"""

helps['cosmosdb mongodb collection redistribute-partition-throughput'] = """
type: command
short-summary: "Redistributes the partition throughput of a mongodb collection."
examples:
- name: Evenly distributes the partition throughput for a mongodb collection among all physical partitions
text: |-
az cosmosdb mongodb collection redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --evenly-distribute
- name: Redistributes the partition throughput for a mongodb collection from source partitions to target partitions
text: |-
az cosmosdb mongodb collection redistribute-partition-throughput --account-name account_name --database-name db_name --name container_name --resource-group rg_name --target-partition-info 8=1200 6=1200' --source-partition-info 9'
"""
42 changes: 41 additions & 1 deletion src/cosmosdb-preview/azext_cosmosdb_preview/_params.py
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,13 @@
validate_mongo_user_definition_id)

from azext_cosmosdb_preview.actions import (
CreateGremlinDatabaseRestoreResource, CreateTableRestoreResource, AddCassandraTableAction, AddSqlContainerAction)
CreateGremlinDatabaseRestoreResource,
CreateTableRestoreResource,
AddCassandraTableAction,
AddSqlContainerAction,
CreateTargetPhysicalPartitionThroughputInfoAction,
CreateSourcePhysicalPartitionThroughputInfoAction,
CreatePhysicalPartitionIdListAction)

from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.models import (
ContinuousTier
Expand Down Expand Up @@ -331,3 +337,37 @@ def load_arguments(self, _):
c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
c.argument('database_name', database_name_type, required=True, help='Name of the mongoDB database')
c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the mongoDB collection')

# Sql container partition retrieve throughput
with self.argument_context('cosmosdb sql container retrieve-partition-throughput') as c:
c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB container')
c.argument('physical_partition_ids', options_list=['--physical-partition-ids', '-p'], nargs='+', action=CreatePhysicalPartitionIdListAction, required=False, help='space separated list of physical partition ids')
c.argument('all_partitions', arg_type=get_three_state_flag(), help="switch to retrieve throughput for all physical partitions")

# Sql container partition redistribute throughput
with self.argument_context('cosmosdb sql container redistribute-partition-throughput') as c:
c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
c.argument('container_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB collection')
c.argument('evenly_distribute', arg_type=get_three_state_flag(), help="switch to distribute throughput equally among all physical partitions")
c.argument('target_partition_info', nargs='+', action=CreateTargetPhysicalPartitionThroughputInfoAction, required=False, help="information about desired target physical partition throughput eg: 0=1200 1=1200")
c.argument('source_partition_info', nargs='+', action=CreateSourcePhysicalPartitionThroughputInfoAction, required=False, help="space separated source physical partition ids eg: 1 2")

# Mongodb collection partition retrieve throughput
with self.argument_context('cosmosdb mongodb collection retrieve-partition-throughput') as c:
c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
c.argument('collection_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB container')
c.argument('physical_partition_ids', options_list=['--physical-partition-ids', '-p'], nargs='+', action=CreatePhysicalPartitionIdListAction, required=False, help='space separated list of physical partition ids')
c.argument('all_partitions', arg_type=get_three_state_flag(), help="switch to retrieve throughput for all physical partitions")

# Mongodb collection partition redistribute throughput
with self.argument_context('cosmosdb mongodb collection redistribute-partition-throughput') as c:
c.argument('account_name', account_name_type, id_part=None, required=True, help='Name of the CosmosDB database account')
c.argument('database_name', database_name_type, required=True, help='Name of the CosmosDB database name')
c.argument('collection_name', options_list=['--name', '-n'], required=True, help='Name of the CosmosDB collection')
c.argument('evenly_distribute', arg_type=get_three_state_flag(), help="switch to distribute throughput equally among all physical partitions")
c.argument('target_partition_info', nargs='+', action=CreateTargetPhysicalPartitionThroughputInfoAction, required=False, help="information about desired target physical partition throughput eg: '0=1200 1=1200'")
c.argument('source_partition_info', nargs='+', action=CreateSourcePhysicalPartitionThroughputInfoAction, required=False, help="space separated source physical partition ids eg: 1 2")
43 changes: 42 additions & 1 deletion src/cosmosdb-preview/azext_cosmosdb_preview/actions.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,7 +13,9 @@
DatabaseRestoreResource,
GremlinDatabaseRestoreResource,
CosmosCassandraDataTransferDataSourceSink,
CosmosSqlDataTransferDataSourceSink
CosmosSqlDataTransferDataSourceSink,
PhysicalPartitionThroughputInfoResource,
PhysicalPartitionId
)

logger = get_logger(__name__)
Expand Down Expand Up @@ -173,3 +175,42 @@ def __call__(self, parser, namespace, values, option_string=None):
namespace.dest_sql_container = sql_container
else:
namespace.sql_container = sql_container

# pylint: disable=protected-access, too-few-public-methods
class CreateTargetPhysicalPartitionThroughputInfoAction(argparse._AppendAction):
def __call__(self, parser, namespace, values, option_string=None):
if namespace.target_partition_info is None:
namespace.target_partition_info = []
if not values:
# pylint: disable=line-too-long
raise CLIError('usage error: --target-partition-info [PhysicalPartitionId1=Throughput1 PhysicalPartitionId2=Throughput2 ...]')
for item in values:
kvp = item.split('=', 1)
if len(kvp) != 2:
raise CLIError('usage error: --target-partition-info [PhysicalPartitionId1=Throughput1 PhysicalPartitionId2=Throughput2 ...]')
namespace.target_partition_info.append(
PhysicalPartitionThroughputInfoResource(id=kvp[0], throughput=kvp[1]))

# pylint: disable=protected-access, too-few-public-methods
class CreateSourcePhysicalPartitionThroughputInfoAction(argparse._AppendAction):
def __call__(self, parser, namespace, values, option_string=None):
if namespace.source_partition_info is None:
namespace.source_partition_info = []
if not values:
# pylint: disable=line-too-long
raise CLIError('usage error: --source-partition-info [PhysicalPartitionId1 PhysicalPartitionId2 ...]')
for item in values:
namespace.source_partition_info.append(
PhysicalPartitionThroughputInfoResource(id=item, throughput=0))

# pylint: disable=protected-access, too-few-public-methods
class CreatePhysicalPartitionIdListAction(argparse._AppendAction):
def __call__(self, parser, namespace, values, option_string=None):
if namespace.physical_partition_ids is None:
namespace.physical_partition_ids = []
if not values:
# pylint: disable=line-too-long
raise CLIError('usage error: --physical-partition-ids [PhysicalPartitionId1 PhysicalPartitionId2 ...]')
for item in values:
namespace.physical_partition_ids.append(
PhysicalPartitionId(id=item))
16 changes: 16 additions & 0 deletions src/cosmosdb-preview/azext_cosmosdb_preview/commands.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,3 +203,19 @@ def load_command_table(self, _):

with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
g.custom_command('merge', 'cli_begin_list_mongo_db_collection_partition_merge', is_preview=True)

# Retrieve partition throughput for Sql containers
with self.command_group('cosmosdb sql container', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
g.custom_command('retrieve-partition-throughput', 'cli_begin_retrieve_sql_container_partition_throughput', is_preview=True)

# Redistribute partition throughput for Sql containers
with self.command_group('cosmosdb sql container', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
g.custom_command('redistribute-partition-throughput', 'cli_begin_redistribute_sql_container_partition_throughput', is_preview=True)
ravgill marked this conversation as resolved.
Show resolved Hide resolved

# Retrieve partition throughput for Mongo collection
with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
g.custom_command('retrieve-partition-throughput', 'cli_begin_retrieve_mongo_container_partition_throughput', is_preview=True)

# Redistribute partition throughput for Mongo collection
with self.command_group('cosmosdb mongodb collection', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
g.custom_command('redistribute-partition-throughput', 'cli_begin_redistribute_mongo_container_partition_throughput', is_preview=True)
Loading