Migrate protobuf

flytecopilot/cmd/sidecar_test.go

@@ -6,8 +6,8 @@ import (
 	"os"
 	"testing"
 
-	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/flyteorg/flyte/v2/flytecopilot/cmd/containerwatcher"
 	"github.com/flyteorg/flyte/v2/flytestdlib/promutils"

flytecopilot/data/download.go

@@ -15,10 +15,10 @@ import (
 	"time"
 
 	"github.com/ghodss/yaml"
-	"github.com/golang/protobuf/jsonpb" //nolint: staticcheck
-	"github.com/golang/protobuf/proto"  //nolint: staticcheck
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/pkg/errors"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/futures"
 	"github.com/flyteorg/flyte/v2/flytestdlib/logger"
@@ -289,7 +289,7 @@ func (d Downloader) handleError(_ context.Context, b *core.Error, toFilePath str
 
 func (d Downloader) handleGeneric(ctx context.Context, b *structpb.Struct, toFilePath string, writeToFile bool) (interface{}, error) {
 	if writeToFile && b != nil {
-		m := jsonpb.Marshaler{}
+		m := protojson.MarshalOptions{}
 		writer, err := os.Create(toFilePath)
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to open file at path %s", toFilePath)
@@ -300,7 +300,12 @@ func (d Downloader) handleGeneric(ctx context.Context, b *structpb.Struct, toFil
 				logger.Errorf(ctx, "failed to close File write stream. Error: %s", err)
 			}
 		}()
-		return b, m.Marshal(writer, b)
+		raw, err := m.Marshal(b)
+		if err != nil {
+			return nil, err
+		}
+		_, err = writer.Write(raw)
+		return b, err
 	}
 	return b, nil
 }
@@ -310,7 +315,6 @@ func (d Downloader) handlePrimitive(primitive *core.Primitive, toFilePath string
 
 	var toByteArray func() ([]byte, error)
 	var v interface{}
-	var err error
 
 	switch primitive.GetValue().(type) {
 	case *core.Primitive_StringValue:
@@ -334,18 +338,18 @@ func (d Downloader) handlePrimitive(primitive *core.Primitive, toFilePath string
 			return []byte(strconv.FormatFloat(primitive.GetFloatValue(), 'f', -1, 64)), nil
 		}
 	case *core.Primitive_Datetime:
-		v = primitive.GetDatetime().AsTime()
-		if err != nil {
+		if err := primitive.GetDatetime().CheckValid(); err != nil {
 			return nil, err
 		}
+		v = primitive.GetDatetime().AsTime()
 		toByteArray = func() ([]byte, error) {
 			return []byte(primitive.GetDatetime().AsTime().Format(time.RFC3339Nano)), nil
 		}
 	case *core.Primitive_Duration:
-		v = primitive.GetDuration().AsDuration()
-		if err != nil {
+		if err := primitive.GetDuration().CheckValid(); err != nil {
 			return nil, err
 		}
+		v = primitive.GetDuration().AsDuration()
 		toByteArray = func() ([]byte, error) {
 			return []byte(primitive.GetDuration().AsDuration().String()), nil
 		}
@@ -533,6 +537,10 @@ func (d Downloader) DownloadInputs(ctx context.Context, inputRef storage.DataRef
 		logger.Errorf(ctx, "Failed to download inputs from [%s], err [%s]", inputRef, err)
 		return errors.Wrapf(err, "failed to download input metadata message from remote store")
 	}
+	if len(inputs.GetLiterals()) == 0 {
+		return nil
+	}
+
 	varMap, lMap, err := d.RecursiveDownload(ctx, inputs, outputDir, true)
 	if err != nil {
 		return errors.Wrapf(err, "failed to download input variable from remote store")

flytecopilot/data/download_test.go

@@ -265,7 +265,7 @@ func TestRecursiveDownload(t *testing.T) {
 		}
 
 		// Mock reading the offloaded metadata
-		err = s.WriteProtobuf(context.Background(), storage.DataReference("s3://container/offloaded"), storage.Options{}, &core.Literal{
+		err = s.WriteProtobuf(context.Background(), "s3://container/offloaded", storage.Options{}, &core.Literal{
 			Value: &core.Literal_Map{
 				Map: &core.LiteralMap{
 					Literals: map[string]*core.Literal{

flytecopilot/data/upload.go

@@ -9,8 +9,8 @@ import (
 	"path/filepath"
 	"reflect"
 
-	"github.com/golang/protobuf/proto" //nolint: staticcheck
 	"github.com/pkg/errors"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/flyteorg/flyte/v2/flyteidl2/clients/go/coreutils"
 	"github.com/flyteorg/flyte/v2/flytestdlib/futures"

flyteidl2/clients/go/coreutils/extract_literal_test.go

@@ -8,8 +8,8 @@ import (
 	"testing"
 	"time"
 
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/stretchr/testify/assert"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
 )

flyteidl2/clients/go/coreutils/literals.go

@@ -11,11 +11,11 @@ import (
 	"strings"
 	"time"
 
-	"github.com/golang/protobuf/jsonpb" //nolint: staticcheck
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/pkg/errors"
 	"github.com/shamaton/msgpack/v2"
+	"google.golang.org/protobuf/encoding/protojson"
 	"google.golang.org/protobuf/types/known/durationpb"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/storage"
@@ -378,8 +378,8 @@ func MakeLiteralForSimpleType(t core.SimpleType, s string) (*core.Literal, error
 	switch t {
 	case core.SimpleType_STRUCT:
 		st := &structpb.Struct{}
-		unmarshaler := jsonpb.Unmarshaler{AllowUnknownFields: true}
-		err := unmarshaler.Unmarshal(strings.NewReader(s), st)
+		unmarshaler := protojson.UnmarshalOptions{DiscardUnknown: true}
+		err := unmarshaler.Unmarshal([]byte(s), st)
 		if err != nil {
 			return nil, errors.Wrapf(err, "failed to load generic type as json.")
 		}

flyteidl2/clients/go/coreutils/literals_test.go

@@ -13,11 +13,11 @@ import (
 
 	"github.com/flyteorg/flyte/v2/flytestdlib/storage"
 	"github.com/go-test/deep"
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/pkg/errors"
 	"github.com/shamaton/msgpack/v2"
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/protobuf/types/known/durationpb"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"

flyteplugins/go/tasks/pluginmachinery/core/phase.go

@@ -4,7 +4,7 @@ import (
 	"fmt"
 	"time"
 
-	structpb "github.com/golang/protobuf/ptypes/struct"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
 )

flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot.go

@@ -7,8 +7,8 @@ import (
 	"strconv"
 	"time"
 
-	"github.com/golang/protobuf/proto" //nolint: staticcheck
 	"github.com/pkg/errors"
+	"google.golang.org/protobuf/proto"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 

flyteplugins/go/tasks/pluginmachinery/flytek8s/copilot_test.go

@@ -7,8 +7,8 @@ import (
 	"testing"
 	"time"
 
-	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/protobuf/proto"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"

flyteplugins/go/tasks/pluginmachinery/flytek8s/k8s_resource_adds_test.go

@@ -6,7 +6,6 @@ import (
 	"reflect"
 	"testing"
 
-	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
 	v12 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"
@@ -64,7 +63,7 @@ func TestGetExecutionEnvVars(t *testing.T) {
 		envVars := GetExecutionEnvVars(mock, tt.consoleURL)
 		assert.Len(t, envVars, tt.expectedEnvVars)
 		if tt.expectedEnvVar != nil {
-			assert.True(t, proto.Equal(&envVars[5], tt.expectedEnvVar))
+			assert.Equal(t, tt.expectedEnvVar, &envVars[5])
 		}
 	}
 }

flyteplugins/go/tasks/pluginmachinery/flytek8s/pod_helper.go

@@ -9,8 +9,8 @@ import (
 	"strings"
 	"time"
 
-	"github.com/golang/protobuf/proto" //nolint: staticcheck
 	"github.com/imdario/mergo"
+	"google.golang.org/protobuf/proto"
 	"google.golang.org/protobuf/types/known/timestamppb"
 	v1 "k8s.io/api/core/v1"
 	"k8s.io/apimachinery/pkg/api/resource"

flyteplugins/go/tasks/pluginmachinery/ioutils/remote_file_output_reader_test.go

@@ -10,7 +10,7 @@ import (
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
-	"google.golang.org/protobuf/runtime/protoiface"
+	"google.golang.org/protobuf/proto"
 )
 
 type MemoryMetadata struct {
@@ -54,7 +54,7 @@ func TestReadOrigin(t *testing.T) {
 			},
 		}
 		store := &storageMocks.ComposedProtobufStore{}
-		store.EXPECT().ReadProtobuf(mock.Anything, mock.Anything, mock.Anything).Run(func(ctx context.Context, ref storage.DataReference, msg protoiface.MessageV1) {
+		store.EXPECT().ReadProtobuf(mock.Anything, mock.Anything, mock.Anything).Run(func(ctx context.Context, ref storage.DataReference, msg proto.Message) {
 			assert.NotNil(t, msg)
 			casted := msg.(*core.ErrorDocument)
 			casted.Error = errorDoc.Error
@@ -89,7 +89,7 @@ func TestReadOrigin(t *testing.T) {
 			},
 		}
 		store := &storageMocks.ComposedProtobufStore{}
-		store.EXPECT().ReadProtobuf(mock.Anything, mock.Anything, mock.Anything).Run(func(ctx context.Context, ref storage.DataReference, msg protoiface.MessageV1) {
+		store.EXPECT().ReadProtobuf(mock.Anything, mock.Anything, mock.Anything).Run(func(ctx context.Context, ref storage.DataReference, msg proto.Message) {
 			assert.NotNil(t, msg)
 			casted := msg.(*core.ErrorDocument)
 			casted.Error = errorDoc.Error

flyteplugins/go/tasks/pluginmachinery/ioutils/task_reader_test.go

@@ -5,8 +5,8 @@ import (
 	"fmt"
 	"testing"
 
-	"github.com/golang/protobuf/proto"
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/protobuf/proto"
 
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/core/mocks"
 	"github.com/flyteorg/flyte/v2/flytestdlib/contextutils"

flyteplugins/go/tasks/pluginmachinery/secret/secrets_test.go

@@ -12,6 +12,8 @@ import (
 
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/secret/config"
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/secret/mocks"
+	secretUtils "github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/utils/secrets"
+	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
 )
 
 func TestSecretsWebhook_Mutate(t *testing.T) {
@@ -23,12 +25,16 @@ func TestSecretsWebhook_Mutate(t *testing.T) {
 	})
 
 	namespace := "test-namespace"
+	secretAnnotations, err := secretUtils.MarshalSecretsToMapStrings([]*core.Secret{
+		{
+			Key: "my_key",
+		},
+	})
+	assert.NoError(t, err)
 	podWithAnnotations := &corev1.Pod{
 		ObjectMeta: v1.ObjectMeta{
-			Namespace: namespace,
-			Annotations: map[string]string{
-				"flyte.secrets/s0": "nnsxsorcnv4v623fperca",
-			},
+			Namespace:   namespace,
+			Annotations: secretAnnotations,
 		},
 	}
 

flyteplugins/go/tasks/pluginmachinery/utils/marshal_utils.go

@@ -3,16 +3,15 @@ package utils
 import (
 	"encoding/json"
 	"fmt"
-	"strings"
 
-	"github.com/golang/protobuf/jsonpb" //nolint: staticcheck
-	"github.com/golang/protobuf/proto"  //nolint: staticcheck
-	structpb "github.com/golang/protobuf/ptypes/struct"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 )
 
-var jsonPbMarshaler = jsonpb.Marshaler{}
-var jsonPbUnmarshaler = &jsonpb.Unmarshaler{
-	AllowUnknownFields: true,
+var jsonPbMarshaler = protojson.MarshalOptions{}
+var jsonPbUnmarshaler = protojson.UnmarshalOptions{
+	DiscardUnknown: true,
 }
 
 // Deprecated: Use flytestdlib/utils.UnmarshalStructToPb instead.
@@ -21,12 +20,12 @@ func UnmarshalStruct(structObj *structpb.Struct, msg proto.Message) error {
 		return fmt.Errorf("nil Struct Object passed")
 	}
 
-	jsonObj, err := jsonPbMarshaler.MarshalToString(structObj)
+	jsonObj, err := jsonPbMarshaler.Marshal(structObj)
 	if err != nil {
 		return err
 	}
 
-	if err = jsonPbUnmarshaler.Unmarshal(strings.NewReader(jsonObj), msg); err != nil {
+	if err = jsonPbUnmarshaler.Unmarshal(jsonObj, msg); err != nil {
 		return err
 	}
 
@@ -39,12 +38,12 @@ func MarshalStruct(in proto.Message, out *structpb.Struct) error {
 		return fmt.Errorf("nil Struct Object passed")
 	}
 
-	jsonObj, err := jsonPbMarshaler.MarshalToString(in)
+	jsonObj, err := jsonPbMarshaler.Marshal(in)
 	if err != nil {
 		return err
 	}
 
-	if err = jsonpb.UnmarshalString(jsonObj, out); err != nil {
+	if err = jsonPbUnmarshaler.Unmarshal(jsonObj, out); err != nil {
 		return err
 	}
 
@@ -53,7 +52,8 @@ func MarshalStruct(in proto.Message, out *structpb.Struct) error {
 
 // Deprecated: Use flytestdlib/utils.MarshalToString instead.
 func MarshalToString(msg proto.Message) (string, error) {
-	return jsonPbMarshaler.MarshalToString(msg)
+	b, err := jsonPbMarshaler.Marshal(msg)
+	return string(b), err
 }
 
 // Deprecated: Use flytestdlib/utils.MarshalObjToStruct instead.
@@ -66,7 +66,7 @@ func MarshalObjToStruct(input interface{}) (*structpb.Struct, error) {
 
 	// Turn JSON into a protobuf struct
 	structObj := &structpb.Struct{}
-	if err := jsonpb.UnmarshalString(string(b), structObj); err != nil {
+	if err := jsonPbUnmarshaler.Unmarshal(b, structObj); err != nil {
 		return nil, err
 	}
 	return structObj, nil

flyteplugins/go/tasks/pluginmachinery/utils/marshal_utils_test.go

@@ -5,8 +5,8 @@ import (
 	"testing"
 
 	"github.com/go-test/deep"
-	structpb "github.com/golang/protobuf/ptypes/struct"
 	"github.com/stretchr/testify/assert"
+	structpb "google.golang.org/protobuf/types/known/structpb"
 	v1 "k8s.io/api/core/v1"
 )
 

flyteplugins/go/tasks/pluginmachinery/utils/secrets/marshaler.go

@@ -2,13 +2,12 @@ package secrets
 
 import (
 	"fmt"
+	"strconv"
+	"strings"
 
 	"github.com/flyteorg/flyte/v2/flyteplugins/go/tasks/pluginmachinery/encoding"
 	"github.com/flyteorg/flyte/v2/gen/go/flyteidl2/core"
-	"github.com/golang/protobuf/proto" //nolint: staticcheck
-
-	"strconv"
-	"strings"
+	"google.golang.org/protobuf/encoding/prototext"
 )
 
 const (
@@ -36,7 +35,7 @@ func decodeSecret(encoded string) (string, error) {
 }
 
 func marshalSecret(s *core.Secret) string {
-	return encodeSecret(proto.MarshalTextString(s))
+	return encodeSecret(prototext.MarshalOptions{Multiline: false}.Format(s))
 }
 
 func unmarshalSecret(encoded string) (*core.Secret, error) {
@@ -46,7 +45,7 @@ func unmarshalSecret(encoded string) (*core.Secret, error) {
 	}
 
 	s := &core.Secret{}
-	err = proto.UnmarshalText(decoded, s)
+	err = prototext.Unmarshal([]byte(decoded), s)
 	return s, err
 }