[CosmosDB] az cosmosdb create/update: Add --enable-prpp-autoscale to enable/disable burst capacity feature (#30415)

* adding vector policy

* CLI changes - prpp autoscale

* Testing prpp autoscale

* adding tests

* Bump up the dependency

* Recordings

* Merge recordings with master

* Merge recordings with Vipul's changes

* Recordings

* Updated recordings

* Test recording

* Re-recording tests, again

* Changed vector embeddings flag due to linter max param length

* The test

* External modules - recordings

---------

Co-authored-by: Vipul Vishal <[email protected]>
Co-authored-by: Shorya Kumar <[email protected]>

Author: 3 people

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

@@ -193,6 +193,12 @@
       Usage:    --enable-burst-capacity true
       Default:  false
       The accepted values for the enable-burst-capacity are true and false.
+  - name: --enable-prpp-autoscale
+    short-summary: Flag to Enable/Disable burst capacity feature
+    long-summary: |
+      Usage:    --enable-prpp-autoscale true
+      Default:  false
+      The accepted values for the --enable-prpp-autoscale are true and false.
 examples:
   - name: Creates a new Azure Cosmos DB database account. (autogenerated)
     text: az cosmosdb create --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup --subscription MySubscription
@@ -957,6 +963,12 @@
       Usage:    --enable-burst-capacity true
       Default:  false
       The accepted values for the enable-burst-capacity are true and false.
+  - name: --enable-prpp-autoscale
+    short-summary: Flag to Enable/Disable burst capacity feature
+    long-summary: |
+      Usage:    --enable-prpp-autoscale true
+      Default:  false
+      The accepted values for the --enable-prpp-autoscale are true and false.
 examples:
   - name: Update an Azure Cosmos DB database account. (autogenerated)
     text: az cosmosdb update --capabilities EnableGremlin --name MyCosmosDBDatabaseAccount --resource-group MyResourceGroup

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

@@ -37,7 +37,12 @@
 from azure.mgmt.cosmosdb.models import (
     ContinuousTier, MinimalTlsVersion)
 
-SQL_GREMLIN_INDEXING_POLICY_EXAMPLE = """--idx "{\\"indexingMode\\": \\"consistent\\", \\"automatic\\": true, \\"includedPaths\\": [{\\"path\\": \\"/*\\"}], \\"excludedPaths\\": [{ \\"path\\": \\"/headquarters/employees/?\\"}, { \\"path\\": \\"/\\\\"_etag\\\\"/?\\"}]}"
+GREMLIN_INDEXING_POLICY_EXAMPLE = """--idx "{\\"indexingMode\\": \\"consistent\\", \\"automatic\\": true, \\"includedPaths\\": [{\\"path\\": \\"/*\\"}], \\"excludedPaths\\": [{ \\"path\\": \\"/headquarters/employees/?\\"}, { \\"path\\": \\"/\\\\"_etag\\\\"/?\\"}]}"
+"""
+SQL_INDEXING_POLICY_EXAMPLE = """--idx "{\\"indexingMode\\": \\"consistent\\", \\"automatic\\": true, \\"includedPaths\\": [{\\"path\\": \\"/*\\"}], \\"excludedPaths\\": [{ \\"path\\": \\"/headquarters/employees/?\\"}, { \\"path\\": \\"/\\\\"_etag\\\\"/?\\"}],\\"vectorIndexes\\": [{\\"path\\": \\"/vector1\\",\\"type\\": \\"flat\\"}]}"
+"""
+
+SQL_VECTOR_EMBEDDING_POLICY_EXAMPLE = """--vector-embeddings "{\\"vectorEmbeddings\\": [{\\"path\\": \\"/vector1\\", \\"dataType\\": \\"float32\\", \\"dimensions\\": 2, \\"distanceFunction\\": \\"dotproduct\\" }]}"
 """
 
 SQL_UNIQUE_KEY_POLICY_EXAMPLE = """--unique-key-policy "{\\"uniqueKeys\\": [{\\"paths\\": [\\"/path/to/key1\\"]}, {\\"paths\\": [\\"/path/to/key2\\"]}]}"
@@ -90,6 +95,7 @@ def load_arguments(self, _):
         c.argument('tables_to_restore', nargs='+', action=CreateTableRestoreResource, arg_group='Restore')
         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('key_uri', help="The URI of the key vault")
 
     for scope in ['cosmosdb create', 'cosmosdb update']:
@@ -121,6 +127,7 @@ def load_arguments(self, _):
             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('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('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")
 
@@ -156,6 +163,7 @@ def load_arguments(self, _):
         c.argument('throughput', type=int, help='Offer Throughput (RU/s)')
         c.argument('partition_key_path', help='Partition Key Path, e.g., \'/properties/name\'')
         c.argument('client_encryption_policy', options_list=['--cep'], type=shell_safe_json_parse, completer=FilesCompleter(), validator=validate_client_encryption_policy, help='Client Encryption Policy, you can enter it as a string or as a file, e.g., --cep @policy-file.json or ' + SQL_CLIENT_ENCRYPTION_POLICY_EXAMPLE)
+        c.argument('vector_embedding_policy', options_list=['--vector-embeddings'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Vector Encryption Policy, you can enter it as a string or as a file, e.g., --vector-embeddings @policy-file.json or ' + SQL_VECTOR_EMBEDDING_POLICY_EXAMPLE)
         c.argument('indexing_policy', type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --indexing-policy @policy-file.json)')
         c.argument('default_ttl', type=int, help='Default TTL. Provide 0 to disable.')
 
@@ -208,8 +216,9 @@ def load_arguments(self, _):
         c.argument('partition_key_path', options_list=['--partition-key-path', '-p'], help='Partition Key Path, e.g., \'/address/zipcode\'')
         c.argument('partition_key_version', type=int, options_list=['--partition-key-version'], help='The version of partition key.')
         c.argument('default_ttl', options_list=['--ttl'], type=int, help='Default TTL. If the value is missing or set to "-1", items don’t expire. If the value is set to "n", items will expire "n" seconds after last modified time.')
-        c.argument('indexing_policy', options_list=['--idx'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --idx @policy-file.json or ' + SQL_GREMLIN_INDEXING_POLICY_EXAMPLE)
+        c.argument('indexing_policy', options_list=['--idx'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --idx @policy-file.json or ' + SQL_INDEXING_POLICY_EXAMPLE)
         c.argument('client_encryption_policy', options_list=['--cep'], type=shell_safe_json_parse, completer=FilesCompleter(), validator=validate_client_encryption_policy, help='Client Encryption Policy, you can enter it as a string or as a file, e.g., --cep @policy-file.json or ' + SQL_CLIENT_ENCRYPTION_POLICY_EXAMPLE)
+        c.argument('vector_embedding_policy', options_list=['--vector-embeddings'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Vector Encryption Policy, you can enter it as a string or as a file, e.g., --vector-embeddings @policy-file.json or ' + SQL_VECTOR_EMBEDDING_POLICY_EXAMPLE)
         c.argument('unique_key_policy', options_list=['--unique-key-policy', '-u'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Unique Key Policy, you can enter it as a string or as a file, e.g., --unique-key-policy @policy-file.json or ' + SQL_UNIQUE_KEY_POLICY_EXAMPLE)
         c.argument('conflict_resolution_policy', options_list=['--conflict-resolution-policy', '-c'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Conflict Resolution Policy, you can enter it as a string or as a file, e.g., --conflict-resolution-policy @policy-file.json or ' + SQL_GREMLIN_CONFLICT_RESOLUTION_POLICY_EXAMPLE)
         c.argument('max_throughput', max_throughput_type)
@@ -289,7 +298,7 @@ def load_arguments(self, _):
         c.argument('graph_name', options_list=['--name', '-n'], help="Graph name")
         c.argument('partition_key_path', options_list=['--partition-key-path', '-p'], help='Partition Key Path, e.g., \'/address/zipcode\'')
         c.argument('default_ttl', options_list=['--ttl'], type=int, help='Default TTL. If the value is missing or set to "-1", items don’t expire. If the value is set to "n", items will expire "n" seconds after last modified time.')
-        c.argument('indexing_policy', options_list=['--idx'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --idx @policy-file.json or ' + SQL_GREMLIN_INDEXING_POLICY_EXAMPLE)
+        c.argument('indexing_policy', options_list=['--idx'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Indexing Policy, you can enter it as a string or as a file, e.g., --idx @policy-file.json or ' + GREMLIN_INDEXING_POLICY_EXAMPLE)
         c.argument('conflict_resolution_policy', options_list=['--conflict-resolution-policy', '-c'], type=shell_safe_json_parse, completer=FilesCompleter(), help='Conflict Resolution Policy, you can enter it as a string or as a file, e.g., --conflict-resolution-policy @policy-file.json or ' + SQL_GREMLIN_CONFLICT_RESOLUTION_POLICY_EXAMPLE)
         c.argument('max_throughput', max_throughput_type)
         c.argument('throughput', help='The throughput of Gremlin graph (RU/s). Default value is 400. Omit this parameter if the database has shared throughput unless the graph should have dedicated throughput.')

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

@@ -141,6 +141,7 @@ def cli_cosmosdb_create(cmd,
                         restore_timestamp=None,
                         enable_partition_merge=None,
                         enable_burst_capacity=None,
+                        enable_prpp_autoscale=None,
                         minimal_tls_version=None):
     """Create a new Azure Cosmos DB database account."""
 
@@ -196,6 +197,7 @@ def cli_cosmosdb_create(cmd,
                                     arm_location=resource_group_location,
                                     enable_partition_merge=enable_partition_merge,
                                     enable_burst_capacity=enable_burst_capacity,
+                                    enable_prpp_autoscale=enable_prpp_autoscale,
                                     minimal_tls_version=minimal_tls_version)
 
 
@@ -241,6 +243,7 @@ def _create_database_account(client,
                              arm_location=None,
                              enable_partition_merge=None,
                              enable_burst_capacity=None,
+                             enable_prpp_autoscale=None,
                              minimal_tls_version=None,
                              disable_ttl=None):
 
@@ -379,6 +382,7 @@ def _create_database_account(client,
         restore_parameters=restore_parameters,
         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
     )
 
@@ -419,6 +423,7 @@ def cli_cosmosdb_update(client,
                         continuous_tier=None,
                         enable_partition_merge=None,
                         enable_burst_capacity=None,
+                        enable_prpp_autoscale=None,
                         minimal_tls_version=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
@@ -516,6 +521,7 @@ def cli_cosmosdb_update(client,
         analytical_storage_configuration=analytical_storage_configuration,
         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)
 
     async_docdb_update = client.begin_update(resource_group_name, account_name, params)
@@ -599,9 +605,10 @@ def _populate_sql_container_definition(sql_container_resource,
                                        client_encryption_policy,
                                        partition_key_version,
                                        conflict_resolution_policy,
-                                       analytical_storage_ttl):
+                                       analytical_storage_ttl,
+                                       vector_embedding_policy):
     if all(arg is None for arg in
-           [partition_key_path, partition_key_version, default_ttl, indexing_policy, unique_key_policy, client_encryption_policy, conflict_resolution_policy, analytical_storage_ttl]):
+           [partition_key_path, partition_key_version, default_ttl, indexing_policy, unique_key_policy, client_encryption_policy, conflict_resolution_policy, analytical_storage_ttl, vector_embedding_policy]):
         return False
 
     if partition_key_path is not None:
@@ -630,6 +637,9 @@ def _populate_sql_container_definition(sql_container_resource,
     if analytical_storage_ttl is not None:
         sql_container_resource.analytical_storage_ttl = analytical_storage_ttl
 
+    if vector_embedding_policy is not None:
+        sql_container_resource.vector_embedding_policy = vector_embedding_policy
+
     return True
 
 
@@ -647,7 +657,8 @@ def cli_cosmosdb_sql_container_create(client,
                                       max_throughput=None,
                                       unique_key_policy=None,
                                       conflict_resolution_policy=None,
-                                      analytical_storage_ttl=None):
+                                      analytical_storage_ttl=None,
+                                      vector_embedding_policy=None):
     """Creates an Azure Cosmos DB SQL container """
     sql_container_resource = SqlContainerResource(id=container_name)
 
@@ -659,7 +670,8 @@ def cli_cosmosdb_sql_container_create(client,
                                        client_encryption_policy,
                                        partition_key_version,
                                        conflict_resolution_policy,
-                                       analytical_storage_ttl)
+                                       analytical_storage_ttl,
+                                       vector_embedding_policy)
 
     options = _get_options(throughput, max_throughput)
 
@@ -681,7 +693,8 @@ def cli_cosmosdb_sql_container_update(client,
                                       container_name,
                                       default_ttl=None,
                                       indexing_policy=None,
-                                      analytical_storage_ttl=None):
+                                      analytical_storage_ttl=None,
+                                      vector_embedding_policy=None):
     """Updates an Azure Cosmos DB SQL container """
     logger.debug('reading SQL container')
     sql_container = client.get_sql_container(resource_group_name, account_name, database_name, container_name)
@@ -692,6 +705,7 @@ def cli_cosmosdb_sql_container_update(client,
     sql_container_resource.default_ttl = sql_container.resource.default_ttl
     sql_container_resource.unique_key_policy = sql_container.resource.unique_key_policy
     sql_container_resource.conflict_resolution_policy = sql_container.resource.conflict_resolution_policy
+    sql_container_resource.vector_embedding_policy = sql_container.resource.vector_embedding_policy
 
     # client encryption policy is immutable
     sql_container_resource.client_encryption_policy = sql_container.resource.client_encryption_policy
@@ -704,7 +718,8 @@ def cli_cosmosdb_sql_container_update(client,
                                           None,
                                           None,
                                           None,
-                                          analytical_storage_ttl):
+                                          analytical_storage_ttl,
+                                          vector_embedding_policy):
         logger.debug('replacing SQL container')
 
     sql_container_create_update_resource = SqlContainerCreateUpdateParameters(
@@ -2145,7 +2160,8 @@ def _populate_collection_definition(collection,
                                     partition_key_path=None,
                                     default_ttl=None,
                                     indexing_policy=None,
-                                    client_encryption_policy=None):
+                                    client_encryption_policy=None,
+                                    vector_embedding_policy=None):
     if all(arg is None for arg in [partition_key_path, default_ttl, indexing_policy]):
         return False
 
@@ -2166,6 +2182,9 @@ def _populate_collection_definition(collection,
     if client_encryption_policy is not None:
         collection['clientEncryptionPolicy'] = client_encryption_policy
 
+    if vector_embedding_policy is not None:
+        collection['vectorIndexingPolicy'] = vector_embedding_policy
+
     return True
 
 
@@ -2176,7 +2195,8 @@ def cli_cosmosdb_collection_create(client,
                                    partition_key_path=None,
                                    default_ttl=None,
                                    indexing_policy=DEFAULT_INDEXING_POLICY,
-                                   client_encryption_policy=None):
+                                   client_encryption_policy=None,
+                                   vector_embedding_policy=None):
     """Creates an Azure Cosmos DB collection """
     collection = {'id': collection_id}
 
@@ -2188,7 +2208,8 @@ def cli_cosmosdb_collection_create(client,
                                     partition_key_path,
                                     default_ttl,
                                     indexing_policy,
-                                    client_encryption_policy)
+                                    client_encryption_policy,
+                                    vector_embedding_policy)
 
     created_collection = client.CreateContainer(_get_database_link(database_id), collection,
                                                 options)