github.com/kubeshop/testkube@v1.17.23/pkg/cloud/data/executor/executor.go (about)

     1  package executor
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"math"
     7  
     8  	"google.golang.org/grpc"
     9  	"google.golang.org/grpc/encoding/gzip"
    10  	"google.golang.org/protobuf/types/known/structpb"
    11  
    12  	"github.com/kubeshop/testkube/pkg/agent"
    13  	"github.com/kubeshop/testkube/pkg/cloud"
    14  )
    15  
    16  type Command string
    17  
    18  //go:generate mockgen -destination=./mock_executor.go -package=executor "github.com/kubeshop/testkube/pkg/cloud/data/executor" Executor
    19  type Executor interface {
    20  	Execute(ctx context.Context, command Command, payload any) (response []byte, err error)
    21  	Close() error
    22  }
    23  
    24  type CloudGRPCExecutor struct {
    25  	client cloud.TestKubeCloudAPIClient
    26  	conn   *grpc.ClientConn
    27  	apiKey string
    28  }
    29  
    30  func NewCloudGRPCExecutor(client cloud.TestKubeCloudAPIClient, grpcConn *grpc.ClientConn, apiKey string) *CloudGRPCExecutor {
    31  	return &CloudGRPCExecutor{client: client, conn: grpcConn, apiKey: apiKey}
    32  }
    33  
    34  func (e *CloudGRPCExecutor) Execute(ctx context.Context, command Command, payload any) (response []byte, err error) {
    35  	jsonPayload, err := json.Marshal(payload)
    36  	if err != nil {
    37  		return nil, err
    38  	}
    39  	s := structpb.Struct{}
    40  	if err := s.UnmarshalJSON(jsonPayload); err != nil {
    41  		return nil, err
    42  	}
    43  	req := cloud.CommandRequest{
    44  		Command: string(command),
    45  		Payload: &s,
    46  	}
    47  	ctx = agent.AddAPIKeyMeta(ctx, e.apiKey)
    48  	opts := []grpc.CallOption{grpc.UseCompressor(gzip.Name), grpc.MaxCallRecvMsgSize(math.MaxInt32)}
    49  	cmdResponse, err := e.client.Call(ctx, &req, opts...)
    50  	if err != nil {
    51  		return nil, err
    52  	}
    53  	return cmdResponse.Response, nil
    54  }
    55  
    56  func (e *CloudGRPCExecutor) Close() error {
    57  	return e.conn.Close()
    58  }