chore: add context to UsageHandler methods (#115)

Because - We want the pipeline trigger context to be propagated to the component execution and usage handler. This commit - Adds a context param to `Execute`, `Check` and `Collect`.
instill-ai · May 10, 2024 · 80c4780 · 80c4780
1 parent ed2068f
commit 80c4780
Show file tree

Hide file tree

Showing 34 changed files with 143 additions and 98 deletions.
diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md
@@ -396,7 +396,7 @@ the `hello` component will only need to override the following methods:
   will return an object that implements the `Execute` method.
   - `ExecutionWrapper` will wrap the execution call with the input and output
     schema validation.
-- `Execute([]*structpb.Struct) ([]*structpb.Struct, error)` is the most
+- `Execute(context.Context []*structpb.Struct) ([]*structpb.Struct, error)` is the most
   important function in the component. All the data manipulation will take place
   here.
 
@@ -463,7 +463,7 @@ func (o *operator) CreateExecution(sysVars map[string]any, task string) (*base.E
 	return &base.ExecutionWrapper{Execution: e}, nil
 }
 
-func (e *execution) Execute(_ []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(context.Context, []*structpb.Struct) ([]*structpb.Struct, error) {
 	return nil, nil
 }
 ```
@@ -498,7 +498,7 @@ func (o *operator) CreateExecution(sysVars map[string]any, task string) (*base.E
 	return &base.ExecutionWrapper{Execution: e}, nil
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	outputs := make([]*structpb.Struct, len(inputs))
 
 	// An execution  might take several inputs. One result will be returned for
@@ -549,6 +549,7 @@ import (
 
 func TestOperator_Execute(t *testing.T) {
 	c := qt.New(t)
+	ctx := context.Background()
 
 	bo := base.BaseOperator{Logger: zap.NewNop()}
 	operator := Init(bo)
@@ -560,7 +561,7 @@ func TestOperator_Execute(t *testing.T) {
 		pbIn, err := structpb.NewStruct(map[string]any{"target": "bolero-wombat"})
 		c.Assert(err, qt.IsNil)
 
-		got, err := exec.Execution.Execute([]*structpb.Struct{pbIn})
+		got, err := exec.Execution.Execute(ctx, []*structpb.Struct{pbIn})
 		c.Check(err, qt.IsNil)
 		c.Assert(got, qt.HasLen, 1)
 

diff --git a/pkg/base/execution.go b/pkg/base/execution.go
@@ -1,6 +1,7 @@
 package base
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"strconv"
@@ -25,11 +26,12 @@ type IExecution interface {
 	GetLogger() *zap.Logger
 	GetTaskInputSchema() string
 	GetTaskOutputSchema() string
+	GetSystemVariables() map[string]any
 
 	UsesSecret() bool
 	UsageHandlerCreator() UsageHandlerCreator
 
-	Execute([]*structpb.Struct) ([]*structpb.Struct, error)
+	Execute(context.Context, []*structpb.Struct) ([]*structpb.Struct, error)
 }
 
 func FormatErrors(inputPath string, e jsonschema.Detailed, errors *[]string) {
@@ -115,18 +117,18 @@ func Validate(data []*structpb.Struct, jsonSchema string, target string) error {
 }
 
 // Execute wraps the execution method with validation and usage collection.
-func (e *ExecutionWrapper) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *ExecutionWrapper) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	if err := Validate(inputs, e.Execution.GetTaskInputSchema(), "inputs"); err != nil {
 		return nil, err
 	}
 
 	newUH := e.Execution.UsageHandlerCreator()
 	h := newUH(e.Execution)
-	if err := h.Check(inputs); err != nil {
+	if err := h.Check(ctx, inputs); err != nil {
 		return nil, err
 	}
 
-	outputs, err := e.Execution.Execute(inputs)
+	outputs, err := e.Execution.Execute(ctx, inputs)
 	if err != nil {
 		return nil, err
 	}
@@ -135,7 +137,7 @@ func (e *ExecutionWrapper) Execute(inputs []*structpb.Struct) ([]*structpb.Struc
 		return nil, err
 	}
 
-	if err := h.Collect(inputs, outputs); err != nil {
+	if err := h.Collect(ctx, inputs, outputs); err != nil {
 		return nil, err
 	}
 

diff --git a/pkg/base/usage.go b/pkg/base/usage.go
@@ -1,21 +1,25 @@
 package base
 
-import "google.golang.org/protobuf/types/known/structpb"
+import (
+	"context"
+
+	"google.golang.org/protobuf/types/known/structpb"
+)
 
 // UsageHandler allows the component execution wrapper to add checks and
 // collect usage metrics around a component execution.
 type UsageHandler interface {
-	Check(inputs []*structpb.Struct) error
-	Collect(inputs, outputs []*structpb.Struct) error
+	Check(ctx context.Context, inputs []*structpb.Struct) error
+	Collect(ctx context.Context, inputs, outputs []*structpb.Struct) error
 }
 
 // UsageHandlerCreator returns a function to initialize a UsageHandler.
 type UsageHandlerCreator func(IExecution) UsageHandler
 
 type noopUsageHandler struct{}
 
-func (h *noopUsageHandler) Check([]*structpb.Struct) error { return nil }
-func (h *noopUsageHandler) Collect(_, _ []*structpb.Struct) error {
+func (h *noopUsageHandler) Check(context.Context, []*structpb.Struct) error { return nil }
+func (h *noopUsageHandler) Collect(_ context.Context, _, _ []*structpb.Struct) error {
 	return nil
 }
 

diff --git a/pkg/connector/airbyte/v0/main.go b/pkg/connector/airbyte/v0/main.go
@@ -1,6 +1,7 @@
 package airbyte
 
 import (
+	"context"
 	_ "embed"
 	"fmt"
 	"sync"
@@ -44,10 +45,10 @@ func (c *connector) CreateExecution(sysVars map[string]any, connection *structpb
 	}}, nil
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(context.Context, []*structpb.Struct) ([]*structpb.Struct, error) {
 	return nil, fmt.Errorf("the Airbyte connector has been removed")
 }
 
-func (c *connector) Test(sysVars map[string]any, connection *structpb.Struct) error {
+func (c *connector) Test(map[string]any, *structpb.Struct) error {
 	return fmt.Errorf("the Airbyte connector has been removed")
 }
diff --git a/pkg/connector/archetypeai/v0/connector_test.go b/pkg/connector/archetypeai/v0/connector_test.go
@@ -1,6 +1,7 @@
 package archetypeai
 
 import (
+	"context"
 	"fmt"
 	"io"
 	"net/http"
@@ -100,6 +101,7 @@ var (
 
 func TestConnector_Execute(t *testing.T) {
 	c := qt.New(t)
+	ctx := context.Background()
 
 	testcases := []struct {
 		name string
@@ -265,7 +267,7 @@ func TestConnector_Execute(t *testing.T) {
 			pbIn, err := base.ConvertToStructpb(tc.in)
 			c.Assert(err, qt.IsNil)
 
-			got, err := exec.Execution.Execute([]*structpb.Struct{pbIn})
+			got, err := exec.Execution.Execute(ctx, []*structpb.Struct{pbIn})
 			if tc.wantErr != "" {
 				c.Check(errmsg.Message(err), qt.Matches, tc.wantErr)
 				return

diff --git a/pkg/connector/archetypeai/v0/main.go b/pkg/connector/archetypeai/v0/main.go
@@ -3,6 +3,7 @@ package archetypeai
 
 import (
 	"bytes"
+	"context"
 	_ "embed"
 	"fmt"
 	"strings"
@@ -81,7 +82,7 @@ func (c *connector) CreateExecution(sysVars map[string]any, connection *structpb
 }
 
 // Execute performs calls the Archetype AI API to execute a task.
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	outputs := make([]*structpb.Struct, len(inputs))
 
 	for i, input := range inputs {

diff --git a/pkg/connector/bigquery/v0/main.go b/pkg/connector/bigquery/v0/main.go
@@ -70,7 +70,7 @@ func getTableName(config *structpb.Struct) string {
 	return config.GetFields()["table_name"].GetStringValue()
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	outputs := []*structpb.Struct{}
 
 	client, err := NewClient(getJSONKey(e.Connection), getProjectID(e.Connection))

diff --git a/pkg/connector/googlecloudstorage/v0/main.go b/pkg/connector/googlecloudstorage/v0/main.go
@@ -63,7 +63,7 @@ func getJSONKey(config *structpb.Struct) string {
 	return config.GetFields()["json_key"].GetStringValue()
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	outputs := []*structpb.Struct{}
 
 	client, err := NewClient(getJSONKey(e.Connection))

diff --git a/pkg/connector/googlesearch/v0/main.go b/pkg/connector/googlesearch/v0/main.go
@@ -66,7 +66,7 @@ func getSearchEngineID(config *structpb.Struct) string {
 	return config.GetFields()["cse_id"].GetStringValue()
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 
 	service, err := NewService(getAPIKey(e.Connection))
 	if err != nil || service == nil {

diff --git a/pkg/connector/huggingface/v0/connector_test.go b/pkg/connector/huggingface/v0/connector_test.go
@@ -1,6 +1,7 @@
 package huggingface
 
 import (
+	"context"
 	"encoding/base64"
 	"fmt"
 	"io"
@@ -222,6 +223,7 @@ func TestConnector_ExecuteSpeechRecognition(t *testing.T) {
 func testTask(c *qt.C, p taskParams) {
 	bc := base.Connector{Logger: zap.NewNop()}
 	connector := Init(bc)
+	ctx := context.Background()
 
 	c.Run("nok - HTTP client error - "+p.task, func(c *qt.C) {
 		c.Parallel()
@@ -238,7 +240,7 @@ func testTask(c *qt.C, p taskParams) {
 		c.Assert(err, qt.IsNil)
 		pbIn.Fields["model"] = structpb.NewStringValue(model)
 
-		_, err = exec.Execution.Execute([]*structpb.Struct{pbIn})
+		_, err = exec.Execution.Execute(ctx, []*structpb.Struct{pbIn})
 		c.Check(err, qt.IsNotNil)
 		c.Check(err, qt.ErrorMatches, ".*no such host")
 		c.Check(errmsg.Message(err), qt.Matches, "Failed to call .*check that the connector configuration is correct.")
@@ -326,7 +328,7 @@ func testTask(c *qt.C, p taskParams) {
 			c.Assert(err, qt.IsNil)
 			pbIn.Fields["model"] = structpb.NewStringValue(model)
 
-			got, err := exec.Execution.Execute([]*structpb.Struct{pbIn})
+			got, err := exec.Execution.Execute(ctx, []*structpb.Struct{pbIn})
 			if tc.wantErr != "" {
 				c.Check(err, qt.IsNotNil)
 				c.Check(errmsg.Message(err), qt.Equals, tc.wantErr)

diff --git a/pkg/connector/huggingface/v0/main.go b/pkg/connector/huggingface/v0/main.go
@@ -2,6 +2,7 @@
 package huggingface
 
 import (
+	"context"
 	_ "embed"
 	"encoding/base64"
 	"encoding/json"
@@ -101,7 +102,7 @@ func wrapSliceInStruct(data []byte, key string) (*structpb.Struct, error) {
 	}, nil
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	client := newClient(e.Connection, e.GetLogger())
 	outputs := []*structpb.Struct{}
 

diff --git a/pkg/connector/instill/v0/main.go b/pkg/connector/instill/v0/main.go
@@ -98,7 +98,7 @@ func getRequestMetadata(vars map[string]any) metadata.MD {
 	)
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(ctx context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 	var err error
 
 	if len(inputs) <= 0 || inputs[0] == nil {
@@ -116,7 +116,7 @@ func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, erro
 	}
 
 	modelNameSplits := strings.Split(inputs[0].GetFields()["model_name"].GetStringValue(), "/")
-	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute)
+	ctx, cancel := context.WithTimeout(ctx, 5*time.Minute)
 	defer cancel()
 
 	ctx = metadata.NewOutgoingContext(ctx, getRequestMetadata(e.SystemVariables))

diff --git a/pkg/connector/numbers/v0/main.go b/pkg/connector/numbers/v0/main.go
@@ -3,6 +3,7 @@ package numbers
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -191,7 +192,7 @@ func (e *execution) registerAsset(data []byte, reg Register) (string, error) {
 	}
 }
 
-func (e *execution) Execute(inputs []*structpb.Struct) ([]*structpb.Struct, error) {
+func (e *execution) Execute(_ context.Context, inputs []*structpb.Struct) ([]*structpb.Struct, error) {
 
 	var outputs []*structpb.Struct