Introducing VNET support for Azure CLI - Cosmos DB (Azure#6135)

* initial work

* naming non-optional parameters

* Updating version

* fix style issues

* fixed positional parameters issue

src/command_modules/azure-cli-cosmosdb/HISTORY.rst

@@ -3,6 +3,10 @@
 Release History
 ===============
 
+0.1.21
+++++++
+* Introducing VNET support for Azure CLI - Cosmos DB
+
 0.1.20
 ++++++
 * Minor fixes

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

@@ -11,7 +11,8 @@
 from azure.cli.core.util import shell_safe_json_parse
 
 from azure.cli.command_modules.cosmosdb._validators import (
-    validate_locations, validate_ip_range_filter, validate_failover_policies, validate_capabilities)
+    validate_locations, validate_ip_range_filter, validate_failover_policies, validate_capabilities,
+    validate_virtual_network_rules)
 
 
 def load_arguments(self, _):
@@ -35,6 +36,8 @@ def load_arguments(self, _):
             c.argument('kind', arg_type=get_enum_type(DatabaseAccountKind), help='The type of Cosmos DB database account to create')
             c.argument('enable_automatic_failover', arg_type=get_three_state_flag(), help='Enables automatic failover of the write region in the rare event that the region is unavailable due to an outage. Automatic failover will result in a new write region for the account and is chosen based on the failover priorities configured for the account.')
             c.argument('capabilities', nargs='+', validator=validate_capabilities, help='set custom capabilities on the Cosmos DB database account.')
+            c.argument('enable_virtual_network', arg_type=get_three_state_flag(), help='Enables virtual network on the Cosmos DB database account')
+            c.argument('virtual_network_rules', nargs='+', validator=validate_virtual_network_rules, help='ACL\'s for virtual network')
 
     with self.argument_context('cosmosdb regenerate-key') as c:
         c.argument('key_kind', arg_type=get_enum_type(KeyKind))

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

@@ -10,7 +10,7 @@ def validate_failover_policies(ns):
     fp_dict = []
     for item in ns.failover_policies:
         comps = item.split('=', 1)
-        fp_dict.append(FailoverPolicy(comps[0], int(comps[1])))
+        fp_dict.append(FailoverPolicy(location_name=comps[0], failover_priority=int(comps[1])))
     ns.failover_policies = fp_dict
 
 
@@ -40,3 +40,13 @@ def validate_capabilities(ns):
         for item in ns.capabilities:
             capabilties_list.append(Capability(name=item))
         ns.capabilities = capabilties_list
+
+
+def validate_virtual_network_rules(ns):
+    """ Extracts multiple space-separated virtual network rules """
+    from azure.mgmt.cosmosdb.models.virtual_network_rule import VirtualNetworkRule
+    if ns.virtual_network_rules is not None:
+        virtual_network_rules_list = []
+        for item in ns.virtual_network_rules:
+            virtual_network_rules_list.append(VirtualNetworkRule(id=item))
+        ns.virtual_network_rules = virtual_network_rules_list

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

@@ -51,12 +51,15 @@ def cli_cosmosdb_create(cmd, client,
                         max_interval=5,
                         ip_range_filter=None,
                         enable_automatic_failover=None,
-                        capabilities=None):
+                        capabilities=None,
+                        enable_virtual_network=None,
+                        virtual_network_rules=None):
     """Create a new Azure Cosmos DB database account."""
     consistency_policy = None
     if default_consistency_level is not None:
-        consistency_policy = ConsistencyPolicy(default_consistency_level, max_staleness_prefix,
-                                               max_interval)
+        consistency_policy = ConsistencyPolicy(default_consistency_level=default_consistency_level,
+                                               max_staleness_prefix=max_staleness_prefix,
+                                               max_interval_in_seconds=max_interval)
 
     from azure.mgmt.resource import ResourceManagementClient
     from azure.cli.core.commands.client_factory import get_mgmt_service_client
@@ -68,14 +71,16 @@ def cli_cosmosdb_create(cmd, client,
         locations.append(Location(location_name=resource_group_location, failover_priority=0))
 
     params = DatabaseAccountCreateUpdateParameters(
-        resource_group_location,
-        locations,
+        location=resource_group_location,
+        locations=locations,
         tags=tags,
         kind=kind,
         consistency_policy=consistency_policy,
         ip_range_filter=ip_range_filter,
+        is_virtual_network_filter_enabled=enable_virtual_network,
         enable_automatic_failover=enable_automatic_failover,
-        capabilities=capabilities)
+        capabilities=capabilities,
+        virtual_network_rules=virtual_network_rules)
 
     async_docdb_create = client.create_or_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_create.result()
@@ -93,7 +98,9 @@ def cli_cosmosdb_update(client,
                         max_interval=None,
                         ip_range_filter=None,
                         enable_automatic_failover=None,
-                        capabilities=None):
+                        capabilities=None,
+                        enable_virtual_network=None,
+                        virtual_network_rules=None):
     """Update an existing Azure Cosmos DB database account. """
     existing = client.get(resource_group_name, account_name)
 
@@ -105,7 +112,9 @@ def cli_cosmosdb_update(client,
            max_staleness_prefix is not None or \
            max_interval is not None or \
            ip_range_filter is not None or \
-           enable_automatic_failover is not None:
+           enable_automatic_failover is not None or \
+           enable_virtual_network is not None or \
+           virtual_network_rules is not None:
             raise CLIError("Cannot set capabilities and update properties at the same time. {0}".format(locations))
 
         else:
@@ -122,7 +131,9 @@ def cli_cosmosdb_update(client,
            max_staleness_prefix is None and\
            max_interval is None and\
            ip_range_filter is None and\
-           enable_automatic_failover is None:
+           enable_automatic_failover is None and\
+           enable_virtual_network is None and\
+           virtual_network_rules is None:
 
             async_docdb_create = client.patch(resource_group_name, account_name, tags=tags, capabilities=capabilities)
             docdb_account = async_docdb_create.result()
@@ -146,8 +157,9 @@ def cli_cosmosdb_update(client,
 
     consistency_policy = None
     if update_consistency_policy:
-        consistency_policy = ConsistencyPolicy(default_consistency_level, max_staleness_prefix,
-                                               max_interval)
+        consistency_policy = ConsistencyPolicy(default_consistency_level=default_consistency_level,
+                                               max_staleness_prefix=max_staleness_prefix,
+                                               max_interval_in_seconds=max_interval)
     else:
         consistency_policy = existing.consistency_policy
 
@@ -162,18 +174,26 @@ def cli_cosmosdb_update(client,
     if enable_automatic_failover is None:
         enable_automatic_failover = existing.enable_automatic_failover
 
+    if enable_virtual_network is None:
+        enable_virtual_network = existing.is_virtual_network_filter_enabled
+
+    if virtual_network_rules is None:
+        virtual_network_rules = existing.virtual_network_rules
+
     if tags is None:
         tags = existing.tags
 
     params = DatabaseAccountCreateUpdateParameters(
-        existing.location,
-        locations,
+        location=existing.location,
+        locations=locations,
         tags=tags,
         kind=existing.kind,
         consistency_policy=consistency_policy,
         ip_range_filter=ip_range_filter,
         enable_automatic_failover=enable_automatic_failover,
-        capabilities=existing.capabilities)
+        capabilities=existing.capabilities,
+        is_virtual_network_filter_enabled=enable_virtual_network,
+        virtual_network_rules=virtual_network_rules)
 
     async_docdb_create = client.create_or_update(resource_group_name, account_name, params)
     docdb_account = async_docdb_create.result()

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

@@ -16,7 +16,7 @@
     cmdclass = {}
 
 
-VERSION = "0.1.20"
+VERSION = "0.1.21"
 # The full list of classifiers is available at
 # https://pypi.python.org/pypi?%3Aaction=list_classifiers
 CLASSIFIERS = [
@@ -34,7 +34,7 @@
 
 DEPENDENCIES = [
 
-    'azure-mgmt-cosmosdb==0.3.1',
+    'azure-mgmt-cosmosdb==0.4.0',
     'azure-cli-core',
     'pydocumentdb>=2.0.1'
 ]