github.com/opensearch-project/opensearch-go/v2@v2.3.0/opensearchutil/json_reader.go (about)

     1  // SPDX-License-Identifier: Apache-2.0
     2  //
     3  // The OpenSearch Contributors require contributions made to
     4  // this file be licensed under the Apache-2.0 license or a
     5  // compatible open source license.
     6  //
     7  // Modifications Copyright OpenSearch Contributors. See
     8  // GitHub history for details.
     9  
    10  // Licensed to Elasticsearch B.V. under one or more contributor
    11  // license agreements. See the NOTICE file distributed with
    12  // this work for additional information regarding copyright
    13  // ownership. Elasticsearch B.V. licenses this file to you under
    14  // the Apache License, Version 2.0 (the "License"); you may
    15  // not use this file except in compliance with the License.
    16  // You may obtain a copy of the License at
    17  //
    18  //    http://www.apache.org/licenses/LICENSE-2.0
    19  //
    20  // Unless required by applicable law or agreed to in writing,
    21  // software distributed under the License is distributed on an
    22  // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
    23  // KIND, either express or implied.  See the License for the
    24  // specific language governing permissions and limitations
    25  // under the License.
    26  
    27  package opensearchutil
    28  
    29  import (
    30  	"bytes"
    31  	"encoding/json"
    32  	"io"
    33  )
    34  
    35  // NewJSONReader encodes v into JSON and returns it as an io.Reader.
    36  //
    37  func NewJSONReader(v interface{}) io.Reader {
    38  	return &JSONReader{val: v, buf: nil}
    39  }
    40  
    41  // JSONEncoder defines the interface for custom JSON encoders.
    42  //
    43  type JSONEncoder interface {
    44  	EncodeJSON(io.Writer) error
    45  }
    46  
    47  // JSONReader represents a reader which takes an interface value,
    48  // encodes it into JSON, and wraps it in an io.Reader.
    49  //
    50  type JSONReader struct {
    51  	val interface{}
    52  	buf interface {
    53  		io.ReadWriter
    54  		io.WriterTo
    55  	}
    56  }
    57  
    58  // Read implements the io.Reader interface.
    59  //
    60  func (r *JSONReader) Read(p []byte) (int, error) {
    61  	if r.buf == nil {
    62  		r.buf = new(bytes.Buffer)
    63  		if err := r.encode(r.buf); err != nil {
    64  			return 0, err
    65  		}
    66  	}
    67  
    68  	return r.buf.Read(p)
    69  }
    70  
    71  // WriteTo implements the io.WriterTo interface.
    72  //
    73  func (r *JSONReader) WriteTo(w io.Writer) (int64, error) {
    74  	cw := countingWriter{Writer: w}
    75  	err := r.encode(&cw)
    76  	return int64(cw.n), err
    77  }
    78  
    79  func (r *JSONReader) encode(w io.Writer) error {
    80  	var err error
    81  
    82  	if e, ok := r.val.(JSONEncoder); ok {
    83  		err = e.EncodeJSON(w)
    84  		if err != nil {
    85  			return err
    86  		}
    87  		return nil
    88  	}
    89  
    90  	return json.NewEncoder(w).Encode(r.val)
    91  }
    92  
    93  type countingWriter struct {
    94  	io.Writer
    95  	n int
    96  }
    97  
    98  func (cw *countingWriter) Write(p []byte) (int, error) {
    99  	n, err := cw.Writer.Write(p)
   100  	if n > 0 {
   101  		cw.n += n
   102  	}
   103  	return n, err
   104  }