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

     1  package logs
     2  
     3  import (
     4  	"fmt"
     5  
     6  	"go.uber.org/zap"
     7  
     8  	"github.com/kubeshop/testkube/pkg/log"
     9  	"github.com/kubeshop/testkube/pkg/logs/pb"
    10  	"github.com/kubeshop/testkube/pkg/logs/repository"
    11  	"github.com/kubeshop/testkube/pkg/logs/state"
    12  )
    13  
    14  func NewLogsServer(repo repository.Factory, state state.Interface) *LogsServer {
    15  	return &LogsServer{
    16  		state:       state,
    17  		repoFactory: repo,
    18  		log:         log.DefaultLogger.With("service", "logs-grpc-server"),
    19  	}
    20  }
    21  
    22  type LogsServer struct {
    23  	pb.UnimplementedLogsServiceServer
    24  	repoFactory   repository.Factory
    25  	state         state.Interface
    26  	log           *zap.SugaredLogger
    27  	traceMessages bool
    28  }
    29  
    30  func (s LogsServer) Logs(req *pb.LogRequest, stream pb.LogsService_LogsServer) error {
    31  	ctx := stream.Context()
    32  
    33  	log := s.log.With("execution_id", req.ExecutionId)
    34  
    35  	// get state of current log stream (pending or finished)
    36  	st, err := s.state.Get(ctx, req.ExecutionId)
    37  	if err != nil {
    38  		return err
    39  	}
    40  
    41  	// get valid repository based on state
    42  	repo, err := s.repoFactory.GetRepository(st)
    43  	if err != nil {
    44  		return err
    45  	}
    46  
    47  	log.Debugw("starting sending log stream", "repo", fmt.Sprintf("%T", repo), "state", st)
    48  
    49  	// stream logs from repository through GRPC channel
    50  	ch, err := repo.Get(ctx, req.ExecutionId)
    51  	if err != nil {
    52  		return err
    53  	}
    54  
    55  	for l := range ch {
    56  		if s.traceMessages {
    57  			log.Debugw("sending log chunk", "log", l)
    58  		}
    59  		if err := stream.Send(pb.MapResponseToPB(l)); err != nil {
    60  			return err
    61  		}
    62  	}
    63  
    64  	log.Debugw("log stream finished")
    65  
    66  	return nil
    67  }
    68  
    69  func (s *LogsServer) WithMessageTracing(enabled bool) *LogsServer {
    70  	s.traceMessages = enabled
    71  	return s
    72  }