[CosmosDB] Adding support for selecting Cosmos DB analytical storage schema type

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

@@ -53,7 +53,7 @@ class ThroughputTypes(str, Enum):
 
 def load_arguments(self, _):
     from knack.arguments import CLIArgumentType
-    from azure.mgmt.cosmosdb.models import KeyKind, DefaultConsistencyLevel, DatabaseAccountKind, TriggerType, TriggerOperation, ServerVersion, NetworkAclBypass, BackupPolicyType
+    from azure.mgmt.cosmosdb.models import KeyKind, DefaultConsistencyLevel, DatabaseAccountKind, TriggerType, TriggerOperation, ServerVersion, NetworkAclBypass, BackupPolicyType, AnalyticalStorageSchemaType
 
     with self.argument_context('cosmosdb') as c:
         c.argument('account_name', arg_type=name_type, help='Name of the Cosmos DB database account', completer=get_resource_name_completion_list('Microsoft.DocumentDb/databaseAccounts'), id_part='name')
@@ -93,6 +93,7 @@ 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('server_version', arg_type=get_enum_type(ServerVersion), help="Valid only for MongoDB accounts.", 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('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')
 
     for scope in ['cosmosdb regenerate-key', 'cosmosdb keys regenerate']:

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

@@ -54,6 +54,7 @@
     PeriodicModeProperties,
     SqlRoleAssignmentCreateUpdateParameters,
     SqlRoleDefinitionCreateUpdateParameters,
+    AnalyticalStorageConfiguration,
     RestoreParameters,
     ContinuousModeBackupPolicy,
     ContinuousBackupRestoreLocation,
@@ -113,6 +114,7 @@ def cli_cosmosdb_create(cmd, client,
                         backup_retention=None,
                         assign_identity=None,
                         default_identity=None,
+                        analytical_storage_schema_type=None,
                         backup_policy_type=None,
                         databases_to_restore=None,
                         is_restore_request=None,
@@ -158,6 +160,7 @@ def cli_cosmosdb_create(cmd, client,
                                     is_restore_request=is_restore_request,
                                     restore_source=restore_source,
                                     restore_timestamp=restore_timestamp_utc,
+                                    analytical_storage_schema_type=analytical_storage_schema_type,
                                     backup_policy_type=backup_policy_type,
                                     backup_interval=backup_interval,
                                     assign_identity=assign_identity,
@@ -196,6 +199,7 @@ def _create_database_account(client,
                              assign_identity=None,
                              default_identity=None,
                              backup_policy_type=None,
+                             analytical_storage_schema_type=None,
                              databases_to_restore=None,
                              is_restore_request=None,
                              restore_source=None,
@@ -251,6 +255,11 @@ def _create_database_account(client,
         )
         backup_policy.periodic_mode_properties = periodic_mode_properties
 
+    analytical_storage_configuration = None
+    if analytical_storage_schema_type is not None:
+        analytical_storage_configuration = AnalyticalStorageConfiguration()
+        analytical_storage_configuration.schema_type = analytical_storage_schema_type
+
     create_mode = CreateMode.restore.value if is_restore_request else CreateMode.default.value
     params = None
     restore_parameters = None
@@ -290,6 +299,7 @@ def _create_database_account(client,
         backup_policy=backup_policy,
         identity=system_assigned_identity,
         default_identity=default_identity,
+        analytical_storage_configuration=analytical_storage_configuration,
         create_mode=create_mode,
         restore_parameters=restore_parameters
     )
@@ -324,6 +334,7 @@ def cli_cosmosdb_update(client,
                         backup_interval=None,
                         backup_retention=None,
                         default_identity=None,
+                        analytical_storage_schema_type=None,
                         backup_policy_type=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
@@ -373,6 +384,11 @@ def cli_cosmosdb_update(client,
         if isinstance(existing.backup_policy, PeriodicModeBackupPolicy):
             backup_policy = ContinuousModeBackupPolicy()
 
+    analytical_storage_configuration = None
+    if analytical_storage_schema_type is not None:
+        analytical_storage_configuration = AnalyticalStorageConfiguration()
+        analytical_storage_configuration.schema_type = analytical_storage_schema_type
+
     params = DatabaseAccountUpdateParameters(
         locations=locations,
         tags=tags,
@@ -390,7 +406,8 @@ def cli_cosmosdb_update(client,
         network_acl_bypass_resource_ids=network_acl_bypass_resource_ids,
         api_properties=api_properties,
         backup_policy=backup_policy,
-        default_identity=default_identity)
+        default_identity=default_identity,
+        analytical_storage_configuration=analytical_storage_configuration)
 
     async_docdb_update = client.begin_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_update.result()