From 12baad140b89b51cc770de05640c219e722cd418 Mon Sep 17 00:00:00 2001
From: Venelin <venelin@pulumi.com>
Date: Tue, 17 Dec 2024 16:52:29 +0200
Subject: [PATCH] Add a TransformPropertyValueLimitDescent property value
 walker

---
 unstable/propertyvalue/propertyvalue.go      | 77 ++++++++++++++++++++
 unstable/propertyvalue/propertyvalue_test.go | 45 ++++++++++++
 2 files changed, 122 insertions(+)

diff --git a/unstable/propertyvalue/propertyvalue.go b/unstable/propertyvalue/propertyvalue.go
index f882bcab4..799916766 100644
--- a/unstable/propertyvalue/propertyvalue.go
+++ b/unstable/propertyvalue/propertyvalue.go
@@ -15,6 +15,8 @@
 package propertyvalue
 
 import (
+	"errors"
+
 	"github.com/pulumi/pulumi/sdk/v3/go/common/resource"
 )
 
@@ -108,6 +110,81 @@ func TransformPropertyValue(
 	return transformer(path, value)
 }
 
+type LimitDescentError struct{}
+
+func (LimitDescentError) Error() string {
+	return "limit descent"
+}
+
+// TransformPropertyValueLimitDescent is a variant of TransformPropertyValue that allows the transformer
+// to return a LimitDescentError to indicate that the recursion should not descend into the value without
+// aborting the whole transformation.
+func TransformPropertyValueLimitDescent(
+	path resource.PropertyPath,
+	transformer func(resource.PropertyPath, resource.PropertyValue) (resource.PropertyValue, error),
+	value resource.PropertyValue,
+) (resource.PropertyValue, error) {
+	value, err := transformer(path, value)
+	if err != nil {
+		if errors.Is(err, LimitDescentError{}) {
+			return value, nil
+		}
+		return resource.NewNullProperty(), err
+	}
+
+	switch {
+	case value.IsArray():
+		// preserve nil arrays
+		if !isNilArray(value) {
+			tvs := []resource.PropertyValue{}
+			for i, v := range value.ArrayValue() {
+				tv, err := TransformPropertyValueLimitDescent(extendPath(path, i), transformer, v)
+				if err != nil {
+					return resource.NewNullProperty(), err
+				}
+				tvs = append(tvs, tv)
+			}
+			value = resource.NewArrayProperty(tvs)
+		}
+	case value.IsObject():
+		// preserve nil objects
+		if !isNilObject(value) {
+			pm := make(resource.PropertyMap)
+			for k, v := range value.ObjectValue() {
+				tv, err := TransformPropertyValueLimitDescent(extendPath(path, string(k)), transformer, v)
+				if err != nil {
+					return resource.NewNullProperty(), err
+				}
+				pm[k] = tv
+			}
+			value = resource.NewObjectProperty(pm)
+		}
+	case value.IsOutput():
+		o := value.OutputValue()
+		te, err := TransformPropertyValueLimitDescent(path, transformer, o.Element)
+		if err != nil {
+			return resource.NewNullProperty(), err
+		}
+		value = resource.NewOutputProperty(resource.Output{
+			Element:      te,
+			Known:        o.Known,
+			Secret:       o.Secret,
+			Dependencies: o.Dependencies,
+		})
+	case value.IsSecret():
+		s := value.SecretValue()
+		te, err := TransformPropertyValueLimitDescent(path, transformer, s.Element)
+		if err != nil {
+			return resource.NewNullProperty(), err
+		}
+		value = resource.NewSecretProperty(&resource.Secret{
+			Element: te,
+		})
+	}
+
+	return value, nil
+}
+
 // Removes any resource.NewSecretProperty wrappers. Removes Secret: true flags from any first-class outputs.
 func RemoveSecrets(pv resource.PropertyValue) resource.PropertyValue {
 	unsecret := func(pv resource.PropertyValue) resource.PropertyValue {
diff --git a/unstable/propertyvalue/propertyvalue_test.go b/unstable/propertyvalue/propertyvalue_test.go
index 7aaca5ebb..eea662f6d 100644
--- a/unstable/propertyvalue/propertyvalue_test.go
+++ b/unstable/propertyvalue/propertyvalue_test.go
@@ -89,3 +89,48 @@ func TestTransformPreservesNilObjects(t *testing.T) {
 	require.True(t, result.IsObject())
 	require.Nil(t, result.ObjectValue())
 }
+
+func TestTransformPropertyValueLimitDescent(t *testing.T) {
+	t.Parallel()
+	t.Run("simple value transformation", func(t *testing.T) {
+		t.Parallel()
+		input := resource.NewStringProperty("hello")
+		transformer := func(_ resource.PropertyPath, v resource.PropertyValue) (resource.PropertyValue, error) {
+			if v.IsString() {
+				return resource.NewStringProperty(v.StringValue() + " world"), nil
+			}
+			return v, nil
+		}
+
+		result, err := TransformPropertyValueLimitDescent(nil, transformer, input)
+		require.NoError(t, err)
+		require.Equal(t, "hello world", result.StringValue())
+	})
+
+	t.Run("limit descent on array", func(t *testing.T) {
+		t.Parallel()
+		input := resource.NewObjectProperty(resource.PropertyMap{
+			"array": resource.NewArrayProperty([]resource.PropertyValue{
+				resource.NewStringProperty("should not transform"),
+			}),
+			"string": resource.NewStringProperty("should"),
+		})
+
+		transformer := func(_ resource.PropertyPath, v resource.PropertyValue) (resource.PropertyValue, error) {
+			if v.IsArray() {
+				return v, LimitDescentError{}
+			}
+			if v.IsString() {
+				return resource.NewStringProperty(v.StringValue() + " transformed"), nil
+			}
+			return v, nil
+		}
+
+		result, err := TransformPropertyValueLimitDescent(nil, transformer, input)
+		require.NoError(t, err)
+		require.True(t, result.IsObject())
+		require.Equal(t, "should transformed", result.ObjectValue()["string"].StringValue())
+		arr := result.ObjectValue()["array"].ArrayValue()
+		require.Equal(t, "should not transform", arr[0].StringValue())
+	})
+}