github.com/yandex/pandora@v0.5.32/core/provider/json.go (about)

     1  package provider
     2  
     3  import (
     4  	"io"
     5  
     6  	jsoniter "github.com/json-iterator/go"
     7  	"github.com/yandex/pandora/core"
     8  	"github.com/yandex/pandora/core/coreutil"
     9  	"github.com/yandex/pandora/lib/ioutil2"
    10  )
    11  
    12  // NewJSONProvider returns generic core.Provider that reads JSON data from source and decodes it
    13  // into ammo returned from newAmmo.
    14  func NewJSONProvider(newAmmo func() core.Ammo, conf JSONProviderConfig) core.Provider {
    15  	return NewCustomJSONProvider(nil, newAmmo, conf)
    16  }
    17  
    18  // NewCustomJSONProvider is like NewJSONProvider, but also allows to wrap JSON decoder, to
    19  // decode data into intermediate struct, but then transform in into desired ammo.
    20  // For example, decode {"body":"some data"} into struct { Data string }, and transform it to
    21  // http.Request.
    22  func NewCustomJSONProvider(wrapDecoder func(deps core.ProviderDeps, decoder AmmoDecoder) AmmoDecoder, newAmmo func() core.Ammo, conf JSONProviderConfig) core.Provider {
    23  	var newDecoder NewAmmoDecoder = func(deps core.ProviderDeps, source io.Reader) (AmmoDecoder, error) {
    24  		decoder := NewJSONAmmoDecoder(source, conf.Buffer.BufferSizeOrDefault())
    25  		if wrapDecoder != nil {
    26  			decoder = wrapDecoder(deps, decoder)
    27  		}
    28  		return decoder, nil
    29  	}
    30  	return NewDecodeProvider(newAmmo, newDecoder, conf.Decode)
    31  }
    32  
    33  type JSONProviderConfig struct {
    34  	Decode DecodeProviderConfig      `config:",squash"`
    35  	Buffer coreutil.BufferSizeConfig `config:",squash"`
    36  }
    37  
    38  func DefaultJSONProviderConfig() JSONProviderConfig {
    39  	return JSONProviderConfig{Decode: DefaultDecodeProviderConfig()}
    40  }
    41  
    42  func NewJSONAmmoDecoder(r io.Reader, buffSize int) AmmoDecoder {
    43  	var readError error
    44  	// HACK(skipor): jsoniter.Iterator don't handle read errors well, but jsoniter.Decoder don't allow to set buffer size.
    45  	var errTrackingReader ioutil2.ReaderFunc = func(p []byte) (n int, err error) {
    46  		n, err = r.Read(p)
    47  		if n > 0 {
    48  			// Need to suppress error, to distinguish parse error in last chunk and read error.
    49  			return n, nil
    50  		}
    51  		if err != nil {
    52  			readError = err
    53  		}
    54  		return n, err
    55  	}
    56  	return &JSONAmmoDecoder{
    57  		iter:         jsoniter.Parse(jsoniter.ConfigFastest, errTrackingReader, buffSize),
    58  		readErrorPtr: &readError,
    59  	}
    60  }
    61  
    62  type JSONAmmoDecoder struct {
    63  	iter         *jsoniter.Iterator
    64  	readErrorPtr *error
    65  }
    66  
    67  func (d *JSONAmmoDecoder) Decode(ammo core.Ammo) error {
    68  	coreutil.ResetReusedAmmo(ammo)
    69  	d.iter.ReadVal(ammo)
    70  	if d.iter.Error != nil {
    71  		if *d.readErrorPtr != nil {
    72  			return *d.readErrorPtr
    73  		}
    74  		return d.iter.Error
    75  	}
    76  	return nil
    77  }