From 5cb2cba9f04977d12ed215e042248dae8f1b73e9 Mon Sep 17 00:00:00 2001
From: Jonathan Amsterdam <jba@google.com>
Date: Sat, 25 May 2024 12:07:05 -0400
Subject: [PATCH] chore(internal/protoveneer): support time

Protos represent time using the timestamppb.Timestamp type.
Support converting that to and from a time.Time.
---
 .../cmd/protoveneer/protoveneer.go            | 21 +++++++++++++++----
 .../protoveneer/testdata/basic/basic.pb.go    |  2 ++
 .../cmd/protoveneer/testdata/basic/golden     |  5 +++++
 internal/protoveneer/support/support.go       |  9 ++++++++
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/internal/protoveneer/cmd/protoveneer/protoveneer.go b/internal/protoveneer/cmd/protoveneer/protoveneer.go
index d34ba8b01a88..93b55ff519a7 100644
--- a/internal/protoveneer/cmd/protoveneer/protoveneer.go
+++ b/internal/protoveneer/cmd/protoveneer/protoveneer.go
@@ -676,11 +676,17 @@ func write(typeInfos []*typeInfo, conf *config, fset *token.FileSet) ([]byte, er
 		return nil, errors.New("missing supportImportPath in config")
 	}
 	prn(`    "%s"`, conf.SupportImportPath)
+	importPaths := map[string]bool{}
 	for _, et := range externalTypes {
-		if et.used && et.importPath != "" {
-			prn(`    "%s"`, et.importPath)
+		if et.used {
+			for _, ip := range et.importPaths {
+				importPaths[ip] = true
+			}
 		}
 	}
+	for ip := range importPaths {
+		prn(`    "%s"`, ip)
+	}
 	pr(")\n\n")
 
 	// Types.
@@ -775,7 +781,7 @@ func (ti *typeInfo) generateFromProto(pr func(string, ...any)) {
 type externalType struct {
 	qualifiedName string
 	replaces      string
-	importPath    string
+	importPaths   []string
 	convertTo     string
 	convertFrom   string
 
@@ -787,7 +793,7 @@ var externalTypes = []*externalType{
 	{
 		qualifiedName: "civil.Date",
 		replaces:      "*date.Date",
-		importPath:    "cloud.google.com/go/civil",
+		importPaths:   []string{"cloud.google.com/go/civil"},
 		convertTo:     "support.CivilDateToProto",
 		convertFrom:   "support.CivilDateFromProto",
 	},
@@ -797,6 +803,13 @@ var externalTypes = []*externalType{
 		convertTo:     "support.MapToStructPB",
 		convertFrom:   "support.MapFromStructPB",
 	},
+	{
+		qualifiedName: "time.Time",
+		replaces:      "*timestamppb.Timestamp",
+		importPaths:   []string{"time", "google.golang.org/protobuf/types/known/timestamppb"},
+		convertTo:     "timestamppb.New",
+		convertFrom:   "support.TimeFromProto",
+	},
 }
 
 var protoTypeToExternalType = map[string]*externalType{}
diff --git a/internal/protoveneer/cmd/protoveneer/testdata/basic/basic.pb.go b/internal/protoveneer/cmd/protoveneer/testdata/basic/basic.pb.go
index f87738971cac..2adb4ada76b7 100755
--- a/internal/protoveneer/cmd/protoveneer/testdata/basic/basic.pb.go
+++ b/internal/protoveneer/cmd/protoveneer/testdata/basic/basic.pb.go
@@ -19,6 +19,7 @@ import (
 	date "google.golang.org/genproto/googleapis/type/date"
 	protoimpl "google.golang.org/protobuf/runtime/protoimpl"
 	"google.golang.org/protobuf/types/known/structpb"
+	timestamppb "google.golang.org/protobuf/types/known/timestamppb"
 )
 
 const (
@@ -120,4 +121,5 @@ type Citation struct {
 	// Output only. Publication date of the attribution.
 	PublicationDate *date.Date `protobuf:"bytes,6,opt,name=publication_date,json=publicationDate,proto3" json:"publication_date,omitempty"`
 	Struct          *structpb.Struct
+	CreateTime      *timestamppb.Timestamp
 }
diff --git a/internal/protoveneer/cmd/protoveneer/testdata/basic/golden b/internal/protoveneer/cmd/protoveneer/testdata/basic/golden
index 3de88757abd0..2a997288d1e8 100644
--- a/internal/protoveneer/cmd/protoveneer/testdata/basic/golden
+++ b/internal/protoveneer/cmd/protoveneer/testdata/basic/golden
@@ -8,6 +8,8 @@ import (
 	"cloud.google.com/go/civil"
 	pb "example.com/basic"
 	"example.com/protoveneer/support"
+	"google.golang.org/protobuf/types/known/timestamppb"
+	"time"
 )
 
 // Blob contains raw media bytes.
@@ -47,6 +49,7 @@ type Citation struct {
 	// Output only. Publication date of the attribution.
 	PublicationDate civil.Date
 	Struct          map[string]any
+	CreateTime      time.Time
 }
 
 func (v *Citation) toProto() *pb.Citation {
@@ -57,6 +60,7 @@ func (v *Citation) toProto() *pb.Citation {
 		Uri:             v.URI,
 		PublicationDate: support.CivilDateToProto(v.PublicationDate),
 		Struct:          support.MapToStructPB(v.Struct),
+		CreateTime:      timestamppb.New(v.CreateTime),
 	}
 }
 
@@ -68,6 +72,7 @@ func (Citation) fromProto(p *pb.Citation) *Citation {
 		URI:             p.Uri,
 		PublicationDate: support.CivilDateFromProto(p.PublicationDate),
 		Struct:          support.MapFromStructPB(p.Struct),
+		CreateTime:      support.TimeFromProto(p.CreateTime),
 	}
 }
 
diff --git a/internal/protoveneer/support/support.go b/internal/protoveneer/support/support.go
index 0463fe4d645a..e9f4b454703a 100644
--- a/internal/protoveneer/support/support.go
+++ b/internal/protoveneer/support/support.go
@@ -22,6 +22,7 @@ import (
 	"cloud.google.com/go/civil"
 	"google.golang.org/genproto/googleapis/type/date"
 	"google.golang.org/protobuf/types/known/structpb"
+	"google.golang.org/protobuf/types/known/timestamppb"
 )
 
 // TransformSlice applies f to each element of from and returns
@@ -110,3 +111,11 @@ func MapFromStructPB(p *structpb.Struct) map[string]any {
 	}
 	return p.AsMap()
 }
+
+// TimeFromProto converts a Timestamp into a time.Time.
+func TimeFromProto(ts *timestamppb.Timestamp) time.Time {
+	if ts == nil {
+		return time.Time{}
+	}
+	return ts.AsTime()
+}