Merged into a single cmdlet

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

@@ -6,6 +6,7 @@
 # pylint: disable=too-many-statements
 # pylint: disable=line-too-long
 from argcomplete.completers import FilesCompleter
+from enum import Enum
 
 from azure.cli.core.commands.parameters import (
     get_resource_name_completion_list, name_type, get_enum_type, get_three_state_flag, tags_type)
@@ -36,6 +37,10 @@
 CASSANDRA_SCHEMA_EXAMPLE = """--schema "{\\"columns\\": [{\\"name\\": \\"columnA\\",\\"type\\": \\"uuid\\"}, {\\"name\\": \\"columnB\\",\\"type\\": \\"Ascii\\"}],\\"partitionKeys\\": [{\\"name\\": \\"columnA\\"}]}"
 """
 
+class ThroughputTypes(str, Enum):
+    autoscale = "autoscale"
+    manual = "manual"
+
 
 def load_arguments(self, _):
     from knack.arguments import CLIArgumentType
@@ -107,6 +112,7 @@ def load_arguments(self, _):
     database_name_type = CLIArgumentType(options_list=['--database-name', '-d'], help='Database name.')
     container_name_type = CLIArgumentType(options_list=['--container-name', '-c'], help='Container name.')
     max_throughput_type = CLIArgumentType(options_list=['--max-throughput'], help='The maximum throughput resource can scale to (RU/s). Provided when the resource is autoscale enabled. The minimum value can be 4000 (RU/s)')
+    throughput_type = CLIArgumentType(options_list=['--throughput-type', '-t'], arg_type = get_enum_type(ThroughputTypes), help='The type of throughput to migrate to.')
 
     with self.argument_context('cosmosdb private-endpoint-connection') as c:
         c.argument('private_endpoint_connection_name', options_list=['--name', '-n'], required=False,
@@ -240,54 +246,83 @@ def load_arguments(self, _):
         c.argument('throughput', type=int, help='The throughput of SQL database (RU/s).')
         c.argument('max_throughput', max_throughput_type)
 
+    with self.argument_context('cosmosdb sql database throughput migrate') as c:
+        c.argument('account_name', account_name_type, id_part=None)
+        c.argument('database_name', options_list=['--name', '-n'], help="Database name")
+        c.argument('throughput_type', throughput_type)
+
     with self.argument_context('cosmosdb sql container throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('database_name', database_name_type)
         c.argument('container_name', options_list=['--name', '-n'], help="Container name")
         c.argument('throughput', type=int, help='The throughput of SQL container (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb sql container throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb mongodb database throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('database_name', options_list=['--name', '-n'], help="Database name")
         c.argument('throughput', type=int, help='The throughput of MongoDB database (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb mongodb database throughput migrate') as c:
+        c.argument('migrate', throughput_type)
 
     with self.argument_context('cosmosdb mongodb collection throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('database_name', database_name_type)
         c.argument('collection_name', options_list=['--name', '-n'], help="Collection name")
         c.argument('throughput', type=int, help='The throughput of MongoDB collection (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb mongodb collection throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb cassandra keyspace throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('keyspace_name', options_list=['--name', '-n'], help="Keyspace name")
         c.argument('throughput', type=int, help='The throughput of Cassandra keyspace (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb cassandra keyspace throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb cassandra table throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('keyspace_name', options_list=['--keyspace-name', '-k'], help="Keyspace name")
         c.argument('table_name', options_list=['--name', '-n'], help="Table name")
         c.argument('throughput', type=int, help='The throughput of Cassandra table (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb cassandra table throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb gremlin database throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('database_name', options_list=['--name', '-n'], help="Database name")
         c.argument('throughput', type=int, help='The throughput of Gremlin database (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb gremlin database throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb gremlin graph throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('database_name', database_name_type)
-        c.argument('graph_name', options_list=['--name', '-n'], help="Grapth name")
+        c.argument('graph_name', options_list=['--name', '-n'], help="Graph name")
         c.argument('throughput', type=int, help='The throughput Gremlin graph (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb gremlin graph throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)
 
     with self.argument_context('cosmosdb table throughput') as c:
         c.argument('account_name', account_name_type, id_part=None)
         c.argument('table_name', options_list=['--name', '-n'], help="Table name")
         c.argument('throughput', type=int, help='The throughput of Table (RU/s).')
         c.argument('max_throughput', max_throughput_type)
+
+    with self.argument_context('cosmosdb table throughput migrate') as c:
+        c.argument('throughput_type', throughput_type)

src/azure-cli/azure/cli/command_modules/cosmosdb/commands.py

@@ -202,58 +202,49 @@ def load_command_table(self, _):
 
     # Offer throughput
     with self.command_group('cosmosdb sql database throughput', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
-        g.command('show', 'get_sql_database_throughput')
-        g.command('migrate-to-autoscale', 'migrate_sql_database_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_sql_database_to_manual_throughput')
+        g.show_command('show', 'get_sql_database_throughput')
+        g.custom_command('migrate', 'cli_cosmosdb_sql_database_throughput_migrate')
         g.custom_command('update', 'cli_cosmosdb_sql_database_throughput_update')
 
     with self.command_group('cosmosdb sql container throughput', cosmosdb_sql_sdk, client_factory=cf_sql_resources) as g:
-        g.command('show', 'get_sql_container_throughput')
-        g.command('migrate-to-autoscale', 'migrate_sql_container_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_sql_container_to_manual_throughput')
+        g.show_command('show', 'get_sql_container_throughput')
         g.custom_command('update', 'cli_cosmosdb_sql_container_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_sql_container_throughput_migrate')
 
     with self.command_group('cosmosdb mongodb database throughput', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
-        g.command('show', 'get_mongo_db_database_throughput')
-        g.command('migrate-to-autoscale', 'migrate_mongo_db_database_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_mongo_db_database_to_manual_throughput')
+        g.show_command('show', 'get_mongo_db_database_throughput')
         g.custom_command('update', 'cli_cosmosdb_mongodb_database_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_mongodb_database_throughput_migrate')
 
     with self.command_group('cosmosdb mongodb collection throughput', cosmosdb_mongo_sdk, client_factory=cf_mongo_db_resources) as g:
-        g.command('show', 'get_mongo_db_collection_throughput')
-        g.command('migrate-to-autoscale', 'migrate_mongo_db_collection_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_mongo_db_collection_to_manual_throughput')
+        g.show_command('show', 'get_mongo_db_collection_throughput')
         g.custom_command('update', 'cli_cosmosdb_mongodb_collection_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_mongodb_collection_throughput_migrate')
 
     with self.command_group('cosmosdb cassandra keyspace throughput', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
-        g.command('show', 'get_cassandra_keyspace_throughput')
-        g.command('migrate-to-autoscale', 'migrate_cassandra_keyspace_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_cassandra_keyspace_to_manual_throughput')
+        g.show_command('show', 'get_cassandra_keyspace_throughput')
         g.custom_command('update', 'cli_cosmosdb_cassandra_keyspace_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_cassandra_keyspace_throughput_migrate')
 
     with self.command_group('cosmosdb cassandra table throughput', cosmosdb_cassandra_sdk, client_factory=cf_cassandra_resources) as g:
-        g.command('show', 'get_cassandra_table_throughput')
-        g.command('migrate-to-autoscale', 'migrate_cassandra_table_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_cassandra_table_to_manual_throughput')
+        g.show_command('show', 'get_cassandra_table_throughput')
         g.custom_command('update', 'cli_cosmosdb_cassandra_table_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_cassandra_table_throughput_migrate')
 
     with self.command_group('cosmosdb gremlin database throughput', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
         g.command('show', 'get_gremlin_database_throughput')
-        g.command('migrate-to-autoscale', 'migrate_gremlin_database_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_gremlin_database_to_manual_throughput')
         g.custom_command('update', 'cli_cosmosdb_gremlin_database_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_gremlin_database_throughput_migrate')
 
     with self.command_group('cosmosdb gremlin graph throughput', cosmosdb_gremlin_sdk, client_factory=cf_gremlin_resources) as g:
-        g.command('show', 'get_gremlin_graph_throughput')
-        g.command('migrate-to-autoscale', 'migrate_gremlin_graph_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_gremlin_graph_to_manual_throughput')
+        g.show_command('show', 'get_gremlin_graph_throughput')
         g.custom_command('update', 'cli_cosmosdb_gremlin_graph_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_gremlin_graph_throughput_migrate')
 
     with self.command_group('cosmosdb table throughput', cosmosdb_table_sdk, client_factory=cf_table_resources) as g:
-        g.command('show', 'get_table_throughput')
-        g.command('migrate-to-autoscale', 'migrate_table_to_autoscale')
-        g.command('migrate-to-manual-throughput', 'migrate_table_to_manual_throughput')
+        g.show_command('show', 'get_table_throughput')
         g.custom_command('update', 'cli_cosmosdb_table_throughput_update')
+        g.custom_command('migrate', 'cli_cosmosdb_table_throughput_migrate')
 
     # virtual network rules
     with self.command_group('cosmosdb network-rule', None, client_factory=cf_db_accounts) as g: