github.com/blend/go-sdk@v1.20220411.3/logger/json_output_formatter.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  	"encoding/json"
    13  	"io"
    14  
    15  	"github.com/blend/go-sdk/bufferutil"
    16  )
    17  
    18  var (
    19  	_ WriteFormatter = (*JSONOutputFormatter)(nil)
    20  )
    21  
    22  // NewJSONOutputFormatter returns a new json event formatter.
    23  func NewJSONOutputFormatter(options ...JSONOutputFormatterOption) *JSONOutputFormatter {
    24  	jf := &JSONOutputFormatter{
    25  		BufferPool: bufferutil.NewPool(DefaultBufferPoolSize),
    26  	}
    27  
    28  	for _, option := range options {
    29  		option(jf)
    30  	}
    31  	return jf
    32  }
    33  
    34  // JSONOutputFormatterOption is an option for json formatters.
    35  type JSONOutputFormatterOption func(*JSONOutputFormatter)
    36  
    37  // OptJSONConfig sets a json formatter from a config.
    38  func OptJSONConfig(cfg JSONConfig) JSONOutputFormatterOption {
    39  	return func(jf *JSONOutputFormatter) {
    40  		jf.Pretty = cfg.Pretty
    41  		jf.PrettyIndent = cfg.PrettyIndentOrDefault()
    42  		jf.PrettyPrefix = cfg.PrettyPrefixOrDefault()
    43  	}
    44  }
    45  
    46  // OptJSONPretty sets the json output formatter to indent output.
    47  func OptJSONPretty() JSONOutputFormatterOption {
    48  	return func(jso *JSONOutputFormatter) { jso.Pretty = true }
    49  }
    50  
    51  // OptJSONPrettyPrefix sets the json output formatter to indent output.
    52  func OptJSONPrettyPrefix(prettyPrefix string) JSONOutputFormatterOption {
    53  	return func(jso *JSONOutputFormatter) { jso.PrettyPrefix = prettyPrefix }
    54  }
    55  
    56  // OptJSONPrettyIndent sets the json output formatter to indent output.
    57  func OptJSONPrettyIndent(prettyIndent string) JSONOutputFormatterOption {
    58  	return func(jso *JSONOutputFormatter) { jso.PrettyIndent = prettyIndent }
    59  }
    60  
    61  // JSONOutputFormatter is a json output formatter.
    62  type JSONOutputFormatter struct {
    63  	BufferPool   *bufferutil.Pool
    64  	Pretty       bool
    65  	PrettyPrefix string
    66  	PrettyIndent string
    67  }
    68  
    69  // PrettyPrefixOrDefault returns the pretty prefix or a default.
    70  func (jw JSONOutputFormatter) PrettyPrefixOrDefault() string {
    71  	if jw.PrettyPrefix != "" {
    72  		return jw.PrettyPrefix
    73  	}
    74  	return ""
    75  }
    76  
    77  // PrettyIndentOrDefault returns the pretty indent or a default.
    78  func (jw JSONOutputFormatter) PrettyIndentOrDefault() string {
    79  	if jw.PrettyIndent != "" {
    80  		return jw.PrettyIndent
    81  	}
    82  	return "\t"
    83  }
    84  
    85  // WriteFormat writes the event to the given output.
    86  func (jw JSONOutputFormatter) WriteFormat(ctx context.Context, output io.Writer, e Event) error {
    87  	buffer := jw.BufferPool.Get()
    88  	defer jw.BufferPool.Put(buffer)
    89  
    90  	encoder := json.NewEncoder(buffer)
    91  	if jw.Pretty {
    92  		encoder.SetIndent(jw.PrettyPrefixOrDefault(), jw.PrettyIndentOrDefault())
    93  	}
    94  	if decomposer, ok := e.(JSONWritable); ok {
    95  		fields := jw.CombineFields(jw.GetScopeFields(ctx, e), decomposer.Decompose())
    96  		if err := encoder.Encode(fields); err != nil {
    97  			return err
    98  		}
    99  	} else if err := encoder.Encode(e); err != nil {
   100  		return err
   101  	}
   102  	_, err := io.Copy(output, buffer)
   103  	return err
   104  }
   105  
   106  // CombineFields combines a variadic set of fields.
   107  func (jw JSONOutputFormatter) CombineFields(fields ...map[string]interface{}) map[string]interface{} {
   108  	if len(fields) == 0 {
   109  		return nil
   110  	}
   111  	output := make(map[string]interface{})
   112  	for _, set := range fields {
   113  		if set == nil {
   114  			continue
   115  		}
   116  		for key, value := range set {
   117  			output[key] = value
   118  		}
   119  	}
   120  	return output
   121  }
   122  
   123  // GetScopeFields gets scope fields from a context.
   124  func (jw JSONOutputFormatter) GetScopeFields(ctx context.Context, e Event) map[string]interface{} {
   125  	output := map[string]interface{}{
   126  		FieldFlag:      e.GetFlag(),
   127  		FieldTimestamp: GetEventTimestamp(ctx, e),
   128  	}
   129  	if path := GetPath(ctx); len(path) > 0 {
   130  		output[FieldScopePath] = path
   131  	}
   132  	if labels := GetLabels(ctx); len(labels) > 0 {
   133  		output[FieldLabels] = labels
   134  	}
   135  	if annotations := GetAnnotations(ctx); len(annotations) > 0 {
   136  		output[FieldAnnotations] = annotations
   137  	}
   138  	return output
   139  }