, serialize nested protobuf fields using protojson

commit_hash:a4a31b2722c6221ba23b61b49d55ff4a6b304eab

Author: kabachook

client/resolvers.go

@@ -7,6 +7,8 @@ import (
 	"github.com/apache/arrow-go/v18/arrow"
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/thoas/go-funk"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/runtime/protoimpl"
 	"google.golang.org/protobuf/types/known/timestamppb"
@@ -41,17 +43,57 @@ func ResolveProtoTimestamp(path string) schema.ColumnResolver {
 	}
 }
 
-func ResolveDouble(path string) schema.ColumnResolver {
+// WrapperValue is a constraint for protobuf wrapper types that have a GetValue method
+type WrapperValue[T any] interface {
+	*wrapperspb.DoubleValue | *wrapperspb.FloatValue | *wrapperspb.StringValue |
+		*wrapperspb.Int64Value | *wrapperspb.Int32Value | *wrapperspb.UInt64Value |
+		*wrapperspb.UInt32Value | *wrapperspb.BoolValue | *wrapperspb.BytesValue
+	GetValue() T
+}
+
+// ResolveWrapperValue is a generic resolver for protobuf wrapper types
+func ResolveWrapperValue[T any, W WrapperValue[T]](path string) schema.ColumnResolver {
+	return func(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource, c schema.Column) error {
+		data := funk.Get(resource.Item, path)
+		if data == nil {
+			return nil
+		}
+		wrapper, ok := data.(W)
+		if !ok {
+			var zero W
+			return fmt.Errorf("unexpected type, wanted \"%T\", have \"%T\"", zero, data)
+		}
+		return resource.Set(c.Name, wrapper.GetValue())
+	}
+}
+
+// ResolveProtoMessage creates a resolver that serializes protobuf messages using protojson
+func ResolveProtoMessage(path string) schema.ColumnResolver {
 	return func(ctx context.Context, meta schema.ClientMeta, resource *schema.Resource, c schema.Column) error {
 		data := funk.Get(resource.Item, path)
 		if data == nil {
 			return nil
 		}
-		double, ok := data.(*wrapperspb.DoubleValue)
+
+		// Check if data implements proto.Message
+		msg, ok := data.(proto.Message)
 		if !ok {
-			return fmt.Errorf("unexpected type, wanted \"*wrapperspb.DoubleValue\", have \"%T\"", data)
+			return fmt.Errorf("unexpected type, wanted proto.Message, have %T", data)
 		}
-		return resource.Set(c.Name, double.GetValue())
+
+		// Serialize using protojson with enums as strings
+		marshaler := protojson.MarshalOptions{
+			UseEnumNumbers:    false,
+			EmitDefaultValues: true,
+			UseProtoNames:     true,
+		}
+
+		jsonBytes, err := marshaler.Marshal(msg)
+		if err != nil {
+			return fmt.Errorf("failed to marshal proto message to JSON: %w", err)
+		}
+
+		return resource.Set(c.Name, jsonBytes)
 	}
 }
 

client/transformers.go

@@ -7,6 +7,7 @@ import (
 	"github.com/cloudquery/plugin-sdk/v4/schema"
 	"github.com/cloudquery/plugin-sdk/v4/transformers"
 	cqtypes "github.com/cloudquery/plugin-sdk/v4/types"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/reflect/protoreflect"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	"google.golang.org/protobuf/types/known/wrapperspb"
@@ -17,11 +18,36 @@ func typeTransformer(field reflect.StructField) (arrow.DataType, error) {
 	case *timestamppb.Timestamp,
 		timestamppb.Timestamp:
 		return arrow.FixedWidthTypes.Timestamp_us, nil
+	case *protoreflect.Enum,
+		protoreflect.Enum:
+		return arrow.BinaryTypes.String, nil
 	case *wrapperspb.DoubleValue,
 		wrapperspb.DoubleValue:
 		return arrow.PrimitiveTypes.Float64, nil
-	case protoreflect.Enum:
+	case *wrapperspb.FloatValue,
+		wrapperspb.FloatValue:
+		return arrow.PrimitiveTypes.Float32, nil
+	case *wrapperspb.StringValue,
+		wrapperspb.StringValue:
 		return arrow.BinaryTypes.String, nil
+	case *wrapperspb.Int64Value,
+		wrapperspb.Int64Value:
+		return arrow.PrimitiveTypes.Int64, nil
+	case *wrapperspb.Int32Value,
+		wrapperspb.Int32Value:
+		return arrow.PrimitiveTypes.Int32, nil
+	case *wrapperspb.UInt64Value,
+		wrapperspb.UInt64Value:
+		return arrow.PrimitiveTypes.Uint64, nil
+	case *wrapperspb.UInt32Value,
+		wrapperspb.UInt32Value:
+		return arrow.PrimitiveTypes.Uint32, nil
+	case *wrapperspb.BoolValue,
+		wrapperspb.BoolValue:
+		return arrow.FixedWidthTypes.Boolean, nil
+	case *wrapperspb.BytesValue,
+		wrapperspb.BytesValue:
+		return arrow.BinaryTypes.Binary, nil
 	case nil:
 		return cqtypes.NewJSONType(), nil
 	default:
@@ -34,11 +60,39 @@ func resolverTransformer(field reflect.StructField, path string) schema.ColumnRe
 	case *timestamppb.Timestamp,
 		timestamppb.Timestamp:
 		return ResolveProtoTimestamp(path)
+	case *protoreflect.Enum,
+		protoreflect.Enum:
+		return ResolveProtoEnum(path)
 	case *wrapperspb.DoubleValue,
 		wrapperspb.DoubleValue:
-		return ResolveDouble(path)
-	case protoreflect.Enum:
-		return ResolveProtoEnum(path)
+		return ResolveWrapperValue[float64, *wrapperspb.DoubleValue](path)
+	case *wrapperspb.FloatValue,
+		wrapperspb.FloatValue:
+		return ResolveWrapperValue[float32, *wrapperspb.FloatValue](path)
+	case *wrapperspb.StringValue,
+		wrapperspb.StringValue:
+		return ResolveWrapperValue[string, *wrapperspb.StringValue](path)
+	case *wrapperspb.Int64Value,
+		wrapperspb.Int64Value:
+		return ResolveWrapperValue[int64, *wrapperspb.Int64Value](path)
+	case *wrapperspb.Int32Value,
+		wrapperspb.Int32Value:
+		return ResolveWrapperValue[int32, *wrapperspb.Int32Value](path)
+	case *wrapperspb.UInt64Value,
+		wrapperspb.UInt64Value:
+		return ResolveWrapperValue[uint64, *wrapperspb.UInt64Value](path)
+	case *wrapperspb.UInt32Value,
+		wrapperspb.UInt32Value:
+		return ResolveWrapperValue[uint32, *wrapperspb.UInt32Value](path)
+	case *wrapperspb.BoolValue,
+		wrapperspb.BoolValue:
+		return ResolveWrapperValue[bool, *wrapperspb.BoolValue](path)
+	case *wrapperspb.BytesValue,
+		wrapperspb.BytesValue:
+		return ResolveWrapperValue[[]byte, *wrapperspb.BytesValue](path)
+	case *proto.Message,
+		proto.Message:
+		return ResolveProtoMessage(path)
 	default:
 		return nil
 	}