github.com/blend/go-sdk@v1.20220411.3/logger/shim_writer.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package logger
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"io"
    14  	"strings"
    15  )
    16  
    17  var (
    18  	_ io.Writer = (*ShimWriter)(nil)
    19  )
    20  
    21  // Constants
    22  const (
    23  	DefaultShimWriterMessageFlag = "shim"
    24  )
    25  
    26  // NewShimWriter returns a new shim writer.
    27  // A "Shim Writer" is mean to bridge situations where you need to pass
    28  // an io.Writer to a given function, and want that function to write to a logger.
    29  // I.e. you can set `cmd.Stdout = NewShimWriter(log)` to have a
    30  // shell command write to a logger for standard out.
    31  func NewShimWriter(log Triggerable, opts ...ShimWriterOption) ShimWriter {
    32  	shim := ShimWriter{
    33  		Context:       context.Background(),
    34  		Log:           log,
    35  		EventProvider: ShimWriterMessageEventProvider(DefaultShimWriterMessageFlag),
    36  	}
    37  	for _, opt := range opts {
    38  		opt(&shim)
    39  	}
    40  	return shim
    41  }
    42  
    43  // OptShimWriterEventProvider sets the event provider for the shim writer.
    44  func OptShimWriterEventProvider(provider func([]byte) Event) ShimWriterOption {
    45  	return func(sw *ShimWriter) { sw.EventProvider = provider }
    46  }
    47  
    48  // OptShimWriterContext sets the context for a given shim writer.
    49  func OptShimWriterContext(ctx context.Context) ShimWriterOption {
    50  	return func(sw *ShimWriter) { sw.Context = ctx }
    51  }
    52  
    53  // ShimWriterMessageEventProvider returns a message event with a given flag
    54  // for a given contents.
    55  func ShimWriterMessageEventProvider(flag string, opts ...MessageEventOption) func([]byte) Event {
    56  	return func(contents []byte) Event {
    57  		return NewMessageEvent(flag, strings.TrimSpace(string(contents)), opts...)
    58  	}
    59  }
    60  
    61  // ShimWriterErrorEventProvider returns an error event with a given flag
    62  // for a given contents.
    63  func ShimWriterErrorEventProvider(flag string, opts ...ErrorEventOption) func([]byte) Event {
    64  	return func(contents []byte) Event {
    65  		return NewErrorEvent(flag, fmt.Errorf(strings.TrimSpace(string(contents))), opts...)
    66  	}
    67  }
    68  
    69  // ShimWriterOption is a mutator for a shim writer.
    70  type ShimWriterOption func(*ShimWriter)
    71  
    72  // ShimWriter is a type that implements io.Writer with
    73  // a logger backend.
    74  type ShimWriter struct {
    75  	Context       context.Context
    76  	Log           Triggerable
    77  	EventProvider func([]byte) Event
    78  }
    79  
    80  // Write implements io.Writer.
    81  func (sw ShimWriter) Write(contents []byte) (count int, err error) {
    82  	ctx := sw.Context
    83  	if ctx == nil {
    84  		ctx = context.Background()
    85  	}
    86  	sw.Log.TriggerContext(ctx, sw.EventProvider(contents))
    87  	count = len(contents)
    88  	return
    89  }