[AppConfig] az appconfig kv: Add snapshot reference support

src/azure-cli/azure/cli/command_modules/appconfig/_constants.py

@@ -63,6 +63,13 @@ class KeyVaultConstants:
     KEYVAULT_CONTENT_TYPE = "application/vnd.microsoft.appconfig.keyvaultref+json;charset=utf-8"
 
 
+class SnapshotReferenceConstants:
+    SNAPSHOT_REFERENCE_CONTENT_TYPE = (
+        'application/json; profile="https://azconfig.io/mime-profiles/snapshot-ref"; charset=utf-8'
+    )
+    SNAPSHOT_NAME_KEY = "snapshot_name"
+
+
 class AIConfigConstants:
     AI_CHAT_COMPLETION_CONTENT_TYPE = "application/vnd.microsoft.appconfig.aichatcompletion+json;charset=utf-8"
 

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

@@ -242,6 +242,8 @@
     text: az appconfig kv list -n MyAppConfiguration --tags tag1=
   - name: List all key-values with tag name "tag1" with null value
     text: az appconfig kv list -n MyAppConfiguration --tags tag1=\\0
+  - name: Resolve a snapshot reference and list all key-values from the referenced snapshot.
+    text: az appconfig kv list -n MyAppConfiguration --key MySnapshotRef --resolve-snapshot-references
 """
 
 helps['appconfig kv lock'] = """
@@ -296,6 +298,18 @@
     text: az appconfig kv set-keyvault --connection-string Endpoint=https://contoso.azconfig.io;Id=xxx;Secret=xxx --key HostSecret --secret-identifier https://contoso.vault.azure.net/Secrets/DummySecret --tags tag1=value1 tag2=value2
 """
 
+helps['appconfig kv set-snapshot-reference'] = """
+type: command
+short-summary: Set a snapshot reference.
+examples:
+  - name: Set a snapshot reference with label MyLabel.
+    text: az appconfig kv set-snapshot-reference -n MyAppConfiguration --key MySnapshotRef --label MyLabel --snapshot-name MySnapshot
+  - name: Set a snapshot reference using login-based authentication.
+    text: az appconfig kv set-snapshot-reference --endpoint https://myappconfiguration.azconfig.io --key MySnapshotRef --snapshot-name MySnapshot --auth-mode login
+  - name: Set a snapshot reference with tags using connection string.
+    text: az appconfig kv set-snapshot-reference --connection-string Endpoint=https://contoso.azconfig.io;Id=xxx;Secret=xxx --key MySnapshotRef --snapshot-name MySnapshot --tags tag1=value1 tag2=value2
+"""
+
 helps['appconfig kv show'] = """
 type: command
 short-summary: Show all attributes of a key-value.

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

@@ -24,7 +24,7 @@
                           validate_feature_query_fields, validate_filter_parameters,
                           validate_separator, validate_secret_identifier,
                           validate_key, validate_feature, validate_feature_key,
-                          validate_identity, validate_auth_mode,
+                          validate_identity, validate_auth_mode, validate_snapshot_reference,
                           validate_resolve_keyvault, validate_export_profile, validate_import_profile,
                           validate_strict_import, validate_export_as_reference, validate_snapshot_filters,
                           validate_snapshot_export, validate_snapshot_import, validate_tag_filters,
@@ -326,6 +326,12 @@ def load_arguments(self, _):
         c.argument('tags', arg_type=tags_type)
         c.argument('secret_identifier', validator=validate_secret_identifier, help="ID of the Key Vault object. Can be found using 'az keyvault {collection} show' command, where collection is key, secret or certificate. To set reference to the latest version of your secret, remove version information from secret identifier.")
 
+    with self.argument_context('appconfig kv set-snapshot-reference') as c:
+        c.argument('key', validator=validate_key, help="Key to be set. Key cannot be a '.' or '..', or contain the '%%' character.")
+        c.argument('label', help="If no label specified, set the key with null label by default")
+        c.argument('tags', arg_type=tags_type)
+        c.argument('snapshot_name', validator=validate_snapshot_reference, help='Name of the snapshot to reference. This is required.')
+
     with self.argument_context('appconfig kv delete') as c:
         c.argument('key', validator=validate_key, help='Support star sign as filters, for instance * means all key and abc* means keys with abc as prefix.')
         c.argument('label', help="If no label specified, delete entry with null label. Support star sign as filters, for instance * means all label and abc* means labels with abc as prefix.")
@@ -341,6 +347,7 @@ def load_arguments(self, _):
         c.argument('tags', arg_type=tags_arg_type, help="If no tags are specified, return all key-values with any tags. Support space-separated tags: key[=value] [key[=value] ...].")
         c.argument('snapshot', help="List all keys in a given snapshot of the App Configuration store. If no snapshot is specified, the keys currently in the store are listed.")
         c.argument('resolve_keyvault', arg_type=get_three_state_flag(), help="Resolve the content of key vault reference. This argument should NOT be specified along with --fields. Instead use --query for customized query.")
+        c.argument('resolve_snapshot_references', arg_type=get_three_state_flag(), help="Resolve snapshot references and return the referenced snapshots' key-values.")
 
     with self.argument_context('appconfig kv restore') as c:
         c.argument('key', help='If no key specified, restore all keys by default. Support star sign as filters, for instance abc* means keys with abc as prefix.')

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

@@ -308,6 +308,11 @@ def validate_key(namespace):
         raise InvalidArgumentValueError("Key is invalid. Key cannot be a '.' or '..', or contain the '%' character.")
 
 
+def validate_snapshot_reference(namespace):
+    if not namespace.snapshot_name or str(namespace.snapshot_name).isspace():
+        raise RequiredArgumentMissingError("--snapshot-name is required and cannot be empty.")
+
+
 def validate_resolve_keyvault(namespace):
     if namespace.resolve_keyvault:
         identifier = getattr(namespace, 'destination', None)

src/azure-cli/azure/cli/command_modules/appconfig/commands.py

@@ -108,6 +108,7 @@ def get_custom_sdk(custom_module, client_factory, table_transformer):
         g.command('import', 'import_config')
         g.command('export', 'export_config')
         g.command('set-keyvault', 'set_keyvault')
+        g.command('set-snapshot-reference', 'set_snapshot_reference')
 
     # FeatureManagement Commands
     with self.command_group('appconfig feature',

src/azure-cli/azure/cli/command_modules/appconfig/keyvalue.py

@@ -50,7 +50,8 @@
                          SearchFilterOptions, StatusCodes,
                          ImportExportProfiles, CompareFieldsMap,
                          JsonDiff, ImportMode,
-                         AIConfigConstants, HttpHeaders)
+                         AIConfigConstants, HttpHeaders,
+                         SnapshotReferenceConstants)
 from ._featuremodels import map_keyvalue_to_featureflag
 from ._json import parse_json_with_comments
 from ._models import (convert_configurationsetting_to_keyvalue, convert_keyvalue_to_configurationsetting)
@@ -648,6 +649,85 @@ def set_keyvault(cmd,
     raise CLIError("Failed to set the keyvault reference '{}' due to a conflicting operation.".format(key))
 
 
+def set_snapshot_reference(cmd,
+                           key,
+                           snapshot_name,
+                           name=None,
+                           label=None,
+                           tags=None,
+                           yes=False,
+                           connection_string=None,
+                           auth_mode="key",
+                           endpoint=None):
+    azconfig_client = get_appconfig_data_client(cmd, name, connection_string, auth_mode, endpoint)
+
+    snapshot_ref_value = json.dumps({SnapshotReferenceConstants.SNAPSHOT_NAME_KEY: snapshot_name}, ensure_ascii=False)
+    retry_times = 3
+    retry_interval = 1
+
+    label = label if label and label != SearchFilterOptions.EMPTY_LABEL else None
+
+    # generate correlation request id for operations in the same activity
+    correlation_request_id = str(uuid.uuid4())
+
+    for i in range(0, retry_times):
+        retrieved_kv = None
+        set_kv = None
+        new_kv = None
+
+        try:
+            retrieved_kv = azconfig_client.get_configuration_setting(key=key, label=label, headers={HttpHeaders.CORRELATION_REQUEST_ID: correlation_request_id})
+        except ResourceNotFoundError:
+            logger.debug("Key '%s' with label '%s' not found. A new snapshot reference will be created.", key, label)
+        except HttpResponseError as exception:
+            raise CLIErrors.AzureResponseError("Failed to retrieve key-values from config store. " + str(exception))
+
+        if retrieved_kv is None:
+            set_kv = ConfigurationSetting(key=key,
+                                          label=label,
+                                          value=snapshot_ref_value,
+                                          content_type=SnapshotReferenceConstants.SNAPSHOT_REFERENCE_CONTENT_TYPE,
+                                          tags=tags)
+        else:
+            set_kv = ConfigurationSetting(key=key,
+                                          label=label,
+                                          value=snapshot_ref_value,
+                                          content_type=SnapshotReferenceConstants.SNAPSHOT_REFERENCE_CONTENT_TYPE,
+                                          tags=retrieved_kv.tags if tags is None else tags,
+                                          read_only=retrieved_kv.read_only,
+                                          etag=retrieved_kv.etag)
+
+        verification_kv = {
+            "key": set_kv.key,
+            "label": set_kv.label,
+            "content_type": set_kv.content_type,
+            "value": set_kv.value,
+            "tags": set_kv.tags
+        }
+        entry = json.dumps(verification_kv, indent=2, sort_keys=True, ensure_ascii=False)
+        confirmation_message = "Are you sure you want to set the snapshot reference: \n" + entry + "\n"
+        user_confirmation(confirmation_message, yes)
+
+        try:
+            if set_kv.etag is None:
+                new_kv = azconfig_client.add_configuration_setting(set_kv, headers={HttpHeaders.CORRELATION_REQUEST_ID: correlation_request_id})
+            else:
+                new_kv = azconfig_client.set_configuration_setting(set_kv, match_condition=MatchConditions.IfNotModified, headers={HttpHeaders.CORRELATION_REQUEST_ID: correlation_request_id})
+            return convert_configurationsetting_to_keyvalue(new_kv)
+
+        except ResourceReadOnlyError:
+            raise CLIError("Failed to update read only snapshot reference. Unlock the key-value before updating it.")
+        except HttpResponseError as exception:
+            if exception.status_code == StatusCodes.PRECONDITION_FAILED:
+                logger.debug('Retrying setting %s times with exception: concurrent setting operations', i + 1)
+                time.sleep(retry_interval)
+            else:
+                raise CLIErrors.AzureResponseError("Failed to set the snapshot reference due to an exception: " + str(exception))
+        except Exception as exception:
+            raise CLIError("Failed to set the snapshot reference due to an exception: " + str(exception))
+    raise CLIError("Failed to set the snapshot reference '{}' due to a conflicting operation.".format(key))
+
+
 def delete_key(cmd,
                key,
                name=None,
@@ -821,6 +901,7 @@ def list_key(cmd,
              top=None,
              all_=False,
              resolve_keyvault=False,
+             resolve_snapshot_references=False,
              auth_mode="key",
              endpoint=None):
     if fields and resolve_keyvault:
@@ -841,9 +922,46 @@ def list_key(cmd,
                                             top=top,
                                             all_=all_,
                                             cli_ctx=cmd.cli_ctx if resolve_keyvault else None)
+
+    if resolve_snapshot_references:
+        keyvalues = __resolve_snapshot_references(azconfig_client, keyvalues)
+
     return keyvalues
 
 
+def __resolve_snapshot_references(azconfig_client, keyvalues):
+    """Return key-values in the referenced snapshot. The result may contain duplicate keys,
+    which the caller is responsible for handling.
+    """
+    resolved_keyvalues = []
+    for keyvalue in keyvalues:
+        content_type = getattr(keyvalue, 'content_type', None)
+        if not (content_type and SnapshotReferenceConstants.SNAPSHOT_REFERENCE_CONTENT_TYPE in content_type):
+            resolved_keyvalues.append(keyvalue)
+            continue
+
+        reference_value = getattr(keyvalue, 'value', None) or '{}'
+        try:
+            snapshot_name = json.loads(reference_value).get(SnapshotReferenceConstants.SNAPSHOT_NAME_KEY)
+        except (json.JSONDecodeError, TypeError):
+            logger.warning("Skipping snapshot reference with key '%s': invalid value format.", getattr(keyvalue, 'key', None))
+            continue
+
+        if not snapshot_name:
+            logger.warning("Skipping snapshot reference with key '%s': missing snapshot name.", getattr(keyvalue, 'key', None))
+            continue
+
+        try:
+            snapshot_keyvalues = __read_kv_from_config_store(azconfig_client, snapshot=snapshot_name)
+        except Exception as ex:  # pylint: disable=broad-except
+            logger.warning("Skipping snapshot reference '%s': %s", snapshot_name, str(ex))
+            continue
+
+        resolved_keyvalues.extend(snapshot_keyvalues)
+
+    return resolved_keyvalues
+
+
 def restore_key(cmd,
                 datetime,
                 key=None,

src/azure-cli/azure/cli/command_modules/appconfig/linter_exclusions.yml

@@ -15,5 +15,10 @@ appconfig update:
       rule_exclusions:
         - option_length_too_long
     azure_front_door_profile:
+      rule_exclusions:
+        - option_length_too_long
+appconfig kv list:
+  parameters:
+    resolve_snapshot_references:
       rule_exclusions:
         - option_length_too_long