[CosmosDB] az cosmosdb fleet: Add new fleet feature (#32390)

Co-authored-by: Darshan Sapaliga <[email protected]>
Co-authored-by: Achint Agrawal <[email protected]>

Author: 3 people

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

@@ -174,3 +174,16 @@ def cf_restorable_tables(cli_ctx, _):
 
 def cf_restorable_table_resources(cli_ctx, _):
     return cf_cosmosdb(cli_ctx).restorable_table_resources
+
+
+# fleet
+def cf_fleet(cli_ctx, _):
+    return cf_cosmosdb(cli_ctx).fleet
+
+
+def cf_fleetspace(cli_ctx, _):
+    return cf_cosmosdb(cli_ctx).fleetspace
+
+
+def cf_fleetspace_account(cli_ctx, _):
+    return cf_cosmosdb(cli_ctx).fleetspace_account

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

@@ -1785,3 +1785,115 @@
     text: |
       az cosmosdb table restore --resource-group resource_group --account-name database_account_name --table-name name_of_table_needs_to_be_restored --restore-timestamp 2020-07-13T16:03:41+0000
 """
+
+helps['cosmosdb fleet'] = """
+type: group
+short-summary: Manage Azure Cosmos DB Fleet resources.
+"""
+
+helps['cosmosdb fleet create'] = """
+type: command
+short-summary: Create a new Cosmos DB Fleet.
+examples:
+  - name: Create a fleet
+    text: |
+      az cosmosdb fleet create \\
+        --resource-group MyResourceGroup \\
+        --fleet-name MyFleet \\
+        --location westus
+"""
+
+helps['cosmosdb fleet list'] = """
+type: command
+short-summary: List Cosmos DB Fleets in a subscription or resource group.
+"""
+
+helps['cosmosdb fleet show'] = """
+type: command
+short-summary: Show details of a specific Cosmos DB Fleet.
+"""
+
+helps['cosmosdb fleet delete'] = """
+type: command
+short-summary: Delete a specific Cosmos DB Fleet.
+"""
+
+helps['cosmosdb fleetspace'] = """
+type: group
+short-summary: Manage Cosmos DB Fleetspace resources.
+"""
+
+helps['cosmosdb fleetspace create'] = """
+type: command
+short-summary: Create a new Fleetspace under a Cosmos DB Fleet.
+examples:
+  - name: Create a fleetspace
+    text: |
+      az cosmosdb fleetspace create \\
+        --resource-group MyResourceGroup \\
+        --fleet-name MyFleet \\
+        --fleetspace-name MyFleetspace \\
+        --body @fleetspace.json
+"""
+
+helps['cosmosdb fleetspace update'] = """
+type: command
+short-summary: Update an existing Cosmos DB Fleetspace.
+examples:
+  - name: Update fleetspace throughput settings
+    text: |
+      az cosmosdb fleetspace update \\
+        --resource-group MyResourceGroup \\
+        --fleet-name MyFleet \\
+        --fleetspace-name MyFleetspace \\
+        --body @fleetspace.json
+"""
+
+helps['cosmosdb fleetspace list'] = """
+type: command
+short-summary: List all Fleetspaces under a Fleet.
+"""
+
+helps['cosmosdb fleetspace show'] = """
+type: command
+short-summary: Show details of a specific Fleetspace.
+"""
+
+helps['cosmosdb fleetspace delete'] = """
+type: command
+short-summary: Delete a Fleetspace from a Fleet.
+"""
+
+helps['cosmosdb fleetspace account'] = """
+type: group
+short-summary: Manage database accounts within a Cosmos DB Fleetspace.
+"""
+
+helps['cosmosdb fleetspace account create'] = """
+type: command
+short-summary: Register an existing Cosmos DB database account to a Fleetspace.
+examples:
+  - name: Register a database account to a fleetspace
+    text: |
+      az cosmosdb fleetspace account create \\
+        --resource-group MyResourceGroup \\
+        --fleet-name MyFleet \\
+        --fleetspace-name MyFleetspace \\
+        --fleetspace-account-name MyAccount \\
+        --body @fleetspaceAccount.json
+"""
+
+helps['cosmosdb fleetspace account list'] = """
+type: command
+short-summary: List all database accounts associated with a Fleetspace.
+"""
+
+helps['cosmosdb fleetspace account show'] = """
+type: command
+short-summary: Show details of a registered database account in a Fleetspace.
+"""
+
+helps['cosmosdb fleetspace account delete'] = """
+type: command
+short-summary: Unregister a database account from a Fleetspace.
+"""

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

@@ -28,14 +28,16 @@
     validate_mongo_role_definition_body,
     validate_mongo_role_definition_id,
     validate_mongo_user_definition_body,
-    validate_mongo_user_definition_id)
+    validate_mongo_user_definition_id,
+    validate_fleetspace_body,
+    validate_fleetspaceAccount_body)
 
 from azure.cli.command_modules.cosmosdb.actions import (
     CreateLocation, CreateDatabaseRestoreResource, CreateGremlinDatabaseRestoreResource, CreateTableRestoreResource, UtcDatetimeAction, InvokeCommandArgumentsAddAction)
 from azure.cli.command_modules.cosmosdb.custom import (
     CosmosKeyTypes)
 from azure.mgmt.cosmosdb.models import (
-    ContinuousTier, MinimalTlsVersion)
+    ContinuousTier, MinimalTlsVersion, DefaultPriorityLevel)
 
 GREMLIN_INDEXING_POLICY_EXAMPLE = """--idx "{\\"indexingMode\\": \\"consistent\\", \\"automatic\\": true, \\"includedPaths\\": [{\\"path\\": \\"/*\\"}], \\"excludedPaths\\": [{ \\"path\\": \\"/headquarters/employees/?\\"}, { \\"path\\": \\"/\\\\"_etag\\\\"/?\\"}]}"
 """
@@ -72,6 +74,28 @@
 MONGO_USER_DEFINITION_EXAMPLE = """--body "{\\"Id\\": \\"be79875a-2cc4-40d5-8958-566017875b39\\",\\"UserName\\": \\"MyUserName\\",\\"Password\\": \\"MyPass\\",\\"CustomData\\": \\"MyCustomData\\",\\"Mechanisms\\": \\"SCRAM-SHA-256\\"\\"DatabaseName\\": \\"MyDb\\",\\"Roles\\": [ {\\"Role\\": \\"myReadRole\\",\\"Db\\": \\"MyDb\\"}]}"
 """
 
+FLEETSPACE_PROPERTIES_EXAMPLE = """--body "{
+    \\"properties\\": {
+        \\"serviceTier\\": \\"GeneralPurpose\\",
+        \\"dataRegions\\": [\\"West US 2\\"],
+        \\"throughputPoolConfiguration\\": {
+            \\"minThroughput\\": 100000,
+            \\"maxThroughput\\": 300000
+        }
+    }
+}"
+"""
+
+FLEETSPACE_ACCOUNT_PROPERTIES_EXAMPLE = """--body "{
+    \\"properties\\": {
+        \\"globalDatabaseAccountProperties\\": {
+            \\"resourceId\\": \\"/subscriptions/00000000-0000-0000-0000-000000000000/resourceGroups/example-rg/providers/Microsoft.DocumentDB/databaseAccounts/example-account\\",
+            \\"armLocation\\": \\"East US\\"
+        }
+    }
+}"
+"""
+
 
 class ThroughputTypes(str, Enum):
     autoscale = "autoscale"
@@ -131,6 +155,8 @@ def load_arguments(self, _):
             c.argument('enable_partition_merge', arg_type=get_three_state_flag(), help="Flag to enable partition merge on the account.")
             c.argument('enable_burst_capacity', arg_type=get_three_state_flag(), help="Flag to enable burst capacity on the account.")
             c.argument('enable_prpp_autoscale', arg_type=get_three_state_flag(), help="Enable or disable PerRegionPerPartitionAutoscale.")
+            c.argument('enable_pbe', arg_type=get_three_state_flag(), help="Flag to enable priority based execution on the account.")
+            c.argument('default_priority_level', arg_type=get_enum_type(DefaultPriorityLevel), help="Default Priority Level of Request if not specified.")
             c.argument('continuous_tier', arg_type=get_enum_type(ContinuousTier), help="The tier of Continuous backup", arg_group='Backup Policy')
             c.argument('minimal_tls_version', arg_type=get_enum_type(MinimalTlsVersion), help="Indicates the minimum allowed TLS version")
 
@@ -407,14 +433,15 @@ def load_arguments(self, _):
         c.argument('target_database_account_name', options_list=['--target-database-account-name', '-n'], help='Name of the new target Cosmos DB database account after the restore')
         c.argument('account_name', completer=None, options_list=['--account-name', '-a'], help='Name of the source Cosmos DB database account for the restore', id_part=None)
         c.argument('restore_timestamp', options_list=['--restore-timestamp', '-t'], action=UtcDatetimeAction, help="The timestamp to which the account has to be restored to.")
-        c.argument('location', arg_type=get_location_type(self.cli_ctx), help="The location of the source account from which restore is triggered. This will also be the write region of the restored account")
+        c.argument('location', arg_type=get_location_type(self.cli_ctx), help="This is the write region of the restored account. This is also the location of the source account where its backups are located if source_backup_location is not provided.")
         c.argument('databases_to_restore', nargs='+', action=CreateDatabaseRestoreResource)
         c.argument('gremlin_databases_to_restore', nargs='+', action=CreateGremlinDatabaseRestoreResource)
         c.argument('tables_to_restore', nargs='+', action=CreateTableRestoreResource)
         c.argument('assign_identity', nargs='*', help="Assign system or user assigned identities separated by spaces. Use '[system]' to refer system assigned identity.")
         c.argument('default_identity', help="The primary identity to access key vault in CMK related features. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more.")
         c.argument('public_network_access', options_list=['--public-network-access', '-p'], arg_type=get_enum_type(['ENABLED', 'DISABLED']), help="Sets public network access in server to either Enabled or Disabled.")
         c.argument('disable_ttl', options_list=['--disable-ttl', '-d'], arg_type=get_three_state_flag(), help="Enable or disable restoring with ttl disabled.")
+        c.argument('source_backup_location', help="This is the location of the source account where backups are located. Provide this value if the source and target are in different locations.", is_preview=True)
 
     # Mongo role definition
     with self.argument_context('cosmosdb mongodb role definition') as c:
@@ -681,3 +708,34 @@ def load_arguments(self, _):
         c.argument('instance_count', options_list=['--count', '-c'], help="Instance Count.")
         c.argument('instance_size', options_list=['--size'], help="Instance Size. Possible values are: Cosmos.D4s, Cosmos.D8s, Cosmos.D16s etc")
         c.argument('dedicated_gateway_type', options_list=['--gateway-type'], arg_type=get_enum_type(['IntegratedCache', 'DistributedQuery']), help="Dedicated Gateway Type. Valid only for SqlDedicatedGateway service kind")
+
+    # Cosmos DB Fleet
+    with self.argument_context('cosmosdb fleet') as c:
+        c.argument('resource_group', options_list=['--resource-group', '-g'], help='Name of the resource group.', required=True)
+        c.argument('fleet_name', options_list=['--fleet-name', '-n'], help='Name of the Fleet resource.', required=True)
+
+    with self.argument_context('cosmosdb fleet create') as c:
+        c.argument('location', options_list=['--location', '-l'], help='Location of the Fleet.', required=True)
+        c.argument('tags', help="Tags in 'key=value key2=value2' format.")
+
+    # Cosmos DB Fleetspace
+    with self.argument_context('cosmosdb fleetspace') as c:
+        c.argument('resource_group', options_list=['--resource-group', '-g'], help='Name of the resource group.', required=True)
+        c.argument('fleet_name', options_list=['--fleet-name'], help='Name of the Cosmos DB Fleet.', required=True)
+        c.argument('fleetspace_name', options_list=['--fleetspace-name', '-n'], help='Name of the Fleetspace resource.', required=True)
+
+    with self.argument_context('cosmosdb fleetspace create') as c:
+        c.argument('fleetspace_body', options_list=['--body', '-b'], validator=validate_fleetspace_body, completer=FilesCompleter(), help="Fleetspace body with properties.serviceTier (required), properties.dataRegions (required), and properties.throughputPoolConfiguration (fields: minThroughput, maxThroughput). You can enter it as a string or as a file, e.g., --body @fleetspace.json or " + FLEETSPACE_PROPERTIES_EXAMPLE)
+
+    with self.argument_context('cosmosdb fleetspace update') as c:
+        c.argument('fleetspace_body', options_list=['--body', '-b'], validator=validate_fleetspace_body, completer=FilesCompleter(), help="Fleetspace body with properties.serviceTier (optional), properties.dataRegions (optional), and properties.throughputPoolConfiguration (fields: minThroughput, maxThroughput). You can enter it as a string or as a file, e.g., --body @fleetspace.json or " + FLEETSPACE_PROPERTIES_EXAMPLE)
+
+    # Cosmos DB Fleetspace account
+    with self.argument_context('cosmosdb fleetspace account') as c:
+        c.argument('resource_group', options_list=['--resource-group', '-g'], help='Name of the resource group.', required=True)
+        c.argument('fleet_name', options_list=['--fleet-name'], help='Name of the Cosmos DB Fleet.', required=True)
+        c.argument('fleetspace_name', options_list=['--fleetspace-name'], help='Name of the Fleetspace resource.', required=True)
+        c.argument('fleetspace_account_name', options_list=['--fleetspace-account-name', '-n'], help='Name of the Fleetspace Account resource.', required=True)
+
+    with self.argument_context('cosmosdb fleetspace account create') as c:
+        c.argument('fleetspace_account_body', options_list=['--body', '-b'], validator=validate_fleetspaceAccount_body, completer=FilesCompleter(), help="Fleetspace Account body with properties.globalDatabaseAccountProperties (fields: armLocation, resourceId). You can enter it as a string or as a file, e.g., --body @fleetspaceAccount.json or " + FLEETSPACE_ACCOUNT_PROPERTIES_EXAMPLE)

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

@@ -531,3 +531,77 @@ def _validate_pk_paths_in_cep(path, partition_key_path, policyFormatVersion, enc
                 raise InvalidArgumentValueError(f"Partition key path:{pkPath} is part of "
                                                 "Client Encryption policy with invalid encryption type. "
                                                 "Only deterministic encryption type is supported.")
+
+
+def validate_fleetspace_body(cmd, ns):
+    from azure.cli.core.util import get_file_json, shell_safe_json_parse
+    import os
+
+    if ns.fleetspace_body is not None:
+        if os.path.exists(ns.fleetspace_body):
+            body = get_file_json(ns.fleetspace_body)
+        else:
+            body = shell_safe_json_parse(ns.fleetspace_body)
+
+        if not isinstance(body, dict):
+            raise InvalidArgumentValueError('Invalid fleetspace body. Must be a JSON object.')
+
+        props = body.get('properties', {})
+        if not isinstance(props, dict):
+            raise InvalidArgumentValueError('Missing or invalid "properties" field in fleetspace body.')
+
+        tp_config = props.get('throughputPoolConfiguration', {})
+        if not isinstance(tp_config, dict):
+            raise InvalidArgumentValueError('Missing or invalid "throughputPoolConfiguration" in properties.')
+
+        # Check for minThroughput and maxThroughput in throughputPoolConfiguration
+        for field in ['minThroughput', 'maxThroughput']:
+            if field not in tp_config:
+                raise InvalidArgumentValueError(f'Missing "{field}" in throughputPoolConfiguration.')
+
+        if not isinstance(tp_config['minThroughput'], int) or tp_config['minThroughput'] <= 0:
+            raise InvalidArgumentValueError('"minThroughput" must be a positive integer.')
+
+        if not isinstance(tp_config['maxThroughput'], int) or tp_config['maxThroughput'] <= 0:
+            raise InvalidArgumentValueError('"maxThroughput" must be a positive integer.')
+
+        # Check for serviceTier and dataRegions at base properties level
+        if 'serviceTier' in props:
+            if not isinstance(props['serviceTier'], str):
+                raise InvalidArgumentValueError('"serviceTier" must be a string.')
+
+        if 'dataRegions' in props:
+            if not isinstance(props['dataRegions'], list) or not all(isinstance(r, str) for r in props['dataRegions']):
+                raise InvalidArgumentValueError('"dataRegions" must be a list of strings.')
+
+        ns.fleetspace_body = body
+
+
+def validate_fleetspaceAccount_body(cmd, ns):
+    from azure.cli.core.util import get_file_json, shell_safe_json_parse
+    import os
+
+    if ns.fleetspace_account_body is not None:
+        if os.path.exists(ns.fleetspace_account_body):
+            body = get_file_json(ns.fleetspace_account_body)
+        else:
+            body = shell_safe_json_parse(ns.fleetspace_account_body)
+
+        if not isinstance(body, dict):
+            raise InvalidArgumentValueError("Fleetspace Account body must be a valid JSON object.")
+
+        props = body.get("properties")
+        if not isinstance(props, dict):
+            raise InvalidArgumentValueError('Missing or invalid "properties" field.')
+
+        gdp = props.get("globalDatabaseAccountProperties")
+        if not isinstance(gdp, dict):
+            raise InvalidArgumentValueError('Missing or invalid "globalDatabaseAccountProperties".')
+
+        if "resourceId" not in gdp or not isinstance(gdp["resourceId"], str) or not gdp["resourceId"].startswith("/subscriptions/"):
+            raise InvalidArgumentValueError('"resourceId" must be a valid ARM resource ID string.')
+
+        if "armLocation" not in gdp or not isinstance(gdp["armLocation"], str):
+            raise InvalidArgumentValueError('"armLocation" must be a valid string.')
+
+        ns.fleetspace_account_body = body

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

@@ -32,7 +32,10 @@
     cf_db_locations,
     cf_cassandra_cluster,
     cf_cassandra_data_center,
-    cf_service
+    cf_service,
+    cf_fleet,
+    cf_fleetspace,
+    cf_fleetspace_account
 )
 
 from azure.cli.command_modules.cosmosdb._format import (
@@ -499,3 +502,41 @@ def load_command_table(self, _):
         g.command('list', 'list')
         g.show_command('show', 'get')
         g.command('delete', 'begin_delete', confirmation=True, supports_no_wait=True)
+
+    setup_fleet_commands(self)
+
+
+def setup_fleet_commands(self):
+    # Fleet operations
+    cosmosdb_fleet_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.cosmosdb.operations#FleetOperations.{}',
+        client_factory=cf_fleet)
+
+    # Fleetspace operations
+    cosmosdb_fleetspace_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.cosmosdb.operations#FleetspaceOperations.{}',
+        client_factory=cf_fleetspace)
+
+    # Fleetspace account operations
+    cosmosdb_fleetspace_account_sdk = CliCommandType(
+        operations_tmpl='azure.mgmt.cosmosdb.operations#FleetspaceAccountOperations.{}',
+        client_factory=cf_fleetspace_account)
+
+    with self.command_group('cosmosdb fleet', cosmosdb_fleet_sdk, client_factory=cf_fleet, is_preview=True) as g:
+        g.custom_command('create', 'cli_cosmosdb_fleet_create')
+        g.custom_command('list', 'cli_list_cosmosdb_fleets')
+        g.show_command('show', 'get')
+        g.command('delete', 'begin_delete', confirmation=True)
+
+    with self.command_group('cosmosdb fleetspace', cosmosdb_fleetspace_sdk, client_factory=cf_fleetspace, is_preview=True) as g:
+        g.custom_command('create', 'cli_cosmosdb_fleetspace_create')
+        g.command('list', 'list')
+        g.show_command('show', 'get')
+        g.custom_command('update', 'cli_cosmosdb_fleetspace_update')
+        g.command('delete', 'begin_delete', confirmation=True)
+
+    with self.command_group('cosmosdb fleetspace account', cosmosdb_fleetspace_account_sdk, client_factory=cf_fleetspace_account, is_preview=True) as g:
+        g.custom_command('create', 'cli_cosmosdb_fleetspace_account_create')
+        g.command('list', 'list')
+        g.show_command('show', 'get')
+        g.command('delete', 'begin_delete', confirmation=True)