feat: add context in module interface (#898)

KusionStack · Mar 11, 2024 · cf74a3a · cf74a3a
1 parent f1ba807
commit cf74a3a
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 6 deletions.
diff --git a/pkg/modules/generators/app_configurations_generator.go b/pkg/modules/generators/app_configurations_generator.go
@@ -1,6 +1,7 @@
 package generators
 
 import (
+	"context"
 	"errors"
 	"fmt"
 
@@ -164,7 +165,7 @@ func (g *appConfigurationGenerator) callModules(projectModuleConfigs map[string]
 
 		// invoke the plugin
 		log.Infof("invoke module:%s with request:%s", t, protoRequest.String())
-		response, err := plugin.Module.Generate(protoRequest)
+		response, err := plugin.Module.Generate(context.Background(), protoRequest)
 		if err != nil {
 			return nil, fmt.Errorf("invoke kusion module: %s failed. %w", t, err)
 		}

diff --git a/pkg/modules/generators/app_configurations_generator_test.go b/pkg/modules/generators/app_configurations_generator_test.go
@@ -1,6 +1,7 @@
 package generators
 
 import (
+	"context"
 	"testing"
 
 	"github.com/bytedance/mockey"
@@ -53,7 +54,7 @@ func TestAppConfigurationGenerator_Generate(t *testing.T) {
 
 type fakeModule struct{}
 
-func (f *fakeModule) Generate(req *proto.GeneratorRequest) (*proto.GeneratorResponse, error) {
+func (f *fakeModule) Generate(ctx context.Context, req *proto.GeneratorRequest) (*proto.GeneratorResponse, error) {
 	res := v1.Resource{
 		ID:   "apps.kusionstack.io/v1alpha1:PodTransitionRule:fakeNs:default-dev-foo",
 		Type: "Kubernetes",

diff --git a/pkg/modules/interfaces.go b/pkg/modules/interfaces.go
@@ -22,7 +22,7 @@ type Generator interface {
 
 // Module is the interface that we're exposing as a kusion module plugin.
 type Module interface {
-	Generate(req *proto.GeneratorRequest) (*proto.GeneratorResponse, error)
+	Generate(ctx context.Context, req *proto.GeneratorRequest) (*proto.GeneratorResponse, error)
 }
 
 // NewGeneratorFunc is a function that returns a Generator.
@@ -32,8 +32,8 @@ type GRPCClient struct {
 	client proto.ModuleClient
 }
 
-func (c *GRPCClient) Generate(req *proto.GeneratorRequest) (*proto.GeneratorResponse, error) {
-	return c.client.Generate(context.Background(), req)
+func (c *GRPCClient) Generate(ctx context.Context, req *proto.GeneratorRequest) (*proto.GeneratorResponse, error) {
+	return c.client.Generate(ctx, req)
 }
 
 type GRPCServer struct {
@@ -49,7 +49,7 @@ func (s *GRPCServer) Generate(ctx context.Context, req *proto.GeneratorRequest)
 			res = &proto.GeneratorResponse{}
 		}
 	}()
-	res, err = s.Impl.Generate(req)
+	res, err = s.Impl.Generate(ctx, req)
 	return
 }