Check for optional/ Ptr types within Unions [go/programgen]

changelog/pending/20221212--programgen-go--Check-for-optional-Ptr-types-within-Union-types.yaml

@@ -0,0 +1,4 @@
+changes:
+- type: fix
+  scope: programgen/go
+  description: Check for optional/ Ptr types within Union types. This fixes a bug in Go programgen where optional outputs are not returned as pointers.

pkg/codegen/go/gen_program_expressions.go

@@ -889,8 +889,18 @@ func (g *generator) argumentTypeName(expr model.Expression, destType model.Type,
 		return g.argumentTypeName(expr, destType.ElementType, isInput)
 	case *model.UnionType:
 		for _, ut := range destType.ElementTypes {
+			isOptional := false
+			// check if the union contains none, which indicates this is an optional value
+			for _, ut := range destType.ElementTypes {
+				if ut.Equals(model.NoneType) {
+					isOptional = true
+				}
+			}
 			switch ut := ut.(type) {
 			case *model.OpaqueType:
+				if isOptional {
+					return g.argumentTypeNamePtr(expr, ut, isInput)
+				}
 				return g.argumentTypeName(expr, ut, isInput)
 			case *model.ConstType:
 				return g.argumentTypeName(expr, ut.Type, isInput)
@@ -907,6 +917,11 @@ func (g *generator) argumentTypeName(expr model.Expression, destType model.Type,
 	return ""
 }
 
+func (g *generator) argumentTypeNamePtr(expr model.Expression, destType model.Type, isInput bool) (result string) {
+	res := g.argumentTypeName(expr, destType, isInput)
+	return "*" + res
+}
+
 func (g *generator) genRelativeTraversal(w io.Writer,
 	traversal hcl.Traversal, parts []model.Traversable, isRootResource bool) {
 
@@ -1012,13 +1027,23 @@ func (g *generator) genApply(w io.Writer, expr *model.FunctionCallExpression) {
 	if retType == "[]string" {
 		typeAssertion = ".(pulumi.StringArrayOutput)"
 	} else {
+		if strings.HasPrefix(retType, "*") {
+			retType = Title(strings.TrimPrefix(retType, "*")) + "Ptr"
+			switch then.Body.(type) {
+			case *model.ScopeTraversalExpression:
+				traversal := then.Body.(*model.ScopeTraversalExpression)
+				traversal.RootName = "&" + traversal.RootName
+				then.Body = traversal
+			}
+		}
 		typeAssertion = fmt.Sprintf(".(%sOutput)", retType)
 		if !strings.Contains(retType, ".") {
 			typeAssertion = fmt.Sprintf(".(pulumi.%sOutput)", Title(retType))
 		}
 	}
 
 	if len(applyArgs) == 1 {
+		then.Signature.ReturnType = model.NewOptionalType(then.Signature.ReturnType)
 		// If we only have a single output, just generate a normal `.Apply`
 		g.Fgenf(w, "%.v.ApplyT(%.v)%s", applyArgs[0], then, typeAssertion)
 	} else {

pkg/codegen/testing/test/testdata/aws-s3-logging-pp/go/aws-s3-logging.go

@@ -21,9 +21,9 @@ func main() {
 		if err != nil {
 			return err
 		}
-		ctx.Export("targetBucket", bucket.Loggings.ApplyT(func(loggings []s3.BucketLogging) (string, error) {
-			return loggings[0].TargetBucket, nil
-		}).(pulumi.StringOutput))
+		ctx.Export("targetBucket", bucket.Loggings.ApplyT(func(loggings []s3.BucketLogging) (*string, error) {
+			return &loggings[0].TargetBucket, nil
+		}).(pulumi.StringPtrOutput))
 		return nil
 	})
 }

pkg/codegen/testing/test/testdata/synthetic-resource-properties-pp/go/synthetic-resource-properties.go

@@ -16,9 +16,9 @@ func main() {
 		ctx.Export("foo", rt.Res1.ApplyT(func(res1 *resourceproperties.Res1) (resourceproperties.Obj2, error) {
 			return res1.Obj1.Res2.Obj2, nil
 		}).(resourceproperties.Obj2Output))
-		ctx.Export("complex", rt.Res1.ApplyT(func(res1 *resourceproperties.Res1) (float64, error) {
-			return res1.Obj1.Res2.Obj2.Answer, nil
-		}).(pulumi.Float64Output))
+		ctx.Export("complex", rt.Res1.ApplyT(func(res1 *resourceproperties.Res1) (*float64, error) {
+			return &res1.Obj1.Res2.Obj2.Answer, nil
+		}).(pulumi.Float64PtrOutput))
 		return nil
 	})
 }