[ACR] Fix #33660: az acr network-rule list: Fix null virtualNetworkResourceId and add virtualNetworkSubnetResourceId

src/azure-cli/azure/cli/command_modules/acr/network_rule.py

@@ -3,77 +3,79 @@
 # Licensed under the MIT License. See License.txt in the project root for license information.
 # --------------------------------------------------------------------------------------------
 
-import json
 from knack.util import CLIError
-from azure.cli.core.util import send_raw_request
+from azure.cli.core.commands import LongRunningOperation
 from azure.cli.core.commands.client_factory import get_subscription_id
+from ._client_factory import cf_acr_registries
 from ._utils import validate_premium_registry
 
 
 NETWORK_RULE_NOT_SUPPORTED = 'Network rules are only supported for managed registries in Premium SKU.'
-# TODO: The networkRuleSet property was unintentionally removed from preview APIs.
-# This was not previously exposed, because of multi-api support in the Python SDK.
-# The Python SDK no longer supports multi-api, highlighting this issue.
-# Until the next preview API release is deployed, this code uses direct REST calls to get and update
-# the registry's networkRuleSet. Once the next preview API is released, this code should be updated
-# to use the SDK methods to get and update the networkRuleSet, and the api version should be updated
-# to reflect the latest API version, rather than hard coded, as it currently is below.
-API_VERSION = "2021-08-01-preview"
-
-
-def _get_registry_url(cli_ctx, resource_group_name, registry_name):
-    """Build the REST API URL for a registry."""
-    subscription_id = get_subscription_id(cli_ctx)
-    return (
-        "/subscriptions/{}/resourceGroups/{}/providers/Microsoft.ContainerRegistry/registries/{}?api-version={}"
-        .format(subscription_id, resource_group_name, registry_name, API_VERSION)
-    )
-
 
-def _get_registry(cli_ctx, resource_group_name, registry_name):
-    """Get registry using REST API."""
-    url = _get_registry_url(cli_ctx, resource_group_name, registry_name)
-    response = send_raw_request(cli_ctx, "GET", url)
-    return response.json()
 
+def _get_enum_value(value):
+    return getattr(value, 'value', value)
 
-def _update_registry(cli_ctx, resource_group_name, registry_name, update_payload):
-    """Update registry using REST API (PATCH)."""
-    url = _get_registry_url(cli_ctx, resource_group_name, registry_name)
-    response = send_raw_request(cli_ctx, "PATCH", url, body=json.dumps(update_payload))
-    return response.json()
 
+def _format_virtual_network_rule(rule):
+    """Preserve the legacy field while surfacing the SDK's subnet resource ID."""
+    subnet_id = rule.virtual_network_subnet_resource_id
+    return {
+        'virtualNetworkResourceId': subnet_id,
+        'virtualNetworkSubnetResourceId': subnet_id,
+        'action': _get_enum_value(rule.action) or 'Allow',
+    }
 
-def _format_registry_response(response):
-    """Format the registry REST response for CLI output."""
-    properties = response.get('properties', {})
-    network_rule_set = properties.get('networkRuleSet', {})
 
-    virtual_network_rules = [
-        {'virtualNetworkResourceId': rule.get('id'), 'action': rule.get('action', 'Allow')}
-        for rule in (network_rule_set.get('virtualNetworkRules') or [])
-    ]
-    ip_rules = [
-        {'ipAddressOrRange': rule.get('value') or rule.get('ipAddressOrRange'), 'action': rule.get('action', 'Allow')}
-        for rule in (network_rule_set.get('ipRules') or [])
-    ]
+def _format_registry_response(registry):
+    """Format the registry SDK model for CLI output."""
+    network_rule_set = registry.network_rule_set
+    virtual_network_rules = []
+    ip_rules = []
+    if network_rule_set:
+        virtual_network_rules = [
+            _format_virtual_network_rule(rule)
+            for rule in (network_rule_set.virtual_network_rules or [])
+        ]
+        ip_rules = [
+            {
+                'ipAddressOrRange': rule.ip_address_or_range,
+                'action': _get_enum_value(rule.action) or 'Allow',
+            }
+            for rule in (network_rule_set.ip_rules or [])
+        ]
 
     return {
-        'name': response.get('name'),
-        'provisioningState': properties.get('provisioningState'),
+        'name': registry.name,
+        'provisioningState': _get_enum_value(registry.provisioning_state),
         'networkRuleSet': {
-            'defaultAction': network_rule_set.get('defaultAction'),
+            'defaultAction': _get_enum_value(network_rule_set.default_action) if network_rule_set else None,
             'virtualNetworkRules': virtual_network_rules,
             'ipRules': ip_rules,
         },
     }
 
 
+def _get_network_rule_set(cmd, registry):
+    if registry.network_rule_set:
+        return registry.network_rule_set
+    NetworkRuleSet = cmd.get_models('NetworkRuleSet')
+    return NetworkRuleSet(default_action='Allow')
+
+
+def _update_registry(cmd, resource_group_name, registry_name, network_rule_set):
+    RegistryUpdateParameters = cmd.get_models('RegistryUpdateParameters')
+    parameters = RegistryUpdateParameters(network_rule_set=network_rule_set)
+    client = cf_acr_registries(cmd.cli_ctx)
+    return LongRunningOperation(cmd.cli_ctx)(
+        client.begin_update(resource_group_name, registry_name, parameters)
+    )
+
+
 def acr_network_rule_list(cmd, registry_name, resource_group_name=None):
-    _, resource_group_name = validate_premium_registry(
+    registry, _ = validate_premium_registry(
         cmd, registry_name, resource_group_name, NETWORK_RULE_NOT_SUPPORTED)
 
-    registry = _get_registry(cmd.cli_ctx, resource_group_name, registry_name)
     network_rule_set = _format_registry_response(registry)['networkRuleSet']
     return {'virtualNetworkRules': network_rule_set['virtualNetworkRules'], 'ipRules': network_rule_set['ipRules']}
 
@@ -84,25 +86,26 @@ def acr_network_rule_add(cmd,
                          vnet_name=None,
                          ip_address=None,
                          resource_group_name=None):
-    _, resource_group_name = validate_premium_registry(
+    registry, resource_group_name = validate_premium_registry(
         cmd, registry_name, resource_group_name, NETWORK_RULE_NOT_SUPPORTED)
 
-    registry = _get_registry(cmd.cli_ctx, resource_group_name, registry_name)
-    rules = registry.get('properties', {}).get('networkRuleSet', {})
+    VirtualNetworkRule, IPRule = cmd.get_models('VirtualNetworkRule', 'IPRule')
+    rules = _get_network_rule_set(cmd, registry)
 
     if subnet or vnet_name:
-        virtual_network_rules = list(rules.get('virtualNetworkRules') or [])
+        virtual_network_rules = list(rules.virtual_network_rules or [])
         subnet_id = _validate_subnet(cmd.cli_ctx, subnet, vnet_name, resource_group_name)
-        virtual_network_rules.append({'id': subnet_id, 'action': 'Allow'})
-        rules['virtualNetworkRules'] = virtual_network_rules
+        virtual_network_rules.append(
+            VirtualNetworkRule(virtual_network_subnet_resource_id=subnet_id, action='Allow')
+        )
+        rules.virtual_network_rules = virtual_network_rules
 
     if ip_address:
-        ip_rules = list(rules.get('ipRules') or [])
-        ip_rules.append({'value': ip_address, 'action': 'Allow'})
-        rules['ipRules'] = ip_rules
+        ip_rules = list(rules.ip_rules or [])
+        ip_rules.append(IPRule(ip_address_or_range=ip_address, action='Allow'))
+        rules.ip_rules = ip_rules
 
-    response = _update_registry(cmd.cli_ctx, resource_group_name, registry_name,
-                                {'properties': {'networkRuleSet': rules}})
+    response = _update_registry(cmd, resource_group_name, registry_name, rules)
     return _format_registry_response(response)
 
 
@@ -112,28 +115,27 @@ def acr_network_rule_remove(cmd,
                             vnet_name=None,
                             ip_address=None,
                             resource_group_name=None):
-    _, resource_group_name = validate_premium_registry(
+    registry, resource_group_name = validate_premium_registry(
         cmd, registry_name, resource_group_name, NETWORK_RULE_NOT_SUPPORTED)
 
-    registry = _get_registry(cmd.cli_ctx, resource_group_name, registry_name)
-    rules = registry.get('properties', {}).get('networkRuleSet', {})
+    rules = _get_network_rule_set(cmd, registry)
 
     if subnet or vnet_name:
-        virtual_network_rules = list(rules.get('virtualNetworkRules') or [])
+        virtual_network_rules = list(rules.virtual_network_rules or [])
         subnet_id = _validate_subnet(cmd.cli_ctx, subnet, vnet_name, resource_group_name).lower()
-        rules['virtualNetworkRules'] = [
-            x for x in virtual_network_rules if x.get('id', '').lower() != subnet_id
+        rules.virtual_network_rules = [
+            x for x in virtual_network_rules
+            if (x.virtual_network_subnet_resource_id or '').lower() != subnet_id
         ]
 
     if ip_address:
-        ip_rules = list(rules.get('ipRules') or [])
-        rules['ipRules'] = [
+        ip_rules = list(rules.ip_rules or [])
+        rules.ip_rules = [
             x for x in ip_rules
-            if (x.get('value') or x.get('ipAddressOrRange')) != ip_address
+            if x.ip_address_or_range != ip_address
         ]
 
-    response = _update_registry(cmd.cli_ctx, resource_group_name, registry_name,
-                                {'properties': {'networkRuleSet': rules}})
+    response = _update_registry(cmd, resource_group_name, registry_name, rules)
     return _format_registry_response(response)