Repository: kythe/kythe
Pull request #4805 (2021-01-27 20:13:50)

fix(tooling): use protojson for Entry stream encoding


Patch:
diff --git a/kythe/go/storage/stream/BUILD b/kythe/go/storage/stream/BUILD
index abd10cd1b..efc24f87d 100644
--- a/kythe/go/storage/stream/BUILD
+++ b/kythe/go/storage/stream/BUILD
@@ -10,8 +10,8 @@ go_library(
         "//kythe/go/util/schema/facts",
         "//kythe/proto:common_go_proto",
         "//kythe/proto:storage_go_proto",
-        "@com_github_golang_protobuf//jsonpb:go_default_library_gen",
-        "@com_github_golang_protobuf//proto:go_default_library",
+        "@org_golang_google_protobuf//encoding/protojson:go_default_library",
+        "@org_golang_google_protobuf//proto:go_default_library",
     ],
 )
 
@@ -21,5 +21,8 @@ go_test(
     srcs = ["stream_test.go"],
     library = "stream",
     visibility = ["//visibility:private"],
-    deps = ["//kythe/go/test/testutil"],
+    deps = [
+        "//kythe/go/test/testutil",
+        "//kythe/go/util/compare",
+    ],
 )
diff --git a/kythe/go/storage/stream/stream.go b/kythe/go/storage/stream/stream.go
index 2997d7322..3770f0894 100644
--- a/kythe/go/storage/stream/stream.go
+++ b/kythe/go/storage/stream/stream.go
@@ -18,7 +18,6 @@
 package stream // import "kythe.io/kythe/go/storage/stream"
 
 import (
-	"bytes"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -27,8 +26,8 @@ import (
 	"kythe.io/kythe/go/platform/delimited"
 	"kythe.io/kythe/go/util/schema/facts"
 
-	"github.com/golang/protobuf/jsonpb"
-	"github.com/golang/protobuf/proto"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	cpb "kythe.io/kythe/proto/common_go_proto"
 	spb "kythe.io/kythe/proto/storage_go_proto"
@@ -127,11 +126,15 @@ func NewJSONReader(r io.Reader) EntryReader {
 	return func(f func(*spb.Entry) error) error {
 		de := json.NewDecoder(r)
 		for {
-			var entry spb.Entry
-			if err := de.Decode(&entry); err == io.EOF {
+			var raw json.RawMessage
+			if err := de.Decode(&raw); err == io.EOF {
 				return nil
 			} else if err != nil {
-				return fmt.Errorf("error decoding JSON Entry: %v", err)
+				return fmt.Errorf("error decoding JSON Entry: %w", err)
+			}
+			var entry spb.Entry
+			if err := protojson.Unmarshal(raw, &entry); err != nil {
+				return fmt.Errorf("error decoding JSON Entry: %w", err)
 			}
 			if err := f(&entry); err != nil {
 				return err
@@ -140,7 +143,7 @@ func NewJSONReader(r io.Reader) EntryReader {
 	}
 }
 
-var marshaler = &jsonpb.Marshaler{OrigName: true}
+var marshaler = protojson.MarshalOptions{UseProtoNames: true}
 
 // richJSONEntry delays the unmarshaling of the fact_value field
 type richJSONEntry struct {
@@ -155,11 +158,11 @@ func StructuredFactValueJSON(e *spb.Entry) (json.RawMessage, error) {
 		if err := proto.Unmarshal(e.FactValue, &ms); err != nil {
 			return nil, err
 		}
-		msg := new(bytes.Buffer)
-		if err := marshaler.Marshal(msg, &ms); err != nil {
+		rec, err := marshaler.Marshal(&ms)
+		if err != nil {
 			return nil, err
 		}
-		return msg.Bytes(), nil
+		return rec, nil
 	}
 	return json.Marshal(e.FactValue)
 }
@@ -176,9 +179,8 @@ func (r *StructuredEntry) UnmarshalJSON(data []byte) error {
 		return err
 	}
 	if jsonEntry.FactName == facts.Code {
-		factReader := bytes.NewReader(jsonEntry.FactValue)
 		var ms cpb.MarkedSource
-		if err := jsonpb.Unmarshal(factReader, &ms); err != nil {
+		if err := protojson.Unmarshal(jsonEntry.FactValue, &ms); err != nil {
 			return err
 		}
 		pb, err := proto.Marshal(&ms)


Test patch:
diff --git a/kythe/go/storage/stream/stream_test.go b/kythe/go/storage/stream/stream_test.go
index e33af221b..84677b312 100644
--- a/kythe/go/storage/stream/stream_test.go
+++ b/kythe/go/storage/stream/stream_test.go
@@ -24,9 +24,10 @@ import (
 
 	"kythe.io/kythe/go/platform/delimited"
 	"kythe.io/kythe/go/test/testutil"
+	"kythe.io/kythe/go/util/compare"
 
-	"github.com/golang/protobuf/jsonpb"
-	"github.com/golang/protobuf/proto"
+	"google.golang.org/protobuf/encoding/protojson"
+	"google.golang.org/protobuf/proto"
 
 	cpb "kythe.io/kythe/proto/common_go_proto"
 	spb "kythe.io/kythe/proto/storage_go_proto"
@@ -76,12 +77,12 @@ func TestStructuredEntry(t *testing.T) {
 	if err != nil {
 		t.Fatalf("Error marshaling MarkedSource: %v", err)
 	}
-	entry := &StructuredEntry{
+	entry := &spb.Entry{
 		Source:    &spb.VName{Signature: "sig"},
 		FactName:  "/kythe/code",
 		FactValue: pbms,
 	}
-	entryJSON, err := json.Marshal(entry)
+	entryJSON, err := json.Marshal((*StructuredEntry)(entry))
 	if err != nil {
 		t.Fatalf("Error marshaling entry: %v", err)
 	}
@@ -92,7 +93,7 @@ func TestStructuredEntry(t *testing.T) {
 	}
 
 	var msOut cpb.MarkedSource
-	if err := jsonpb.UnmarshalString(string(rawOut.FactValue), &msOut); err != nil {
+	if err := protojson.Unmarshal(rawOut.FactValue, &msOut); err != nil {
 		t.Fatalf("Error unmarshaling fact_value: %v", err)
 	}
 
@@ -105,8 +106,8 @@ func TestStructuredEntry(t *testing.T) {
 		t.Fatalf("Error unmarshaling StructuredEntry: %v", err)
 	}
 
-	if !proto.Equal(entry, &entryOut) {
-		t.Errorf("Roundtrip Marshal/Unmarshal failed: \n%v\n%v", entry, &entryOut)
+	if diff := compare.ProtoDiff(entry, (*spb.Entry)(&entryOut)); diff != "" {
+		t.Errorf("Roundtrip Marshal/Unmarshal failed: \n%v\n%v\n%s", entry, &entryOut, diff)
 	}
 }