Dts/cross account copy

src/cosmosdb-preview/azext_cosmosdb_preview/_help.py

@@ -778,6 +778,109 @@
                az cosmosdb dts cancel --account-name "ddb1" --job-name "j1" -g "rg1"
 """
 
+helps['cosmosdb copy'] = """
+    type: group
+    short-summary: Manage copy job
+"""
+
+helps['cosmosdb copy create'] = """
+    type: command
+    short-summary: "Creates a Copy Job."
+    parameters:
+      - name: --src-cassandra
+        short-summary: "Source Cassandra table"
+        long-summary: |
+            Usage: --src-cassandra keyspace=XX table=XX'
+            keyspace: Keyspace name of CosmosDB Cassandra.
+            table: Table name of CosmosDB Cassandra.
+      - name: --dest-cassandra
+        short-summary: "Destination Cassandra table"
+        long-summary: |
+            Usage: --dest-cassandra keyspace=XX table=XX'
+            keyspace: Keyspace name of CosmosDB Cassandra.
+            table: Table name of CosmosDB Cassandra.
+      - name: --src-nosql
+        short-summary: "Source NoSql container"
+        long-summary: |
+            Usage: --src-nosql database=XX container=XX'
+            database: Database name of CosmosDB Sql.
+            container: Container name of CosmosDB Sql.
+      - name: --dest-nosql
+        short-summary: "Destination NoSql container"
+        long-summary: |
+            Usage: --dest-nosql database=XX container=XX'
+            database: Database name of CosmosDB NoSql.
+            container: Container name of CosmosDB NoSql.
+      - name: --src-mongo
+        short-summary: "Source Mongo collection"
+        long-summary: |
+            Usage: --src-mongo database=XX collection=XX'
+            database: Database name of CosmosDB Mongo.
+            collection: Collection name of CosmosDB Mongo.
+      - name: --dest-mongo
+        short-summary: "Destination Mongo collection"
+        long-summary: |
+            Usage: --dest-mongo database=XX collection=XX'
+            database: Database name of CosmosDB Mongo.
+            collection: Collection name of CosmosDB Mongo.
+
+    examples:
+      - name: Copy NoSql container
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "db1" --dest-account "db1" --src-nosql database=db1 container=c1 --dest-nosql database=db2 container=c2
+      - name: Copy Cassandra table
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "db1" --dest-account "db1" --src-cassandra keyspace=k1 table=t1 --dest-cassandra keyspace=k2 table=t2
+      - name: Copy Mongo collection
+        text: |-
+          az cosmosdb copy create -g "rg1" --job-name "j1" --src-account "db1" --dest-account "db1" --src-mongo database=d1 collection=c1 --dest-mongo database=d2 collection=c2
+"""
+
+helps['cosmosdb copy list'] = """
+    type: command
+    short-summary: "Get a list of Copy Jobs."
+    examples:
+      - name: List all jobs
+        text: |-
+               az cosmosdb dts list --dest-account "db1" -g "rg1"
+"""
+
+helps['cosmosdb copy show'] = """
+    type: command
+    short-summary: "Get a Copy Job."
+    examples:
+      - name: Show details of job j1
+        text: |-
+               az cosmosdb dts show --dest-account "db1" --job-name "j1" -g "rg1"
+"""
+
+helps['cosmosdb copy pause'] = """
+    type: command
+    short-summary: "Pause a Copy Job."
+    examples:
+      - name: Pause job j1
+        text: |-
+               az cosmosdb dts pause --dest-account "db1" --job-name "j1" -g "rg1"
+"""
+
+helps['cosmosdb copy resume'] = """
+    type: command
+    short-summary: "Resume a Copy Job."
+    examples:
+      - name: Resume job j1
+        text: |-
+               az cosmosdb dts resume --dest-account "db1" --job-name "j1" -g "rg1"
+"""
+
+helps['cosmosdb copy cancel'] = """
+    type: command
+    short-summary: "Cancels a Copy Job."
+    examples:
+      - name: Cancel job j1
+        text: |-
+               az cosmosdb dts cancel --dest-account "db1" --job-name "j1" -g "rg1"
+"""
+
 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
@@ -290,9 +291,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 +325,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', 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 +345,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)
 
@@ -420,7 +426,7 @@ def load_arguments(self, _):
     with self.argument_context('cosmosdb dts') as c:
         c.argument('account_name', account_name_type, id_part=None, help='Name of the CosmosDB database account.')
 
-    job_name_type = CLIArgumentType(options_list=['--job-name', '-n'], help='Name of the Data Transfer Job. A random job name will be generated if not passed.')
+    job_name_type = CLIArgumentType(options_list=['--job-name', '-n'], help='Name of the copy job. A random job name will be generated if not passed.')
     with self.argument_context('cosmosdb dts copy') as c:
         c.argument('job_name', job_name_type)
         c.argument('source_cassandra_table', nargs='+', action=AddCassandraTableAction, help='Source cassandra table')
@@ -431,13 +437,37 @@ def load_arguments(self, _):
         c.argument('dest_sql_container', nargs='+', action=AddSqlContainerAction, help='Destination sql container')
         c.argument('worker_count', type=int, help='Worker count')
 
+    with self.argument_context('cosmosdb copy') as c:
+        c.argument('src_account', help='Name of the CosmosDB source database account.', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')
+        c.argument('dest_account', help='Name of the CosmosDB destination database account.', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')
+
+    with self.argument_context('cosmosdb copy create') as c:
+        c.argument('job_name', job_name_type)
+        c.argument('src_cassandra', nargs='+', arg_group='CosmosDB for Cassandra Table Copy', action=AddCassandraTableAction, help='Source Cassandra table details')
+        c.argument('src_mongo', nargs='+', arg_group='CosmosDB for MongoDB Collection Copy', action=AddMongoCollectionAction, help='Source Mongo collection details')
+        c.argument('src_nosql', nargs='+', arg_group='Cosmos DB for NoSQL Container Copy', action=AddSqlContainerAction, help='Source NoSql container details')
+        c.argument('dest_cassandra', nargs='+', arg_group='CosmosDB for Cassandra Table Copy', action=AddCassandraTableAction, help='Destination Cassandra table details')
+        c.argument('dest_mongo', nargs='+', arg_group='CosmosDB for MongoDB Collection Copy', action=AddMongoCollectionAction, help='Destination Mongo collection details')
+        c.argument('dest_nosql', nargs='+', arg_group='Cosmos DB for NoSQL Container Copy', action=AddSqlContainerAction, help='Destination NoSql container details')
+        c.argument('host_copy_on_src', 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", "--dest-account", "--src-account"], required=True, help='CosmosDB account name where the job is created.')
+
     for scope in [
-            'cosmosdb dts show',
-            'cosmosdb dts pause',
-            'cosmosdb dts resume',
-            'cosmosdb dts cancel']:
+            '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 Data Transfer Job.')
+            c.argument('job_name', help='Name of the Copy Job.', id_part='child_name_1', 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,12 +213,16 @@ 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
 

src/cosmosdb-preview/azext_cosmosdb_preview/commands.py

@@ -197,6 +197,20 @@ def load_command_table(self, _):
         g.command('resume', 'resume')
         g.command('cancel', 'cancel')
 
+    # Data Transfer Service
+    cosmosdb_copy_job = CliCommandType(
+        operations_tmpl='azext_cosmosdb_preview.vendored_sdks.azure_mgmt_cosmosdb.operations._data_transfer_jobs_operations#DataTransferJobsOperations.{}',
+        client_factory=cf_data_transfer_job
+    )
+
+    with self.command_group('cosmosdb copy', cosmosdb_copy_job, client_factory=cf_data_transfer_job, is_preview=True) as g:
+        g.custom_command('create', 'cosmosdb_copy_job')
+        g.command('list', 'list_by_database_account')
+        g.show_command('show', 'get')
+        g.command('pause', 'pause')
+        g.command('resume', 'resume')
+        g.command('cancel', 'cancel')
+
     # Merge partitions for Sql containers
     cosmosdb_sql_sdk = CliCommandType(
         operations_tmpl='azure.mgmt.cosmosdb.operations#SqlResourcesOperations.{}',