feat: add context in module interface (#898)

KusionStack · Mar 11, 2024 · 729b89b · 729b89b
1 parent d845aa1
commit 729b89b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 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/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
 }