Skip to content

Commit 35c1b96

Browse files
gregfurmangregfurman
committed
add(sdk/go): Context to workflow requests (#2765)
1 parent 8d035d7 commit 35c1b96

11 files changed

Lines changed: 50 additions & 47 deletions

File tree

‎cmd/hatchet-cli/cli/trigger.go‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -620,19 +620,19 @@ func triggerWorkflowWithClient(hatchetClient client.Client, workflowName string,
620620
}
621621
resultChan := make(chan result, 1)
622622

623+
// Wait for result with timeout
624+
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
625+
defer cancel()
626+
623627
go func() {
624-
workflow, err := hatchetClient.Admin().RunWorkflow(workflowName, inputData)
628+
workflow, err := hatchetClient.Admin().RunWorkflow(ctx, workflowName, inputData)
625629
if err != nil {
626630
resultChan <- result{err: fmt.Errorf("could not trigger workflow: %w", err)}
627631
return
628632
}
629633
resultChan <- result{runID: workflow.RunId()}
630634
}()
631635

632-
// Wait for result with timeout
633-
ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
634-
defer cancel()
635-
636636
select {
637637
case res := <-resultChan:
638638
if res.err != nil {

‎pkg/client/admin.go‎

Lines changed: 31 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -77,22 +77,22 @@ type RunDetails struct {
7777
type AdminClient interface {
7878
// Deprecated: PutWorkflow is part of the legacy v0 workflow definition system.
7979
// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go
80-
PutWorkflow(workflow *types.Workflow, opts ...PutOptFunc) error
80+
PutWorkflow(ctx context.Context, workflow *types.Workflow, opts ...PutOptFunc) error //nolint:staticcheck // SA1019: Retained for backwards compatability
8181
// Deprecated: PutWorkflowV1 is an internal method used by the new Go SDK.
8282
// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead of calling this directly. Migration guide: https://docs.hatchet.run/home/migration-guide-go
83-
PutWorkflowV1(workflow *v1contracts.CreateWorkflowVersionRequest, opts ...PutOptFunc) error
83+
PutWorkflowV1(ctx context.Context, workflow *v1contracts.CreateWorkflowVersionRequest, opts ...PutOptFunc) error
8484

85-
ScheduleWorkflow(workflowName string, opts ...ScheduleOptFunc) error
85+
ScheduleWorkflow(ctx context.Context, workflowName string, opts ...ScheduleOptFunc) error
8686

8787
// RunWorkflow triggers a workflow run and returns the run id
88-
RunWorkflow(workflowName string, input interface{}, opts ...RunOptFunc) (*Workflow, error)
88+
RunWorkflow(ctx context.Context, workflowName string, input interface{}, opts ...RunOptFunc) (*Workflow, error)
8989

90-
BulkRunWorkflow(workflows []*WorkflowRun) ([]string, error)
90+
BulkRunWorkflow(ctx context.Context, workflows []*WorkflowRun) ([]string, error)
9191

92-
RunChildWorkflow(workflowName string, input interface{}, opts *ChildWorkflowOpts) (string, error)
93-
RunChildWorkflows(workflows []*RunChildWorkflowsOpts) ([]string, error)
92+
RunChildWorkflow(ctx context.Context, workflowName string, input interface{}, opts *ChildWorkflowOpts) (string, error)
93+
RunChildWorkflows(ctx context.Context, workflows []*RunChildWorkflowsOpts) ([]string, error)
9494

95-
PutRateLimit(key string, opts *types.RateLimitOpts) error
95+
PutRateLimit(ctx context.Context, key string, opts *types.RateLimitOpts) error
9696

9797
GetRunDetails(ctx context.Context, externalId uuid.UUID) (*RunDetails, error)
9898
}
@@ -149,7 +149,7 @@ func defaultPutOpts() *putOpts {
149149

150150
// Deprecated: PutWorkflow is part of the legacy v0 workflow definition system.
151151
// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead. Migration guide: https://docs.hatchet.run/home/migration-guide-go
152-
func (a *adminClientImpl) PutWorkflow(workflow *types.Workflow, fs ...PutOptFunc) error {
152+
func (a *adminClientImpl) PutWorkflow(ctx context.Context, workflow *types.Workflow, fs ...PutOptFunc) error {
153153
opts := defaultPutOpts()
154154

155155
for _, f := range fs {
@@ -162,7 +162,7 @@ func (a *adminClientImpl) PutWorkflow(workflow *types.Workflow, fs ...PutOptFunc
162162
return fmt.Errorf("could not get put opts: %w", err)
163163
}
164164

165-
_, err = a.client.PutWorkflow(a.ctx.newContext(context.Background()), req)
165+
_, err = a.client.PutWorkflow(a.ctx.newContext(ctx), req)
166166

167167
if err != nil {
168168
return fmt.Errorf("could not create workflow %s: %w", workflow.Name, err)
@@ -173,14 +173,14 @@ func (a *adminClientImpl) PutWorkflow(workflow *types.Workflow, fs ...PutOptFunc
173173

174174
// Deprecated: PutWorkflowV1 is an internal method used by the new Go SDK.
175175
// Use the new Go SDK at github.com/hatchet-dev/hatchet/sdks/go instead of calling this directly. Migration guide: https://docs.hatchet.run/home/migration-guide-go
176-
func (a *adminClientImpl) PutWorkflowV1(workflow *v1contracts.CreateWorkflowVersionRequest, fs ...PutOptFunc) error {
176+
func (a *adminClientImpl) PutWorkflowV1(ctx context.Context, workflow *v1contracts.CreateWorkflowVersionRequest, fs ...PutOptFunc) error {
177177
opts := defaultPutOpts()
178178

179179
for _, f := range fs {
180180
f(opts)
181181
}
182182

183-
_, err := a.v1Client.PutWorkflow(a.ctx.newContext(context.Background()), workflow)
183+
_, err := a.v1Client.PutWorkflow(a.ctx.newContext(ctx), workflow)
184184

185185
if err != nil {
186186
return fmt.Errorf("could not create workflow %s: %w", workflow.Name, err)
@@ -213,7 +213,7 @@ func defaultScheduleOpts() *scheduleOpts {
213213
return &scheduleOpts{}
214214
}
215215

216-
func (a *adminClientImpl) ScheduleWorkflow(workflowName string, fs ...ScheduleOptFunc) error {
216+
func (a *adminClientImpl) ScheduleWorkflow(ctx context.Context, workflowName string, fs ...ScheduleOptFunc) error {
217217
opts := defaultScheduleOpts()
218218

219219
for _, f := range fs {
@@ -238,7 +238,7 @@ func (a *adminClientImpl) ScheduleWorkflow(workflowName string, fs ...ScheduleOp
238238

239239
workflowName = client.ApplyNamespace(workflowName, &a.namespace)
240240

241-
_, err = a.client.ScheduleWorkflow(a.ctx.newContext(context.Background()), &admincontracts.ScheduleWorkflowRequest{
241+
_, err = a.client.ScheduleWorkflow(a.ctx.newContext(ctx), &admincontracts.ScheduleWorkflowRequest{
242242
Name: workflowName,
243243
Schedules: pbSchedules,
244244
Input: string(inputBytes),
@@ -285,7 +285,7 @@ func WithPriority(priority int32) RunOptFunc {
285285
// }
286286
// }
287287

288-
func (a *adminClientImpl) RunWorkflow(workflowName string, input interface{}, options ...RunOptFunc) (*Workflow, error) {
288+
func (a *adminClientImpl) RunWorkflow(ctx context.Context, workflowName string, input interface{}, options ...RunOptFunc) (*Workflow, error) {
289289
inputBytes, err := json.Marshal(input)
290290

291291
if err != nil {
@@ -306,7 +306,7 @@ func (a *adminClientImpl) RunWorkflow(workflowName string, input interface{}, op
306306
}
307307
}
308308

309-
res, err := a.client.TriggerWorkflow(a.ctx.newContext(context.Background()), request)
309+
res, err := a.client.TriggerWorkflow(a.ctx.newContext(ctx), request)
310310

311311
if err != nil {
312312
if status.Code(err) == codes.AlreadyExists {
@@ -330,7 +330,7 @@ func (a *adminClientImpl) RunWorkflow(workflowName string, input interface{}, op
330330
}, nil
331331
}
332332

333-
func (a *adminClientImpl) BulkRunWorkflow(workflows []*WorkflowRun) ([]string, error) {
333+
func (a *adminClientImpl) BulkRunWorkflow(ctx context.Context, workflows []*WorkflowRun) ([]string, error) {
334334

335335
triggerWorkflowRequests := make([]*admincontracts.TriggerWorkflowRequest, len(workflows))
336336

@@ -358,7 +358,7 @@ func (a *adminClientImpl) BulkRunWorkflow(workflows []*WorkflowRun) ([]string, e
358358
Workflows: triggerWorkflowRequests,
359359
}
360360

361-
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(context.Background()), &r)
361+
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(ctx), &r)
362362

363363
if err != nil {
364364
return nil, fmt.Errorf("could not bulk trigger workflows: %w", err)
@@ -368,7 +368,7 @@ func (a *adminClientImpl) BulkRunWorkflow(workflows []*WorkflowRun) ([]string, e
368368

369369
}
370370

371-
func (a *adminClientImpl) RunChildWorkflow(workflowName string, input interface{}, opts *ChildWorkflowOpts) (string, error) {
371+
func (a *adminClientImpl) RunChildWorkflow(ctx context.Context, workflowName string, input interface{}, opts *ChildWorkflowOpts) (string, error) {
372372
inputBytes, err := json.Marshal(input)
373373

374374
if err != nil {
@@ -387,7 +387,7 @@ func (a *adminClientImpl) RunChildWorkflow(workflowName string, input interface{
387387

388388
metadata := string(metadataBytes)
389389

390-
res, err := a.client.TriggerWorkflow(a.ctx.newContext(context.Background()), &admincontracts.TriggerWorkflowRequest{
390+
res, err := a.client.TriggerWorkflow(a.ctx.newContext(ctx), &admincontracts.TriggerWorkflowRequest{
391391
Name: workflowName,
392392
Input: string(inputBytes),
393393
ParentId: &opts.ParentId,
@@ -419,7 +419,7 @@ type RunChildWorkflowsOpts struct {
419419
Opts *ChildWorkflowOpts
420420
}
421421

422-
func (a *adminClientImpl) RunChildWorkflows(workflows []*RunChildWorkflowsOpts) ([]string, error) {
422+
func (a *adminClientImpl) RunChildWorkflows(ctx context.Context, workflows []*RunChildWorkflowsOpts) ([]string, error) {
423423

424424
triggerWorkflowRequests := make([]*admincontracts.TriggerWorkflowRequest, len(workflows))
425425

@@ -463,7 +463,7 @@ func (a *adminClientImpl) RunChildWorkflows(workflows []*RunChildWorkflowsOpts)
463463

464464
}
465465

466-
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(context.Background()), &admincontracts.BulkTriggerWorkflowRequest{
466+
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(ctx), &admincontracts.BulkTriggerWorkflowRequest{
467467
Workflows: triggerWorkflowRequests,
468468
})
469469

@@ -475,7 +475,7 @@ func (a *adminClientImpl) RunChildWorkflows(workflows []*RunChildWorkflowsOpts)
475475
return res.WorkflowRunIds, nil
476476
}
477477

478-
func (a *adminClientImpl) PutRateLimit(key string, opts *types.RateLimitOpts) error {
478+
func (a *adminClientImpl) PutRateLimit(ctx context.Context, key string, opts *types.RateLimitOpts) error {
479479
if err := a.v.Validate(opts); err != nil {
480480
return fmt.Errorf("could not validate rate limit opts: %w", err)
481481
}
@@ -496,7 +496,7 @@ func (a *adminClientImpl) PutRateLimit(key string, opts *types.RateLimitOpts) er
496496
putParams.Duration = admincontracts.RateLimitDuration_SECOND
497497
}
498498

499-
_, err := a.client.PutRateLimit(a.ctx.newContext(context.Background()), putParams)
499+
_, err := a.client.PutRateLimit(a.ctx.newContext(ctx), putParams)
500500

501501
if err != nil {
502502
return fmt.Errorf("could not upsert rate limit: %w", err)
@@ -762,21 +762,21 @@ func (a *adminClientImpl) getAdditionalMetaBytes(opt *map[string]string) ([]byte
762762
return metadataBytes, nil
763763
}
764764

765-
func (h *adminClientImpl) saveOrLoadListener() (*WorkflowRunsListener, error) {
766-
h.listenerMu.Lock()
767-
defer h.listenerMu.Unlock()
765+
func (a *adminClientImpl) saveOrLoadListener() (*WorkflowRunsListener, error) {
766+
a.listenerMu.Lock()
767+
defer a.listenerMu.Unlock()
768768

769-
if h.listener != nil {
770-
return h.listener, nil
769+
if a.listener != nil {
770+
return a.listener, nil
771771
}
772772

773-
listener, err := h.subscriber.SubscribeToWorkflowRunEvents(context.Background())
773+
listener, err := a.subscriber.SubscribeToWorkflowRunEvents(context.Background())
774774

775775
if err != nil {
776776
return nil, fmt.Errorf("failed to subscribe to workflow run events: %w", err)
777777
}
778778

779-
h.listener = listener
779+
a.listener = listener
780780

781781
return listener, nil
782782
}

‎pkg/client/client.go‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -384,7 +384,7 @@ func newFromOpts(opts *ClientOpts) (Client, error) {
384384

385385
// if init workflows is set, then we need to initialize the workflows
386386
if opts.initWorkflows {
387-
if err := initWorkflows(opts.filesLoader, admin); err != nil {
387+
if err := initWorkflows(context.TODO(), opts.filesLoader, admin); err != nil {
388388
return nil, fmt.Errorf("could not init workflows: %w", err)
389389
}
390390
}
@@ -460,11 +460,11 @@ func (c *clientImpl) RunnableActions() []string {
460460
return c.runnableActions
461461
}
462462

463-
func initWorkflows(fl filesLoaderFunc, adminClient AdminClient) error {
463+
func initWorkflows(ctx context.Context, fl filesLoaderFunc, adminClient AdminClient) error {
464464
files := fl()
465465

466466
for _, file := range files {
467-
if err := adminClient.PutWorkflow(file); err != nil {
467+
if err := adminClient.PutWorkflow(ctx, file); err != nil {
468468
return fmt.Errorf("could not create workflow: %w", err)
469469
}
470470
}

‎pkg/v1/features/ratelimits.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ func NewRateLimitsClient(
5757
//
5858
// Upsert creates or updates a rate limit with the provided options.
5959
func (c *rlClientImpl) Upsert(opts CreateRatelimitOpts) error {
60-
return (*c.admin).PutRateLimit(opts.Key, &types.RateLimitOpts{
60+
return (*c.admin).PutRateLimit(context.Background(), opts.Key, &types.RateLimitOpts{
6161
Max: opts.Limit,
6262
Duration: opts.Duration,
6363
})

‎pkg/v1/workflow/declaration.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -542,7 +542,7 @@ func (w *workflowDeclarationImpl[I, O]) RunBulkNoWait(ctx context.Context, input
542542
}
543543
}
544544

545-
run, err := w.v0.Admin().BulkRunWorkflow(toRun)
545+
run, err := w.v0.Admin().BulkRunWorkflow(ctx, toRun)
546546
if err != nil {
547547
return nil, err
548548
}
@@ -556,7 +556,7 @@ func (w *workflowDeclarationImpl[I, O]) RunBulkNoWait(ctx context.Context, input
556556
// RunNoWait executes the workflow with the provided input without waiting for it to complete.
557557
// Instead it returns a run ID that can be used to check the status of the workflow.
558558
func (w *workflowDeclarationImpl[I, O]) RunNoWait(ctx context.Context, input I, opts ...v0Client.RunOptFunc) (*v0Client.Workflow, error) {
559-
run, err := w.v0.Admin().RunWorkflow(w.Name, input, opts...)
559+
run, err := w.v0.Admin().RunWorkflow(ctx, w.Name, input, opts...)
560560
if err != nil {
561561
return nil, err
562562
}

‎pkg/worker/context.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -512,6 +512,7 @@ func (h *hatchetContext) SpawnWorkflow(workflowName string, input any, opts *Spa
512512
h.indexMu.Unlock()
513513

514514
workflowRunId, err := h.client().Admin().RunChildWorkflow(
515+
h,
515516
workflowName,
516517
input,
517518
&client.ChildWorkflowOpts{
@@ -590,6 +591,7 @@ func (h *hatchetContext) SpawnWorkflows(childWorkflows []*SpawnWorkflowsOpts) ([
590591
}
591592

592593
workflowRunIds, err := h.client().Admin().RunChildWorkflows(
594+
h,
593595
triggerWorkflows,
594596
)
595597

‎pkg/worker/service.go‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package worker
22

33
import (
4+
"context"
45
"fmt"
56

67
"github.com/hatchet-dev/hatchet/pkg/client/compute"
@@ -43,7 +44,7 @@ func (s *Service) On(t triggerConverter, workflow workflowConverter) error {
4344
apiWorkflow.Triggers = *wt
4445

4546
// create the workflow via the API
46-
err := s.worker.client.Admin().PutWorkflow(&apiWorkflow)
47+
err := s.worker.client.Admin().PutWorkflow(context.TODO(), &apiWorkflow)
4748

4849
if err != nil {
4950
return err

‎pkg/worker/worker.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -433,7 +433,7 @@ func (w *Worker) RegisterWorkflowV1(workflow *contracts.CreateWorkflowVersionReq
433433

434434
w.registered_workflows[namespaced] = true
435435

436-
return w.client.Admin().PutWorkflowV1(workflow)
436+
return w.client.Admin().PutWorkflowV1(context.TODO(), workflow)
437437
}
438438

439439
// Deprecated: On is part of the legacy v0 workflow definition system.

‎sdks/go/client.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -640,7 +640,7 @@ func (c *Client) RunNoWait(ctx context.Context, workflowName string, input any,
640640
AdditionalMetadata: additionalMetadata,
641641
})
642642
} else {
643-
v0Workflow, err = c.legacyClient.Admin().RunWorkflow(workflowName, input, v0Opts...)
643+
v0Workflow, err = c.legacyClient.Admin().RunWorkflow(ctx, workflowName, input, v0Opts...)
644644
}
645645

646646
if err != nil {

‎sdks/go/features/ratelimits.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,7 @@ func NewRateLimitsClient(
4646

4747
// Upsert creates or updates a rate limit with the provided options.
4848
func (c *RateLimitsClient) Upsert(opts CreateRatelimitOpts) error {
49-
if err := c.admin.PutRateLimit(opts.Key, &types.RateLimitOpts{
49+
if err := c.admin.PutRateLimit(context.TODO(), opts.Key, &types.RateLimitOpts{
5050
Max: opts.Limit,
5151
Duration: opts.Duration,
5252
}); err != nil {

0 commit comments

Comments
 (0)