@@ -18,12 +18,15 @@ import (
1818 "fmt"
1919
2020 "github.com/golang/protobuf/protoc-gen-go/descriptor"
21+ "github.com/googleapis/gapic-generator-go/internal/pbinfo"
22+ "google.golang.org/genproto/googleapis/cloud/extendedops"
23+ "google.golang.org/protobuf/proto"
2124)
2225
2326// customOp represents a custom operation type for long running operations.
2427type customOp struct {
25- message * descriptor.DescriptorProto
26- generated bool
28+ message * descriptor.DescriptorProto
29+ handles [] * descriptor. ServiceDescriptorProto
2730}
2831
2932// isCustomOp determines if the given method should return a custom operation wrapper.
@@ -63,39 +66,219 @@ func (g *generator) customOpPointerType() (string, error) {
6366// customOpInit builds a string containing the Go code for initializing the
6467// operation wrapper type with the Go identifier for a variable that is the
6568// proto-defined operation type.
66- func (g * generator ) customOpInit (p string ) string {
69+ func (g * generator ) customOpInit (h , p string ) string {
6770 opName := g .aux .customOp .message .GetName ()
6871
69- s := fmt .Sprintf ("&%s{proto: %s}" , opName , p )
72+ s := fmt .Sprintf ("&%s{&%s{c: c.operationClient, proto: %s}} " , opName , h , p )
7073
7174 return s
7275}
7376
74- // customOperationType generates the custom operation wrapper type using the
75- // generators current printer. This should only be called once per package.
77+ // customOperationType generates the custom operation wrapper type and operation
78+ // service handle implementations using the generators current printer. This
79+ // should only be called once per package.
7680func (g * generator ) customOperationType () error {
7781 op := g .aux .customOp
7882 if op == nil {
7983 return nil
8084 }
8185 opName := op .message .GetName ()
86+ handleInt := lowerFirst (opName + "Handle" )
8287
8388 ptyp , err := g .customOpPointerType ()
8489 if err != nil {
8590 return err
8691 }
92+ _ , opImp , err := g .descInfo .NameSpec (op .message )
93+ if err != nil {
94+ return err
95+ }
96+ g .imports [opImp ] = true
97+
98+ statusField := operationField (op .message , extendedops .OperationResponseMapping_STATUS )
99+ if statusField == nil {
100+ return fmt .Errorf ("operation message %s is missing an annotated status field" , op .message .GetName ())
101+ }
102+
103+ opNameField := operationField (op .message , extendedops .OperationResponseMapping_NAME )
104+ if opNameField == nil {
105+ return fmt .Errorf ("operation message %s is missing an annotated name field" , op .message .GetName ())
106+ }
107+
108+ opNameGetter := fieldGetter (opNameField .GetName ())
87109
88110 p := g .printf
89111
90- p ("// %s represents a long running operation for this API." , opName )
112+ p ("// %s represents a long- running operation for this API." , opName )
91113 p ("type %s struct {" , opName )
92- p (" proto %s" , ptyp )
114+ p (" %s" , handleInt )
93115 p ("}" )
94116 p ("" )
95- p ("// Proto returns the raw type this wraps." )
96- p ("func (o *%s) Proto() %s {" , opName , ptyp )
97- p (" return o.proto" )
117+
118+ // Done
119+ p ("// Done reports whether the long-running operation has completed." )
120+ p ("func (o *%s) Done() bool {" , opName )
121+ p (g .customOpStatusCheck (statusField ))
98122 p ("}" )
123+ p ("" )
124+
125+ // Name
126+ p ("// Name returns the name of the long-running operation." )
127+ p ("// The name is assigned by the server and is unique within the service from which the operation is created." )
128+ p ("func (o *%s) Name() string {" , opName )
129+ p (" return o.Proto()%s" , opNameGetter )
130+ p ("}" )
131+ p ("" )
132+
133+ p ("type %s interface {" , handleInt )
134+ p (" // Poll retrieves the operation." )
135+ p (" Poll(ctx context.Context, opts ...gax.CallOption) error" )
136+ p ("" )
137+ p (" // Proto returns the long-running operation message." )
138+ p (" Proto() %s" , ptyp )
139+ p ("}" )
140+ p ("" )
141+ g .imports [pbinfo.ImportSpec {Path : "context" }] = true
142+ g .imports [pbinfo.ImportSpec {Name : "gax" , Path : "github.com/googleapis/gax-go/v2" }] = true
143+
144+ for _ , handle := range op .handles {
145+ s := pbinfo .ReduceServName (handle .GetName (), opImp .Name )
146+ n := lowerFirst (s + "Handle" )
147+
148+ // Look up polling method and its input.
149+ var get * descriptor.MethodDescriptorProto
150+ for _ , m := range handle .GetMethod () {
151+ if m .GetName () == "Get" {
152+ get = m
153+ break
154+ }
155+ }
156+ getInput := g .descInfo .Type [get .GetInputType ()]
157+ inNameField := operationResponseField (getInput .(* descriptor.DescriptorProto ), opNameField .GetName ())
158+
159+ // type
160+ p ("// Implements the %s interface for %s." , handleInt , handle .GetName ())
161+ p ("type %s struct {" , n )
162+ p (" c *%sClient" , s )
163+ p (" proto %s" , ptyp )
164+ p ("}" )
165+ p ("" )
166+
167+ // Poll
168+ p ("// Poll retrieves the latest data for the long-running operation." )
169+ p ("func (h *%s) Poll(ctx context.Context, opts ...gax.CallOption) error {" , n )
170+ p (" resp, err := h.c.Get(ctx, &%s.%s{%s: h.proto%s}, opts...)" , opImp .Name , upperFirst (getInput .GetName ()), upperFirst (inNameField .GetName ()), opNameGetter )
171+ p (" if err != nil {" )
172+ p (" return err" )
173+ p (" }" )
174+ p (" h.proto = resp" )
175+ p (" return nil" )
176+ p ("}" )
177+ p ("" )
178+
179+ // Proto
180+ p ("// Proto returns the raw type this wraps." )
181+ p ("func (h *%s) Proto() %s {" , n , ptyp )
182+ p (" return h.proto" )
183+ p ("}" )
184+ p ("" )
185+ }
186+
187+ return nil
188+ }
189+
190+ // loadCustomOpServices maps the service declared as a google.cloud.operation_service
191+ // to the service that owns the method(s) declaring it.
192+ func (g * generator ) loadCustomOpServices (servs []* descriptor.ServiceDescriptorProto ) {
193+ handles := g .aux .customOp .handles
194+ for _ , serv := range servs {
195+ for _ , meth := range serv .GetMethod () {
196+ if opServ := g .customOpService (meth ); opServ != nil {
197+ g .customOpServices [serv ] = opServ
198+ if ! containsService (handles , opServ ) {
199+ handles = append (handles , opServ )
200+ }
201+ break
202+ }
203+ }
204+ }
205+ g .aux .customOp .handles = handles
206+ }
207+
208+ // customOpService loads the ServiceDescriptorProto for the google.cloud.operation_service
209+ // named on the given method.
210+ func (g * generator ) customOpService (m * descriptor.MethodDescriptorProto ) * descriptor.ServiceDescriptorProto {
211+ opServName := proto .GetExtension (m .GetOptions (), extendedops .E_OperationService ).(string )
212+ if opServName == "" {
213+ return nil
214+ }
215+
216+ file := g .descInfo .ParentFile [m ]
217+ fqn := fmt .Sprintf (".%s.%s" , file .GetPackage (), opServName )
218+
219+ return g .descInfo .Serv [fqn ]
220+ }
221+
222+ // customOpStatusCheck constructs a return statement that checks if the operation's Status
223+ // field indicates it is done.
224+ func (g * generator ) customOpStatusCheck (st * descriptor.FieldDescriptorProto ) string {
225+ ret := fmt .Sprintf ("return o.Proto()%s" , fieldGetter (st .GetName ()))
226+ if st .GetType () == descriptor .FieldDescriptorProto_TYPE_ENUM {
227+ done := g .customOpStatusEnumDone ()
228+ ret = fmt .Sprintf ("%s == %s" , ret , done )
229+ }
230+
231+ return ret
232+ }
233+
234+ // customOpStatusEnumDone constructs the Go name of the operation's status enum
235+ // DONE value.
236+ func (g * generator ) customOpStatusEnumDone () string {
237+ op := g .aux .customOp .message
238+
239+ // Ignore the error here, it would failed much earlier if the
240+ // operation type was unresolvable.
241+ _ , imp , _ := g .descInfo .NameSpec (op )
99242
243+ // Ignore the nil case here, it would failed earlier if the
244+ // status field was not present.
245+ statusField := operationField (op , extendedops .OperationResponseMapping_STATUS )
246+ statusEnum := g .descInfo .Type [statusField .GetTypeName ()]
247+
248+ enum := fmt .Sprintf ("%s_DONE" , g .nestedName (g .descInfo .ParentElement [statusEnum ]))
249+
250+ s := fmt .Sprintf ("%s.%s" , imp .Name , enum )
251+
252+ return s
253+ }
254+
255+ // operationField is a helper for loading the target google.cloud.operation_field annotation value
256+ // if present on a field in the given message.
257+ func operationField (m * descriptor.DescriptorProto , target extendedops.OperationResponseMapping ) * descriptor.FieldDescriptorProto {
258+ for _ , f := range m .GetField () {
259+ mapping := proto .GetExtension (f .GetOptions (), extendedops .E_OperationField ).(extendedops.OperationResponseMapping )
260+ if mapping == target {
261+ return f
262+ }
263+ }
100264 return nil
101265}
266+
267+ // operationResponseField is a helper for finding the message field that declares the target field name
268+ // in the google.cloud.operation_response_field annotation.
269+ func operationResponseField (m * descriptor.DescriptorProto , target string ) * descriptor.FieldDescriptorProto {
270+ for _ , f := range m .GetField () {
271+ mapping := proto .GetExtension (f .GetOptions (), extendedops .E_OperationResponseField ).(string )
272+ if mapping == target {
273+ return f
274+ }
275+ }
276+ return nil
277+ }
278+
279+ // handleName is a helper for constructing a operation handle name from the
280+ // operation service name and Go package name.
281+ func handleName (s , pkg string ) string {
282+ s = pbinfo .ReduceServName (s , pkg )
283+ return lowerFirst (s + "Handle" )
284+ }
0 commit comments