Snapshot reference support

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,21 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 	azappcfg.tracingOptions.UseAIConfiguration = useAIConfiguration
 	azappcfg.tracingOptions.UseAIChatCompletionConfiguration = useAIChatCompletionConfiguration
 
+	if len(snapshotRefs) > 0 {
+		var loadSnapshot snapshotSettingsLoader
+		if client, ok := settingsClient.(*selectorSettingsClient); ok {
+			loadSnapshot = func(ctx context.Context, snapshotName string) ([]azappconfig.Setting, error) {
+				return loadSnapshotSettings(ctx, client.client, snapshotName)
+			}
+		}
+
+		if loadSnapshot != nil {
+			if err := azappcfg.loadSettingsFromSnapshotRefs(ctx, loadSnapshot, 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 +456,78 @@ func (azappcfg *AzureAppConfiguration) loadKeyValues(ctx context.Context, settin
 	return nil
 }
 
+func (azappcfg *AzureAppConfiguration) loadSettingsFromSnapshotRefs(ctx context.Context, loadSnapshot snapshotSettingsLoader, snapshotRefs map[string]string, kvSettings map[string]any, keyVaultRefs map[string]string) error {
+	var useAIConfiguration, useAIChatCompletionConfiguration bool
+	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)
+		}
+
+		// Load the snapshot settings
+		settingsFromSnapshot, err := loadSnapshot(ctx, snapshotName)
+		if err != nil {
+			return fmt.Errorf("failed to load snapshot settings: key=%s, error=%w", key, err)
+		}
+
+		for _, setting := range settingsFromSnapshot {
+			if setting.Key == nil {
+				continue
+			}
+
+			trimmedKey := azappcfg.trimPrefix(*setting.Key)
+			if len(trimmedKey) == 0 {
+				log.Printf("Key of the setting '%s' is trimmed to the empty string, just ignore it", *setting.Key)
+				continue
+			}
+
+			if setting.ContentType == nil || setting.Value == nil {
+				kvSettings[trimmedKey] = setting.Value
+				continue
+			}
+
+			contentType := strings.TrimSpace(strings.ToLower(*setting.ContentType))
+			if contentType == featureFlagContentType {
+				continue
+			}
+
+			if contentType == secretReferenceContentType {
+				keyVaultRefs[trimmedKey] = *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[trimmedKey] = setting.Value
+						continue
+					}
+				}
+				kvSettings[trimmedKey] = v
+				if isAIConfigurationContentType(setting.ContentType) {
+					useAIConfiguration = true
+				}
+				if isAIChatCompletionContentType(setting.ContentType) {
+					useAIChatCompletionConfiguration = true
+				}
+			} else {
+				kvSettings[trimmedKey] = setting.Value
+			}
+		}
+	}
+
+	azappcfg.tracingOptions.UseAIConfiguration = useAIConfiguration
+	azappcfg.tracingOptions.UseAIChatCompletionConfiguration = useAIChatCompletionConfiguration
+
+	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 +1110,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/internal/tracing/tracing_test.go

@@ -145,12 +145,70 @@ func TestCreateCorrelationContextHeader(t *testing.T) {
 		assert.Contains(t, features, AIChatCompletionConfigurationTag)
 	})
 
+	t.Run("with snapshot reference", func(t *testing.T) {
+		ctx := context.Background()
+		options := Options{
+			UseSnapshotReference: true,
+		}
+
+		header := CreateCorrelationContextHeader(ctx, options)
+
+		corrContext := header.Get(CorrelationContextHeader)
+		assert.Contains(t, corrContext, FeaturesKey+"="+SnapshotReferenceTag)
+	})
+
+	t.Run("with snapshot reference not set", func(t *testing.T) {
+		ctx := context.Background()
+		options := Options{
+			UseSnapshotReference: false,
+		}
+
+		header := CreateCorrelationContextHeader(ctx, options)
+
+		corrContext := header.Get(CorrelationContextHeader)
+		assert.NotContains(t, corrContext, SnapshotReferenceTag)
+	})
+
+	t.Run("with snapshot reference and other features", func(t *testing.T) {
+		ctx := context.Background()
+		options := Options{
+			UseAIConfiguration:   true,
+			UseSnapshotReference: true,
+		}
+
+		header := CreateCorrelationContextHeader(ctx, options)
+
+		corrContext := header.Get(CorrelationContextHeader)
+		assert.Contains(t, corrContext, FeaturesKey+"=")
+
+		// Extract the Features part
+		parts := strings.Split(corrContext, DelimiterComma)
+		var featuresPart string
+		for _, part := range parts {
+			if strings.HasPrefix(part, FeaturesKey+"=") {
+				featuresPart = part
+				break
+			}
+		}
+
+		// Check both tags are in the features part
+		assert.Contains(t, featuresPart, AIConfigurationTag)
+		assert.Contains(t, featuresPart, SnapshotReferenceTag)
+
+		// Check the delimiter is correct
+		features := strings.Split(strings.TrimPrefix(featuresPart, FeaturesKey+"="), DelimiterPlus)
+		assert.Len(t, features, 2)
+		assert.Contains(t, features, AIConfigurationTag)
+		assert.Contains(t, features, SnapshotReferenceTag)
+	})
+
 	t.Run("with all options", func(t *testing.T) {
 		options := Options{
 			Host:                             HostTypeAzureFunction,
 			KeyVaultConfigured:               true,
 			UseAIConfiguration:               true,
 			UseAIChatCompletionConfiguration: true,
+			UseSnapshotReference:             true,
 		}
 
 		header := CreateCorrelationContextHeader(context.Background(), options)
@@ -172,9 +230,10 @@ func TestCreateCorrelationContextHeader(t *testing.T) {
 			}
 		}
 
-		// Check both AI tags are in the features part
+		// Check all feature tags are in the features part
 		assert.Contains(t, featuresPart, AIConfigurationTag)
 		assert.Contains(t, featuresPart, AIChatCompletionConfigurationTag)
+		assert.Contains(t, featuresPart, SnapshotReferenceTag)
 
 		// Verify the header format
 		assert.Equal(t, 4, strings.Count(corrContext, DelimiterComma)+1, "Should have 4 parts")

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

@@ -50,6 +50,9 @@ type eTagsClient interface {
 	checkIfETagChanged(ctx context.Context) (bool, error)
 }
 
+// snapshotSettingsLoader is a function type that loads settings from a snapshot by name.
+type snapshotSettingsLoader func(ctx context.Context, snapshotName string) ([]azappconfig.Setting, error)
+
 type refreshClient struct {
 	loader    settingsClient
 	monitor   eTagsClient
@@ -86,24 +89,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 +201,31 @@ 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 {
+		var respErr *azcore.ResponseError
+		if errors.As(err, &respErr) && respErr.StatusCode == 404 {
+			return settings, nil // treat non-existing snapshot as empty
+		}
+		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
+}