From 0d1c2ca431587eaf666c67fa13b5c4dd43cebff0 Mon Sep 17 00:00:00 2001
From: Martin Albert <martin.albert.187@gmail.com>
Date: Wed, 24 Apr 2024 07:23:32 +0200
Subject: [PATCH 1/2] refactor(traits): parse different structpb values instead
 of general-purpose map

---
 pkg/types/resource/resource.go   | 18 ++++++++++++++++++
 pkg/types/resource/role_trait.go | 12 +++++++++---
 2 files changed, 27 insertions(+), 3 deletions(-)

diff --git a/pkg/types/resource/resource.go b/pkg/types/resource/resource.go
index 585d9bea..144f7090 100644
--- a/pkg/types/resource/resource.go
+++ b/pkg/types/resource/resource.go
@@ -8,6 +8,7 @@ import (
 	v2 "github.com/conductorone/baton-sdk/pb/c1/connector/v2"
 	"github.com/conductorone/baton-sdk/pkg/annotations"
 	"google.golang.org/protobuf/proto"
+	"google.golang.org/protobuf/types/known/structpb"
 )
 
 type ResourceOption func(*v2.Resource) error
@@ -337,3 +338,20 @@ func NewSecretResource(
 
 	return ret, nil
 }
+
+func ToProtoValue(val interface{}) (*structpb.Value, error) {
+	switch v := val.(type) {
+	case string, bool, float64:
+		return structpb.NewValue(val)
+	case int:
+		return structpb.NewNumberValue(float64(v)), nil
+	case *structpb.Value:
+		return v, nil
+	case *structpb.Struct:
+		return structpb.NewStructValue(v), nil
+	case *structpb.ListValue:
+		return structpb.NewListValue(v), nil
+	default:
+		return nil, fmt.Errorf("type %T is not supported", val)
+	}
+}
diff --git a/pkg/types/resource/role_trait.go b/pkg/types/resource/role_trait.go
index 44fe8096..575e2170 100644
--- a/pkg/types/resource/role_trait.go
+++ b/pkg/types/resource/role_trait.go
@@ -12,9 +12,15 @@ type RoleTraitOption func(gt *v2.RoleTrait) error
 
 func WithRoleProfile(profile map[string]interface{}) RoleTraitOption {
 	return func(rt *v2.RoleTrait) error {
-		p, err := structpb.NewStruct(profile)
-		if err != nil {
-			return err
+		p := &structpb.Struct{Fields: make(map[string]*structpb.Value, len(profile))}
+
+		for key, val := range profile {
+			pv, err := ToProtoValue(val)
+			if err != nil {
+				return fmt.Errorf("error converting profile data: %w", err)
+			}
+
+			p.Fields[key] = pv
 		}
 
 		rt.Profile = p

From 8e6c1816c5b7c4503f78abfb676ef86d2d72fbe3 Mon Sep 17 00:00:00 2001
From: Martin Albert <martin.albert.187@gmail.com>
Date: Wed, 24 Apr 2024 07:24:05 +0200
Subject: [PATCH 2/2] chore(traits): add helpers for retrieving different
 structpb values

---
 pkg/types/resource/traits.go | 57 ++++++++++++++++++++++++++++++++++++
 1 file changed, 57 insertions(+)

diff --git a/pkg/types/resource/traits.go b/pkg/types/resource/traits.go
index 470c2fc5..d80c8fc6 100644
--- a/pkg/types/resource/traits.go
+++ b/pkg/types/resource/traits.go
@@ -41,3 +41,60 @@ func GetProfileInt64Value(profile *structpb.Struct, k string) (int64, bool) {
 
 	return int64(s.NumberValue), true
 }
+
+// GetProfileBoolValue returns a bool and true if the value is found.
+func GetProfileBoolValue(profile *structpb.Struct, k string) (bool, bool) {
+	if profile == nil {
+		return false, false
+	}
+
+	v, ok := profile.Fields[k]
+	if !ok {
+		return false, false
+	}
+
+	s, ok := v.Kind.(*structpb.Value_BoolValue)
+	if !ok {
+		return false, false
+	}
+
+	return s.BoolValue, true
+}
+
+// GetProfileListValue returns a list of pointers to structpb.Value and true if the value is found.
+func GetProfileListValue(profile *structpb.Struct, k string) ([]*structpb.Value, bool) {
+	if profile == nil {
+		return nil, false
+	}
+
+	v, ok := profile.Fields[k]
+	if !ok {
+		return nil, false
+	}
+
+	s, ok := v.Kind.(*structpb.Value_ListValue)
+	if !ok {
+		return nil, false
+	}
+
+	return s.ListValue.Values, true
+}
+
+// GetProfileStructValue returns a pointer to structpb.Struct and true if the value is found.
+func GetProfileStructValue(profile *structpb.Struct, k string) (*structpb.Struct, bool) {
+	if profile == nil {
+		return nil, false
+	}
+
+	v, ok := profile.Fields[k]
+	if !ok {
+		return nil, false
+	}
+
+	s, ok := v.Kind.(*structpb.Value_StructValue)
+	if !ok {
+		return nil, false
+	}
+
+	return s.StructValue, true
+}