github.com/kubeshop/testkube@v1.17.23/pkg/logs/adapter/cloud.go (about)

     1  package adapter
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/pkg/errors"
     8  	"go.uber.org/zap"
     9  	"google.golang.org/grpc/metadata"
    10  
    11  	"github.com/kubeshop/testkube/pkg/log"
    12  	"github.com/kubeshop/testkube/pkg/logs/events"
    13  	"github.com/kubeshop/testkube/pkg/logs/pb"
    14  )
    15  
    16  var _ Adapter = &CloudAdapter{}
    17  
    18  // NewCloudConsumer creates new CloudSubscriber which will send data to local MinIO bucket
    19  func NewCloudAdapter(grpcConn pb.CloudLogsServiceClient, agentApiKey string) *CloudAdapter {
    20  	return &CloudAdapter{
    21  		client:      grpcConn,
    22  		agentApiKey: agentApiKey,
    23  		logger:      log.DefaultLogger.With("service", "logs-cloud-adapter"),
    24  	}
    25  }
    26  
    27  type CloudAdapter struct {
    28  	client      pb.CloudLogsServiceClient
    29  	streams     sync.Map
    30  	agentApiKey string
    31  	logger      *zap.SugaredLogger
    32  }
    33  
    34  func (s *CloudAdapter) Init(ctx context.Context, id string) error {
    35  
    36  	// write metadata to the stream context
    37  	md := metadata.Pairs("api-key", s.agentApiKey, "execution-id", id)
    38  	ctx = metadata.NewOutgoingContext(ctx, md)
    39  
    40  	stream, err := s.client.Stream(ctx)
    41  	if err != nil {
    42  		return errors.Wrap(err, "can't init stream")
    43  	}
    44  
    45  	s.streams.Store(id, stream)
    46  
    47  	return nil
    48  }
    49  
    50  func (s *CloudAdapter) Notify(ctx context.Context, id string, e events.Log) error {
    51  	c, err := s.getStreamClient(id)
    52  	if err != nil {
    53  		return errors.Wrap(err, "can't get stream client for id: "+id)
    54  	}
    55  
    56  	return c.Send(pb.MapToPB(e))
    57  }
    58  
    59  func (s *CloudAdapter) Stop(ctx context.Context, id string) error {
    60  	c, err := s.getStreamClient(id)
    61  	if err != nil {
    62  		return errors.Wrap(err, "can't get stream client for id: "+id)
    63  	}
    64  
    65  	resp, err := c.CloseAndRecv()
    66  	if err != nil {
    67  		return errors.Wrap(err, "closing log stream error")
    68  	}
    69  	s.logger.Debugw("closing response", "resp", resp, "id", id)
    70  
    71  	s.streams.Delete(id)
    72  	return nil
    73  }
    74  
    75  func (s *CloudAdapter) Name() string {
    76  	return "cloud"
    77  }
    78  
    79  func (s *CloudAdapter) getStreamClient(id string) (client pb.CloudLogsService_StreamClient, err error) {
    80  	c, ok := s.streams.Load(id)
    81  	if !ok {
    82  		return nil, errors.New("can't find initialized stream")
    83  	}
    84  
    85  	return c.(pb.CloudLogsService_StreamClient), nil
    86  }