VNET extension

src/rdbms_vnet/azext_rdbms/_client_factory.py

@@ -75,5 +75,6 @@ def get_postgresql_management_client(cli_ctx, **_):
 def cf_mysql_virtual_network_rules_operations(cli_ctx, _):
     return get_mysql_management_client(cli_ctx).virtual_network_rules
 
+
 def cf_postgres_virtual_network_rules_operations(cli_ctx, _):
     return get_postgresql_management_client(cli_ctx).virtual_network_rules

src/rdbms_vnet/azext_rdbms/_params.py

@@ -5,7 +5,7 @@
 
 # pylint: disable=line-too-long
 
-from azure.cli.core.commands.parameters import get_three_state_flag,get_resource_name_completion_list, tags_type, get_location_type, get_enum_type   # pylint: disable=line-too-long
+from azure.cli.core.commands.parameters import get_three_state_flag, get_resource_name_completion_list, tags_type, get_location_type, get_enum_type   # pylint: disable=line-too-long
 from azext_rdbms import mysql, postgresql
 from azure.cli.command_modules.sql._validators import validate_subnet
 
@@ -14,27 +14,21 @@
 
 def load_arguments(self, _):    # pylint: disable=too-many-statements
 
-    server_completers = {
-        'mysql': get_resource_name_completion_list('Microsoft.DBForMySQL/servers'),
-        'postgres': get_resource_name_completion_list('Microsoft.DBForPostgreSQL/servers')
-    }
-
     for scope in ['mysql server vnet-rule', 'postgres server vnet-rule']:
-          with self.argument_context(scope) as c:
-              c.argument('server_name', options_list=['--server-name', '-s'])
-              c.argument('virtual_network_rule_name', options_list=['--name', '-n'], id_part='child_name_1', help='The name of the vnet rule.')
-              c.argument('virtual_network_subnet_id', options_list=['--subnet'], help='Name or ID of the subnet that allows access to an Azure Postgres Server. If subnet name is provided, --vnet-name must be provided.')
-              c.argument('ignore_missing_vnet_service_endpoint', options_list=['--ignore-missing-endpoint', '-i'], help='Create vnet rule before virtual network has vnet service endpoint enabled', arg_type=get_three_state_flag())
-
-          with self.argument_context('postgres server vnet-rule create') as c:
-              c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
-
-          with self.argument_context('postgres server vnet-rule update') as c:
-              c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
-
-          with self.argument_context('mysql server vnet-rule create') as c:
-              c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
-
-          with self.argument_context('mysql server vnet-rule update') as c:
-              c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
+        with self.argument_context(scope) as c:
+                c.argument('server_name', options_list=['--server-name', '-s'])
+                c.argument('virtual_network_rule_name', options_list=['--name', '-n'], id_part='child_name_1', help='The name of the vnet rule.')
+                c.argument('virtual_network_subnet_id', options_list=['--subnet'], help='Name or ID of the subnet that allows access to an Azure Postgres Server. If subnet name is provided, --vnet-name must be provided.')
+                c.argument('ignore_missing_vnet_service_endpoint', options_list=['--ignore-missing-endpoint', '-i'], help='Create vnet rule before virtual network has vnet service endpoint enabled', arg_type=get_three_state_flag())
+
+        with self.argument_context('postgres server vnet-rule create') as c:
+                c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
+
+        with self.argument_context('postgres server vnet-rule update') as c:
+                c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
+
+        with self.argument_context('mysql server vnet-rule create') as c:
+                c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)
 
+        with self.argument_context('mysql server vnet-rule update') as c:
+                c.extra('vnet_name', options_list=['--vnet-name'], help='The virtual network name', validator=validate_subnet)

src/rdbms_vnet/azext_rdbms/commands.py

@@ -7,7 +7,7 @@
 
 from azext_rdbms._client_factory import (
     cf_mysql_virtual_network_rules_operations,
-	cf_postgres_virtual_network_rules_operations)
+    cf_postgres_virtual_network_rules_operations)
 
 
 # pylint: disable=too-many-locals, too-many-statements, line-too-long
@@ -27,22 +27,20 @@ def load_command_table(self, _):
         client_factory=cf_postgres_virtual_network_rules_operations
     )
     with self.command_group('mysql server vnet-rule', mysql_vnet_sdk) as g:
-         g.command('create', 'create_or_update')
-         g.command('delete', 'delete', confirmation=True)
-         g.command('show', 'get')
-         g.command('list', 'list_by_server')
-         #g.generic_update_command('update')
-         g.generic_update_command('update', getter_name='_custom_vnet_update_get', getter_type=rdbms_custom,
+        g.command('create', 'create_or_update')
+        g.command('delete', 'delete', confirmation=True)
+        g.command('show', 'get')
+        g.command('list', 'list_by_server')
+        # g.generic_update_command('update')
+        g.generic_update_command('update', getter_name='_custom_vnet_update_get', getter_type=rdbms_custom,
                                  setter_name='_custom_vnet_update_set', setter_type=rdbms_custom, setter_arg_name='parameters')
 
     with self.command_group('postgres server vnet-rule', postgres_vnet_sdk) as g:
-         g.command('create', 'create_or_update')
-         g.command('delete', 'delete', confirmation=True)
-         g.command('show', 'get')
-         g.command('list', 'list_by_server')
-         #g.generic_update_command('update')
-         g.generic_update_command('update',
-                       getter_name='_custom_vnet_update_get', getter_type=rdbms_custom,
-                       setter_name='_custom_vnet_update_set', setter_type=rdbms_custom, setter_arg_name='parameters')
-                       #custom_func_name='_custom_vnet_update', custom_func_type=rdbms_custom)
-         #g.custom_command('update', '_custom_vnet_update')
+        g.command('create', 'create_or_update')
+        g.command('delete', 'delete', confirmation=True)
+        g.command('show', 'get')
+        g.command('list', 'list_by_server')
+        # g.generic_update_command('update')
+        g.generic_update_command('update',
+                                 getter_name='_custom_vnet_update_get', getter_type=rdbms_custom,
+                                 setter_name='_custom_vnet_update_set', setter_type=rdbms_custom, setter_arg_name='parameters')

src/rdbms_vnet/azext_rdbms/custom.py

@@ -11,9 +11,9 @@
 from ._client_factory import get_mysql_management_client, get_postgresql_management_client
 
 
-
 def _custom_vnet_update_get(client, resource_group_name, server_name, virtual_network_rule_name):
     return client.get(resource_group_name, server_name, virtual_network_rule_name)
 
+
 def _custom_vnet_update_set(client, resource_group_name, server_name, virtual_network_rule_name, virtual_network_subnet_id, ignore_missing_vnet_service_endpoint=None, parameters=None):
     return client.create_or_update(resource_group_name, server_name, virtual_network_rule_name, virtual_network_subnet_id, ignore_missing_vnet_service_endpoint)

src/rdbms_vnet/azext_rdbms/tests/test_rdbms_commands.py

@@ -354,50 +354,50 @@ def _test_vnet_firewall_mgmt(self, resource_group, server, database_engine):
 
         subnet_name_1 = 'clitestsubnet1'
         subnet_prefix_1 = '10.0.0.0/24'
-        
+
         subnet_name_2 = 'clitestsubnet2'
         subnet_prefix_2 = '10.0.1.0/24'
 
-        # pre create the dependent resources here 
+        # pre create the dependent resources here
         # create vnet and subnet
         execute(self.cli_ctx, 'az network vnet create -n {} -g {} -l {} '
-                 '--address-prefix {} --subnet-name {} --subnet-prefix {}'.format(vnet_name, resource_group, location, address_prefix, subnet_name_1,
-                         subnet_prefix_1))
+                '--address-prefix {} --subnet-name {} --subnet-prefix {}'.format(vnet_name, resource_group, location, address_prefix, subnet_name_1,
+                                                                                 subnet_prefix_1))
         # add one more subnet
         execute(self.cli_ctx, 'az network vnet subnet create --vnet-name {} -g {} '
-                 '--address-prefix {} -n {}'.format(vnet_name, resource_group, subnet_prefix_2, subnet_name_2)) 
+                '--address-prefix {} -n {}'.format(vnet_name, resource_group, subnet_prefix_2, subnet_name_2))
         # test vnet-rule create
         self.cmd('{} server vnet-rule create -n {} -g {} -s {} '
                  '--vnet-name {} --subnet {} --ignore-missing-endpoint {}'
                  .format(database_engine, vnet_firewall_rule_1, resource_group, server,
                          vnet_name, subnet_name_1, ignore_missing_endpoint),
                  checks=[
-                        JMESPathCheck('name', vnet_firewall_rule_1),
-                        JMESPathCheck('resourceGroup', resource_group),
-                        JMESPathCheck('state', 'Ready')])
+                         JMESPathCheck('name', vnet_firewall_rule_1),
+                         JMESPathCheck('resourceGroup', resource_group),
+                         JMESPathCheck('state', 'Ready')])
 
         # test vnet-rule show
         self.cmd('{} server vnet-rule show -n {} -g {} -s {}'
                  .format(database_engine, vnet_firewall_rule_1, resource_group, server),
                  checks=[
-                        JMESPathCheck('name', vnet_firewall_rule_1),
-                        JMESPathCheck('resourceGroup', resource_group),
-                        JMESPathCheck('state', 'Ready')])  
+                         JMESPathCheck('name', vnet_firewall_rule_1),
+                         JMESPathCheck('resourceGroup', resource_group),
+                         JMESPathCheck('state', 'Ready')])
 
         # test create one more vnet rule .
         self.cmd('{} server vnet-rule create -n {} -g {} -s {} '
                  '--vnet-name {} --subnet {} --ignore-missing-endpoint {}'
                  .format(database_engine, vnet_firewall_rule_2, resource_group, server,
                          vnet_name, subnet_name_2, ignore_missing_endpoint),
                  checks=[
-                        JMESPathCheck('name', vnet_firewall_rule_2),
-                        JMESPathCheck('resourceGroup', resource_group),
-                        JMESPathCheck('state', 'Ready')])
+                         JMESPathCheck('name', vnet_firewall_rule_2),
+                         JMESPathCheck('resourceGroup', resource_group),
+                         JMESPathCheck('state', 'Ready')])
 
         # test vnet-rule list
         self.cmd('{} server vnet-rule list -g {} -s {}'
                  .format(database_engine, resource_group, server),
-                 checks=[JMESPathCheck('length(@)', 2)])  
+                 checks=[JMESPathCheck('length(@)', 2)])
 
         self.cmd('{} server vnet-rule delete --name {} -g {} --server {} --yes'
                  .format(database_engine, vnet_firewall_rule_1, resource_group, server), checks=NoneCheck())
@@ -406,8 +406,7 @@ def _test_vnet_firewall_mgmt(self, resource_group, server, database_engine):
         self.cmd('{} server vnet-rule delete -n {} -g {} -s {} --yes'
                  .format(database_engine, vnet_firewall_rule_2, resource_group, server), checks=NoneCheck())
         self.cmd('{} server vnet-rule list -g {} --server {}'
-                 .format(database_engine, resource_group, server), checks=[NoneCheck()])        
-
+                 .format(database_engine, resource_group, server), checks=[NoneCheck()])
 
     def _test_db_mgmt(self, resource_group, server, database_engine):
         self.cmd('{} db list -g {} -s {}'.format(database_engine, resource_group, server),