Merge pull request #5 from Achint-Agrawal/PBEStable

Added support for priority based execution

Author: pjohari-ms

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

@@ -35,7 +35,7 @@
 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\\\\"/?\\"}]}"
 """
@@ -131,6 +131,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_priority_based_execution', 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")
 

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

@@ -143,7 +143,9 @@ def cli_cosmosdb_create(cmd,
                         enable_partition_merge=None,
                         enable_burst_capacity=None,
                         enable_prpp_autoscale=None,
-                        minimal_tls_version=None):
+                        minimal_tls_version=None,
+                        enable_priority_based_execution=None,
+                        default_priority_level=None):
     """Create a new Azure Cosmos DB database account."""
 
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
@@ -199,7 +201,9 @@ def cli_cosmosdb_create(cmd,
                                     enable_partition_merge=enable_partition_merge,
                                     enable_burst_capacity=enable_burst_capacity,
                                     enable_prpp_autoscale=enable_prpp_autoscale,
-                                    minimal_tls_version=minimal_tls_version)
+                                    minimal_tls_version=minimal_tls_version,
+                                    enable_priority_based_execution=enable_priority_based_execution,
+                                    default_priority_level=default_priority_level)
 
 
 # pylint: disable=too-many-statements
@@ -247,8 +251,11 @@ def _create_database_account(client,
                              enable_prpp_autoscale=None,
                              minimal_tls_version=None,
                              disable_ttl=None,
+                             enable_priority_based_execution=None,
+                             default_priority_level=None,
                              source_backup_location=None):
 
+
     consistency_policy = None
     if default_consistency_level is not None:
         consistency_policy = ConsistencyPolicy(default_consistency_level=default_consistency_level,
@@ -386,6 +393,8 @@ def _create_database_account(client,
         enable_burst_capacity=enable_burst_capacity,
         enable_per_region_per_partition_autoscale=enable_prpp_autoscale,
         minimal_tls_version=minimal_tls_version,
+        enable_priority_based_execution=enable_priority_based_execution,
+        default_priority_level=default_priority_level,
         source_backup_location=source_backup_location
     )
 
@@ -427,7 +436,9 @@ def cli_cosmosdb_update(client,
                         enable_partition_merge=None,
                         enable_burst_capacity=None,
                         enable_prpp_autoscale=None,
-                        minimal_tls_version=None):
+                        minimal_tls_version=None,
+                        enable_priority_based_execution=None,
+                        default_priority_level=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
 
@@ -525,7 +536,9 @@ def cli_cosmosdb_update(client,
         enable_partition_merge=enable_partition_merge,
         enable_burst_capacity=enable_burst_capacity,
         enable_per_region_per_partition_autoscale=enable_prpp_autoscale,
-        minimal_tls_version=minimal_tls_version)
+        minimal_tls_version=minimal_tls_version,
+        enable_priority_based_execution=enable_priority_based_execution,
+        default_priority_level=default_priority_level)
 
     async_docdb_update = client.begin_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_update.result()