[cosmosdb-preview] Add support for Priority Based Execution, cross ac…

…count container copy, api-version 2023-09-15-preview (Azure#6703)

* init changes for 20230915-preview api

* Changes to add enablePriorityBasedExecution and defaultPriorityLevel

* Removed isPreview flag from CMK related properties and added key-uri to cosmosdb update operation.

* Added test and test recording for priority based execution

* corrected resource group name and whitespace corrections

* white space corrections

* Add cosmosdb copy command

* Add tests

* Dts/cross account copy (#3)

* Deprecate old command

* Add test recording for dts

---------

Co-authored-by: Nitesh Vijay <[email protected]>

* Fix build error

* Fix Static Analysis

* Use boolean type

* Change help text

* Update help text

* Updated Vendored SDK

* Updated test account name in test_cosmosdb-merge-scenario.py

* Whitespace correction and added linter exclution for enable pbe

* Updated linter_exclusions

* Remove id part

* Corrected account name in test_cosmosdb-merge-scenario.py and generated test recordings

* Generated test recording for test_cosmosdb_materializedview

* Add support for MI Cassandra in 2023-09-15

* Add -n for jobname

* Fix breaking change

* Modified setup.py and history.rst

* Fix static check

* Whitespace correction

* Re-recorded tests

* rerecorded test test_cosmosdb_sql_materializedview_scenario

* removed linter exclusion for enable_priority_based_execution

* Rerecorded all tests

* Updated HISTORY.rst

* Rerecorded tests

* rerecorded test_cosmosdb_sql_materializedview_scenario

* rerecorded test_cosmosdb_sql_materializedview_scenario

* rerecorded tests

---------

Co-authored-by: Achint Agrawal <[email protected]>
Co-authored-by: AdrianSibajaRetana <[email protected]>
Co-authored-by: Nitesh Vijay <[email protected]>
Co-authored-by: Nitesh Vijay <[email protected]>
Co-authored-by: Guanzhou <[email protected]>

src/cosmosdb-preview/HISTORY.rst

@@ -2,6 +2,12 @@
 
 Release History
 ===============
+0.26.0
+* Add '--enable-priority-based-execution' and '--default-priority-level' parameter for create/update database account.
+* Add support for cross account container copy
+* Add Non-CMK to CMK support
+* Removed preview tag from CMK related properties
+
 0.25.0
 * Add support for performing database merge for Sql and MongoDB database account.
 

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -778,6 +778,112 @@
                az cosmosdb dts cancel --account-name "ddb1" --job-name "j1" -g "rg1"
 """
 
+helps['cosmosdb copy'] = """
+    type: group
+    short-summary: Manage container copy job
+"""
+
+helps['cosmosdb copy create'] = """
+    type: command
+    short-summary: "Creates a container copy job."
+    parameters:
+      - name: --src-cassandra
+        short-summary: "Source table"
+        long-summary: |
+            Usage: --src-cassandra keyspace=XX table=XX'
+            keyspace: Keyspace name.
+            table: Table name.
+      - name: --dest-cassandra
+        short-summary: "Destination table"
+        long-summary: |
+            Usage: --dest-cassandra keyspace=XX table=XX'
+            keyspace: Keyspace name.
+            table: Table name.
+      - name: --src-nosql
+        short-summary: "Source container"
+        long-summary: |
+            Usage: --src-nosql database=XX container=XX'
+            database: Database name.
+            container: Container name.
+      - name: --dest-nosql
+        short-summary: "Destination container"
+        long-summary: |
+            Usage: --dest-nosql database=XX container=XX'
+            database: Database name.
+            container: Container name.
+      - name: --src-mongo
+        short-summary: "Source collection"
+        long-summary: |
+            Usage: --src-mongo database=XX collection=XX'
+            database: Database name.
+            collection: Collection name.
+      - name: --dest-mongo
+        short-summary: "Destination collection"
+        long-summary: |
+            Usage: --dest-mongo database=XX collection=XX'
+            database: Database name.
+            collection: Collection name.
+
+    examples:
+      - name: Copy Azure Cosmos DB API for NoSQL container in same account
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "acc1" --dest-account "acc1" --src-nosql database=db1 container=c1 --dest-nosql database=db2 container=c2
+      - name: Copy Azure Cosmos DB API for NoSQL container in different account
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "acc1" --dest-account "acc2" --src-nosql database=db1 container=c1 --dest-nosql database=db2 container=c2
+      - name: Copy Azure Cosmos DB API for Apache Cassandra table
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "acc1" --dest-account "acc1" --src-cassandra keyspace=k1 table=t1 --dest-cassandra keyspace=k2 table=t2
+      - name: Copy Azure Cosmos DB API for MongoDB collection
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "acc1" --dest-account "acc1" --src-mongo database=d1 collection=c1 --dest-mongo database=d2 collection=c2
+"""
+
+helps['cosmosdb copy list'] = """
+    type: command
+    short-summary: "Get a container copy job."
+    examples:
+      - name: List all jobs
+        text: |-
+               az cosmosdb copy list -g "rg1" --account-name "acc1"
+"""
+
+helps['cosmosdb copy show'] = """
+    type: command
+    short-summary: "Get a container copy job."
+    examples:
+      - name: Show details of job j1
+        text: |-
+               az cosmosdb copy show -g "rg1" --account-name "acc1" --job-name "j1"
+"""
+
+helps['cosmosdb copy pause'] = """
+    type: command
+    short-summary: "Pause a container copy job."
+    examples:
+      - name: Pause job j1
+        text: |-
+               az cosmosdb copy pause -g "rg1" --account-name "acc1" --job-name "j1"
+"""
+
+helps['cosmosdb copy resume'] = """
+    type: command
+    short-summary: "Resume a container copy job."
+    examples:
+      - name: Resume job j1
+        text: |-
+               az cosmosdb copy resume -g "rg1" --account-name "acc1" --job-name "j1"
+"""
+
+helps['cosmosdb copy cancel'] = """
+    type: command
+    short-summary: "Cancel a container copy job."
+    examples:
+      - name: Cancel job j1
+        text: |-
+               az cosmosdb copy cancel -g "rg1" --account-name "acc1" --job-name "j1"
+"""
+
 helps['cosmosdb sql container merge'] = """
     type: command
     short-summary: "Merges the partitions of a sql container."

src/cosmosdb-preview/azext_cosmosdb_preview/_params.py

@@ -4,6 +4,7 @@
 # --------------------------------------------------------------------------------------------
 # pylint: disable=line-too-long, too-many-statements
 
+import argparse
 from argcomplete.completers import FilesCompleter
 
 from azext_cosmosdb_preview._validators import (
@@ -28,7 +29,7 @@
     CreatePhysicalPartitionIdListAction)
 
 from azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.models import (
-    ContinuousTier
+    ContinuousTier, DefaultPriorityLevel
 )
 
 from azure.cli.core.util import shell_safe_json_parse
@@ -108,11 +109,16 @@ def load_arguments(self, _):
             'managed-cassandra cluster update',
             'managed-cassandra cluster show',
             'managed-cassandra cluster delete',
+            'managed-cassandra cluster deallocate',
             'managed-cassandra cluster backup list',
             'managed-cassandra cluster backup show']:
         with self.argument_context(scope) as c:
             c.argument('cluster_name', options_list=['--cluster-name', '-c'], help="Cluster Name", required=True)
 
+        # Managed Cassandra Cluster
+        with self.argument_context('managed-cassandra cluster deallocate') as c:
+            c.argument('force', options_list=['--force', '-f'], help="Force to deallocate the cluster", required=False)
+
     # Managed Cassandra Cluster
     for scope in [
             'managed-cassandra cluster create',
@@ -128,6 +134,8 @@ def load_arguments(self, _):
             c.argument('gossip_certificates', help="A list of certificates that should be accepted by on-premise data centers.")
             c.argument('external_seed_nodes', nargs='+', validator=validate_seednodes, help="A list of ip addresses of the seed nodes of on-premise data centers.")
             c.argument('identity_type', options_list=['--identity-type'], arg_type=get_enum_type(['None', 'SystemAssigned']), help="Type of identity used for Customer Managed Disk Key.")
+            c.argument("cluster_type", options_list=['--cluster-type'], help="Type of the cluster, can be Production or NonProduction. If set to Production, operations on cluster might have restrictions.", arg_type=get_enum_type(['Production', 'NonProduction']))
+            c.argument('extensions', nargs='*', help="A set of extensions that will be effective on the cluster. It will replace the entire set of extensions with new set. Use \"\" to remove all. Now Available: cassandra-lucene-index.")
 
     # Managed Cassandra Cluster
     with self.argument_context('managed-cassandra cluster create') as c:
@@ -290,9 +298,9 @@ def load_arguments(self, _):
     # CosmosDB account create with gremlin and tables to restore
     with self.argument_context('cosmosdb create') as c:
         c.argument('account_name', completer=None)
-        c.argument('key_uri', help="The URI of the key vault", is_preview=True)
         c.argument('enable_free_tier', arg_type=get_three_state_flag(), help="If enabled the account is free-tier.", is_preview=True)
-        c.argument('assign_identity', nargs='*', help="Assign system or user assigned identities separated by spaces. Use '[system]' to refer system assigned identity.", is_preview=True)
+        c.argument('assign_identity', nargs='*', help="Assign system or user assigned identities separated by spaces. Use '[system]' to refer system assigned identity.")
+        c.argument('key_uri', help="The URI of the key vault")
         c.argument('is_restore_request', options_list=['--is-restore-request', '-r'], arg_type=get_three_state_flag(), help="Restore from an existing/deleted account.", is_preview=True, arg_group='Restore')
         c.argument('restore_source', help="The restorable-database-account Id of the source account from which the account has to be restored. Required if --is-restore-request is set to true.", is_preview=True, arg_group='Restore')
         c.argument('restore_timestamp', action=UtcDatetimeAction, help="The timestamp to which the account has to be restored to. Required if --is-restore-request is set to true.", is_preview=True, arg_group='Restore')
@@ -324,12 +332,17 @@ def load_arguments(self, _):
             c.argument('backup_retention', type=int, help="the time(in hours) for which each backup is retained (only for accounts with periodic mode backups)", arg_group='Backup Policy')
             c.argument('backup_redundancy', arg_type=get_enum_type(BackupStorageRedundancy), help="The redundancy type of the backup Storage account", arg_group='Backup Policy')
             c.argument('server_version', arg_type=get_enum_type(ServerVersion), help="Valid only for MongoDB accounts.")
-            c.argument('default_identity', help="The primary identity to access key vault in CMK related features. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more.", is_preview=True)
+            c.argument('default_identity', help="The primary identity to access key vault in CMK related features. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more.")
             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('continuous_tier', arg_type=get_enum_type(ContinuousTier), help="The tier of Continuous backup", arg_group='Backup Policy')
             c.argument('enable_materialized_views', options_list=['--enable-materialized-views', '--enable-mv'], arg_type=get_three_state_flag(), help="Flag to enable MaterializedViews on the account.", is_preview=True)
             c.argument('enable_burst_capacity', arg_type=get_three_state_flag(), help="Flag to enable burst capacity on the account.", is_preview=True)
+            c.argument('enable_priority_based_execution', options_list=['--enable-priority-based-execution', '--enable-pbe'], arg_type=get_three_state_flag(), help="Flag to enable priority based execution on the account.", is_preview=True)
+            c.argument('default_priority_level', arg_type=get_enum_type(DefaultPriorityLevel), help="Default Priority Level of Request if not specified.", is_preview=True)
+
+    with self.argument_context('cosmosdb update') as c:
+        c.argument('key_uri', help="The URI of the key vault", is_preview=True)
 
     with self.argument_context('cosmosdb restore') as c:
         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')
@@ -339,8 +352,8 @@ def load_arguments(self, _):
         c.argument('databases_to_restore', nargs='+', action=CreateDatabaseRestoreResource)
         c.argument('gremlin_databases_to_restore', nargs='+', action=CreateGremlinDatabaseRestoreResource, is_preview=True)
         c.argument('tables_to_restore', nargs='+', action=CreateTableRestoreResource, is_preview=True)
-        c.argument('assign_identity', nargs='*', help="Assign system or user assigned identities separated by spaces. Use '[system]' to refer system assigned identity.", is_preview=True)
-        c.argument('default_identity', help="The primary identity to access key vault in CMK related features. e.g. 'FirstPartyIdentity', 'SystemAssignedIdentity' and more.", is_preview=True)
+        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('enable_public_network', options_list=['--enable-public-network', '-e'], arg_type=get_three_state_flag(), help="Enable or disable public network access to server.", is_preview=True)
         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)
 
@@ -437,7 +450,37 @@ def load_arguments(self, _):
             'cosmosdb dts resume',
             'cosmosdb dts cancel']:
         with self.argument_context(scope) as c:
-            c.argument('job_name', options_list=['--job-name', '-n'], help='Name of the Data Transfer Job.')
+            c.argument('job_name', options_list=['--job-name', '-n'], help='Name of the Data Transfer Job.', required=True)
+
+    with self.argument_context('cosmosdb copy create') as c:
+        c.argument('job_name', job_name_type)
+        c.argument('src_account', help='Name of the Azure Cosmos DB source database account.', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')
+        c.argument('dest_account', help='Name of the Azure Cosmos DB destination database account.', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')
+        c.argument('src_cassandra', nargs='+', arg_group='Azure Cosmos DB API for Apache Cassandra table copy', action=AddCassandraTableAction, help='Source Cassandra table details')
+        c.argument('src_mongo', nargs='+', arg_group='Azure Cosmos DB API for MongoDB collection copy', action=AddMongoCollectionAction, help='Source Mongo collection details')
+        c.argument('src_nosql', nargs='+', arg_group='Azure Cosmos DB API for NoSQL container copy', action=AddSqlContainerAction, help='Source NoSql container details')
+        c.argument('dest_cassandra', nargs='+', arg_group='Azure Cosmos DB API for Apache Cassandra table copy', action=AddCassandraTableAction, help='Destination Cassandra table details')
+        c.argument('dest_mongo', nargs='+', arg_group='Azure Cosmos DB API for MongoDB collection copy', action=AddMongoCollectionAction, help='Destination Mongo collection details')
+        c.argument('dest_nosql', nargs='+', arg_group='Azure Cosmos DB API for NoSQL container copy', action=AddSqlContainerAction, help='Destination NoSql container details')
+        c.argument('host_copy_on_src', arg_type=get_three_state_flag(), help=argparse.SUPPRESS)
+        c.argument('worker_count', type=int, help=argparse.SUPPRESS)
+
+    for scope in [
+            'cosmosdb copy list',
+            'cosmosdb copy show',
+            'cosmosdb copy pause',
+            'cosmosdb copy resume',
+            'cosmosdb copy cancel']:
+        with self.argument_context(scope) as c:
+            c.argument('account_name', options_list=["--account-name", "-a"], id_part=None, required=True, help='Azure Cosmos DB account name where the job is created. Use --dest-account value from create job command.')
+
+    for scope in [
+            'cosmosdb copy show',
+            'cosmosdb copy pause',
+            'cosmosdb copy resume',
+            'cosmosdb copy cancel']:
+        with self.argument_context(scope) as c:
+            c.argument('job_name', options_list=['--job-name', '-n'], help='Name of the container copy job.', required=True)
 
     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)')
 

src/cosmosdb-preview/azext_cosmosdb_preview/actions.py

@@ -135,6 +135,10 @@ def __call__(self, parser, namespace, values, option_string=None):
             namespace.source_cassandra_table = cassandra_table
         elif option_string == "--dest-cassandra-table":
             namespace.dest_cassandra_table = cassandra_table
+        elif option_string == "--src-cassandra":
+            namespace.src_cassandra = cassandra_table
+        elif option_string == "--dest-cassandra":
+            namespace.dest_cassandra = cassandra_table
         else:
             namespace.cassandra_table = cassandra_table
 
@@ -174,6 +178,8 @@ def __call__(self, parser, namespace, values, option_string=None):
             namespace.source_mongo = mongo_collection
         elif option_string == "--dest-mongo":
             namespace.dest_mongo = mongo_collection
+        elif option_string == "--src-mongo":
+            namespace.src_mongo = mongo_collection
         else:
             namespace.mongo_collection = mongo_collection
 
@@ -207,14 +213,18 @@ def __call__(self, parser, namespace, values, option_string=None):
         if container_name is None:
             raise CLIError(f'usage error: missing key container in {option_string} component')
 
-        sql_container = CosmosSqlDataTransferDataSourceSink(database_name=database_name, container_name=container_name)
+        nosql_container = CosmosSqlDataTransferDataSourceSink(database_name=database_name, container_name=container_name)
 
         if option_string == "--source-sql-container":
-            namespace.source_sql_container = sql_container
+            namespace.source_sql_container = nosql_container
         elif option_string == "--dest-sql-container":
-            namespace.dest_sql_container = sql_container
+            namespace.dest_sql_container = nosql_container
+        elif option_string == "--src-nosql":
+            namespace.src_nosql = nosql_container
+        elif option_string == "--dest-nosql":
+            namespace.dest_nosql = nosql_container
         else:
-            namespace.sql_container = sql_container
+            namespace.sql_container = nosql_container
 
 
 # pylint: disable=protected-access, too-few-public-methods