snapshot reference support

Author: linglingye001

azureappconfiguration/azureappconfiguration.go

@@ -385,6 +385,7 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 	var useAIConfiguration, useAIChatCompletionConfiguration bool
 	kvSettings := make(map[string]any, len(settingsResponse.settings))
 	keyVaultRefs := make(map[string]string)
+	snapshotRefs := make(map[string]string)
 	for trimmedKey, setting := range rawSettings {
 		if setting.ContentType == nil || setting.Value == nil {
 			kvSettings[trimmedKey] = setting.Value
@@ -396,6 +397,9 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 			continue // ignore feature flag while getting key value settings
 		case secretReferenceContentType:
 			keyVaultRefs[trimmedKey] = *setting.Value
+		case snapshotReferenceContentType:
+			snapshotRefs[trimmedKey] = *setting.Value
+			azappcfg.tracingOptions.UseSnapshotReference = true
 		default:
 			if isJsonContentType(setting.ContentType) {
 				var v any
@@ -424,6 +428,10 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 	azappcfg.tracingOptions.UseAIConfiguration = useAIConfiguration
 	azappcfg.tracingOptions.UseAIChatCompletionConfiguration = useAIChatCompletionConfiguration
 
+	if err := azappcfg.loadSettingsFromSnapshotRefs(ctx, settingsClient, snapshotRefs, kvSettings, keyVaultRefs); err != nil {
+		return err
+	}
+
 	secrets, err := azappcfg.loadKeyVaultSecrets(ctx, keyVaultRefs)
 	if err != nil {
 		return fmt.Errorf("failed to load Key Vault secrets: %w", err)
@@ -437,6 +445,58 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 	return nil
 }
 
+func (azappcfg *AzureAppConfiguration) loadSettingsFromSnapshotRefs(ctx context.Context, settingsClient settingsClient, snapshotRefs map[string]string, kvSettings map[string]any, keyVaultRefs map[string]string) error {
+	for key, snapshotRef := range snapshotRefs {
+		// Parse the snapshot reference
+		snapshotName, err := parseSnapshotReference(snapshotRef)
+		if err != nil {
+			return fmt.Errorf("invalid format for Snapshot reference setting %s: %w", key, err)
+		}
+
+		if client, ok := settingsClient.(*selectorSettingsClient); ok {
+			// Load the snapshot settings
+			settingsFromSnapshot, err := loadSnapshotSettings(ctx, client.client, snapshotName)
+			if err != nil {
+				return fmt.Errorf("failed to load snapshot settings: key=%s, error=%s", key, err.Error())
+			}
+
+			for _, setting := range settingsFromSnapshot {
+				if setting.ContentType == nil || setting.Value == nil || setting.Key == nil {
+					continue
+				}
+
+				if *setting.ContentType == featureFlagContentType {
+					continue
+				}
+
+				if *setting.ContentType == secretReferenceContentType {
+					keyVaultRefs[*setting.Key] = *setting.Value
+					continue
+				}
+
+				// Handle JSON content types (similar to regular key-value loading)
+				if isJsonContentType(setting.ContentType) {
+					var v any
+					if err := json.Unmarshal([]byte(*setting.Value), &v); err != nil {
+						// If the value is not valid JSON, try to remove comments and parse again
+						if err := json.Unmarshal(jsonc.StripComments([]byte(*setting.Value)), &v); err != nil {
+							// If still invalid, log the error and treat it as a plain string
+							log.Printf("Failed to unmarshal JSON value from snapshot: key=%s, error=%s", *setting.Key, err.Error())
+							kvSettings[*setting.Key] = setting.Value
+							continue
+						}
+					}
+					kvSettings[*setting.Key] = v
+				} else {
+					kvSettings[*setting.Key] = setting.Value
+				}
+			}
+		}
+	}
+
+	return nil
+}
+
 func (azappcfg *AzureAppConfiguration) loadKeyVaultSecrets(ctx context.Context, keyVaultRefs map[string]string) (map[string]any, error) {
 	secrets := make(map[string]any)
 	if len(keyVaultRefs) == 0 {
@@ -1019,3 +1079,20 @@ func isFailoverable(err error) bool {
 
 	return false
 }
+
+// "{\"snapshot_name\":\"referenced-snapshot\"}"
+func parseSnapshotReference(ref string) (string, error) {
+	var snapshotRef struct {
+		SnapshotName string `json:"snapshot_name"`
+	}
+
+	if err := json.Unmarshal([]byte(ref), &snapshotRef); err != nil {
+		return "", fmt.Errorf("failed to parse snapshot reference: %w", err)
+	}
+
+	if snapshotRef.SnapshotName == "" {
+		return "", fmt.Errorf("snapshot_name is empty in snapshot reference")
+	}
+
+	return snapshotRef.SnapshotName, nil
+}

azureappconfiguration/constants.go

@@ -14,14 +14,15 @@ const (
 
 // General configuration constants
 const (
-	defaultLabel                       = "\x00"
-	wildCard                           = "*"
-	defaultSeparator                   = "."
-	secretReferenceContentType  string = "application/vnd.microsoft.appconfig.keyvaultref+json;charset=utf-8"
-	featureFlagContentType      string = "application/vnd.microsoft.appconfig.ff+json;charset=utf-8"
-	featureFlagKeyPrefix        string = ".appconfig.featureflag/"
-	featureManagementSectionKey string = "feature_management"
-	featureFlagSectionKey       string = "feature_flags"
+	defaultLabel                        = "\x00"
+	wildCard                            = "*"
+	defaultSeparator                    = "."
+	secretReferenceContentType   string = "application/vnd.microsoft.appconfig.keyvaultref+json;charset=utf-8"
+	snapshotReferenceContentType string = "application/json; profile=\"https://azconfig.io/mime-profiles/snapshot-ref\"; charset=utf-8"
+	featureFlagContentType       string = "application/vnd.microsoft.appconfig.ff+json;charset=utf-8"
+	featureFlagKeyPrefix         string = ".appconfig.featureflag/"
+	featureManagementSectionKey  string = "feature_management"
+	featureFlagSectionKey        string = "feature_flags"
 )
 
 // Feature flag constants

azureappconfiguration/internal/tracing/tracing.go

@@ -43,6 +43,7 @@ const (
 	FailoverRequestTag               = "Failover"
 	ReplicaCountKey                  = "ReplicaCount"
 	LoadBalancingEnabledTag          = "LB"
+	SnapshotReferenceTag             = "SnapshotRef"
 
 	// Feature flag usage tracing
 	FMGoVerEnv           = "MS_FEATURE_MANAGEMENT_GO_VERSION"
@@ -74,6 +75,7 @@ type Options struct {
 	KeyVaultRefreshConfigured        bool
 	UseAIConfiguration               bool
 	UseAIChatCompletionConfiguration bool
+	UseSnapshotReference             bool
 	IsFailoverRequest                bool
 	ReplicaCount                     int
 	IsLoadBalancingEnabled           bool
@@ -143,6 +145,10 @@ func CreateCorrelationContextHeader(ctx context.Context, options Options) http.H
 		features = append(features, LoadBalancingEnabledTag)
 	}
 
+	if options.UseSnapshotReference {
+		features = append(features, SnapshotReferenceTag)
+	}
+
 	if len(features) > 0 {
 		featureStr := FeaturesKey + "=" + strings.Join(features, DelimiterPlus)
 		output = append(output, featureStr)

azureappconfiguration/options.go

@@ -91,7 +91,7 @@ type Selector struct {
 // comparableKey returns a comparable representation of the Selector that can be used as a map key.
 // This method creates a deterministic string representation by sorting the TagFilter slice.
 func (s Selector) comparableKey() comparableSelector {
-	cs :=  comparableSelector{
+	cs := comparableSelector{
 		KeyFilter:    s.KeyFilter,
 		LabelFilter:  s.LabelFilter,
 		SnapshotName: s.SnapshotName,

azureappconfiguration/settings_client.go

@@ -86,24 +86,11 @@ func (s *selectorSettingsClient) getSettings(ctx context.Context) (*settingsResp
 
 			pageETags[filter.comparableKey()] = eTags
 		} else {
-			snapshot, err := s.client.GetSnapshot(ctx, filter.SnapshotName, nil)
+			snapshotSettings, err := loadSnapshotSettings(ctx, s.client, filter.SnapshotName)
 			if err != nil {
 				return nil, err
 			}
-
-			if snapshot.CompositionType == nil || *snapshot.CompositionType != azappconfig.CompositionTypeKey {
-				return nil, fmt.Errorf("composition type for the selected snapshot '%s' must be 'key'", filter.SnapshotName)
-			}
-
-			pager := s.client.NewListSettingsForSnapshotPager(filter.SnapshotName, nil)
-			for pager.More() {
-				page, err := pager.NextPage(ctx)
-				if err != nil {
-					return nil, err
-				} else if page.Settings != nil {
-					settings = append(settings, page.Settings...)
-				}
-			}
+			settings = append(settings, snapshotSettings...)
 		}
 	}
 
@@ -211,3 +198,27 @@ func (c *pageETagsClient) checkIfETagChanged(ctx context.Context) (bool, error)
 
 	return false, nil
 }
+
+func loadSnapshotSettings(ctx context.Context, client *azappconfig.Client, snapshotName string) ([]azappconfig.Setting, error) {
+	settings := make([]azappconfig.Setting, 0)
+	snapshot, err := client.GetSnapshot(ctx, snapshotName, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	if snapshot.CompositionType == nil || *snapshot.CompositionType != azappconfig.CompositionTypeKey {
+		return nil, fmt.Errorf("composition type for the selected snapshot '%s' must be 'key'", snapshotName)
+	}
+
+	pager := client.NewListSettingsForSnapshotPager(snapshotName, nil)
+	for pager.More() {
+		page, err := pager.NextPage(ctx)
+		if err != nil {
+			return nil, err
+		} else if page.Settings != nil {
+			settings = append(settings, page.Settings...)
+		}
+	}
+
+	return settings, nil
+}