Skip to content

Commit 8b01d31

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

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),
@@ -327,7 +327,7 @@ func desiredWorkerLabelsToProto(labels map[string]*types.DesiredWorkerLabel) map
327327
// }
328328
// }
329329

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

333333
if err != nil {
@@ -348,7 +348,7 @@ func (a *adminClientImpl) RunWorkflow(workflowName string, input interface{}, op
348348
}
349349
}
350350

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

353353
if err != nil {
354354
if status.Code(err) == codes.AlreadyExists {
@@ -372,7 +372,7 @@ func (a *adminClientImpl) RunWorkflow(workflowName string, input interface{}, op
372372
}, nil
373373
}
374374

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

377377
triggerWorkflowRequests := make([]*v1contracts.TriggerWorkflowRequest, len(workflows))
378378

@@ -400,7 +400,7 @@ func (a *adminClientImpl) BulkRunWorkflow(workflows []*WorkflowRun) ([]string, e
400400
Workflows: triggerWorkflowRequests,
401401
}
402402

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

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

411411
}
412412

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

416416
if err != nil {
@@ -429,7 +429,7 @@ func (a *adminClientImpl) RunChildWorkflow(workflowName string, input interface{
429429

430430
metadata := string(metadataBytes)
431431

432-
res, err := a.client.TriggerWorkflow(a.ctx.newContext(context.Background()), &v1contracts.TriggerWorkflowRequest{
432+
res, err := a.client.TriggerWorkflow(a.ctx.newContext(ctx), &admincontracts.TriggerWorkflowRequest{
433433
Name: workflowName,
434434
Input: string(inputBytes),
435435
ParentId: &opts.ParentId,
@@ -462,7 +462,7 @@ type RunChildWorkflowsOpts struct {
462462
Opts *ChildWorkflowOpts
463463
}
464464

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

467467
triggerWorkflowRequests := make([]*v1contracts.TriggerWorkflowRequest, len(workflows))
468468

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

508508
}
509509

510-
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(context.Background()), &admincontracts.BulkTriggerWorkflowRequest{
510+
res, err := a.client.BulkTriggerWorkflow(a.ctx.newContext(ctx), &admincontracts.BulkTriggerWorkflowRequest{
511511
Workflows: triggerWorkflowRequests,
512512
})
513513

@@ -519,7 +519,7 @@ func (a *adminClientImpl) RunChildWorkflows(workflows []*RunChildWorkflowsOpts)
519519
return res.WorkflowRunIds, nil
520520
}
521521

522-
func (a *adminClientImpl) PutRateLimit(key string, opts *types.RateLimitOpts) error {
522+
func (a *adminClientImpl) PutRateLimit(ctx context.Context, key string, opts *types.RateLimitOpts) error {
523523
if err := a.v.Validate(opts); err != nil {
524524
return fmt.Errorf("could not validate rate limit opts: %w", err)
525525
}
@@ -540,7 +540,7 @@ func (a *adminClientImpl) PutRateLimit(key string, opts *types.RateLimitOpts) er
540540
putParams.Duration = admincontracts.RateLimitDuration_SECOND
541541
}
542542

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

545545
if err != nil {
546546
return fmt.Errorf("could not upsert rate limit: %w", err)
@@ -806,21 +806,21 @@ func (a *adminClientImpl) getAdditionalMetaBytes(opt *map[string]string) ([]byte
806806
return metadataBytes, nil
807807
}
808808

809-
func (h *adminClientImpl) saveOrLoadListener() (*WorkflowRunsListener, error) {
810-
h.listenerMu.Lock()
811-
defer h.listenerMu.Unlock()
809+
func (a *adminClientImpl) saveOrLoadListener() (*WorkflowRunsListener, error) {
810+
a.listenerMu.Lock()
811+
defer a.listenerMu.Unlock()
812812

813-
if h.listener != nil {
814-
return h.listener, nil
813+
if a.listener != nil {
814+
return a.listener, nil
815815
}
816816

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

819819
if err != nil {
820820
return nil, fmt.Errorf("failed to subscribe to workflow run events: %w", err)
821821
}
822822

823-
h.listener = listener
823+
a.listener = listener
824824

825825
return listener, nil
826826
}

‎pkg/client/client.go‎

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

397397
// if init workflows is set, then we need to initialize the workflows
398398
if opts.initWorkflows {
399-
if err := initWorkflows(opts.filesLoader, admin); err != nil {
399+
if err := initWorkflows(context.TODO(), opts.filesLoader, admin); err != nil {
400400
return nil, fmt.Errorf("could not init workflows: %w", err)
401401
}
402402
}
@@ -472,11 +472,11 @@ func (c *clientImpl) RunnableActions() []string {
472472
return c.runnableActions
473473
}
474474

475-
func initWorkflows(fl filesLoaderFunc, adminClient AdminClient) error {
475+
func initWorkflows(ctx context.Context, fl filesLoaderFunc, adminClient AdminClient) error {
476476
files := fl()
477477

478478
for _, file := range files {
479-
if err := adminClient.PutWorkflow(file); err != nil {
479+
if err := adminClient.PutWorkflow(ctx, file); err != nil {
480480
return fmt.Errorf("could not create workflow: %w", err)
481481
}
482482
}

‎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
@@ -539,7 +539,7 @@ func (w *workflowDeclarationImpl[I, O]) RunBulkNoWait(ctx context.Context, input
539539
}
540540
}
541541

542-
run, err := w.v0.Admin().BulkRunWorkflow(toRun)
542+
run, err := w.v0.Admin().BulkRunWorkflow(ctx, toRun)
543543
if err != nil {
544544
return nil, err
545545
}
@@ -553,7 +553,7 @@ func (w *workflowDeclarationImpl[I, O]) RunBulkNoWait(ctx context.Context, input
553553
// RunNoWait executes the workflow with the provided input without waiting for it to complete.
554554
// Instead it returns a run ID that can be used to check the status of the workflow.
555555
func (w *workflowDeclarationImpl[I, O]) RunNoWait(ctx context.Context, input I, opts ...v0Client.RunOptFunc) (*v0Client.Workflow, error) {
556-
run, err := w.v0.Admin().RunWorkflow(w.Name, input, opts...)
556+
run, err := w.v0.Admin().RunWorkflow(ctx, w.Name, input, opts...)
557557
if err != nil {
558558
return nil, err
559559
}

‎pkg/worker/context.go‎

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

516516
workflowRunId, err := h.client().Admin().RunChildWorkflow(
517+
h,
517518
workflowName,
518519
input,
519520
&client.ChildWorkflowOpts{
@@ -595,6 +596,7 @@ func (h *hatchetContext) SpawnWorkflows(childWorkflows []*SpawnWorkflowsOpts) ([
595596
}
596597

597598
workflowRunIds, err := h.client().Admin().RunChildWorkflows(
599+
h,
598600
triggerWorkflows,
599601
)
600602

‎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)