github.com/minio/minio@v0.0.0-20240328213742-3f72439b8a27/internal/s3select/json/reader.go (about)

     1  // Copyright (c) 2015-2021 MinIO, Inc.
     2  //
     3  // This file is part of MinIO Object Storage stack
     4  //
     5  // This program is free software: you can redistribute it and/or modify
     6  // it under the terms of the GNU Affero General Public License as published by
     7  // the Free Software Foundation, either version 3 of the License, or
     8  // (at your option) any later version.
     9  //
    10  // This program is distributed in the hope that it will be useful
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13  // GNU Affero General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Affero General Public License
    16  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  package json
    19  
    20  import (
    21  	"io"
    22  	"sync"
    23  
    24  	"github.com/minio/minio/internal/s3select/sql"
    25  
    26  	"github.com/bcicen/jstream"
    27  )
    28  
    29  // Reader - JSON record reader for S3Select.
    30  type Reader struct {
    31  	args       *ReaderArgs
    32  	decoder    *jstream.Decoder
    33  	valueCh    chan *jstream.MetaValue
    34  	readCloser io.ReadCloser
    35  }
    36  
    37  // Read - reads single record.
    38  func (r *Reader) Read(dst sql.Record) (sql.Record, error) {
    39  	v, ok := <-r.valueCh
    40  	if !ok {
    41  		if err := r.decoder.Err(); err != nil {
    42  			return nil, errJSONParsingError(err)
    43  		}
    44  		return nil, io.EOF
    45  	}
    46  
    47  	var kvs jstream.KVS
    48  	if v.ValueType == jstream.Object {
    49  		// This is a JSON object type (that preserves key
    50  		// order)
    51  		kvs = v.Value.(jstream.KVS)
    52  	} else {
    53  		// To be AWS S3 compatible Select for JSON needs to
    54  		// output non-object JSON as single column value
    55  		// i.e. a map with `_1` as key and value as the
    56  		// non-object.
    57  		kvs = jstream.KVS{jstream.KV{Key: "_1", Value: v.Value}}
    58  	}
    59  
    60  	dstRec, ok := dst.(*Record)
    61  	if !ok {
    62  		dstRec = &Record{}
    63  	}
    64  	dstRec.KVS = kvs
    65  	dstRec.SelectFormat = sql.SelectFmtJSON
    66  	return dstRec, nil
    67  }
    68  
    69  // Close - closes underlying reader.
    70  func (r *Reader) Close() error {
    71  	// Close the input.
    72  	err := r.readCloser.Close()
    73  	for range r.valueCh {
    74  		// Drain values so we don't leak a goroutine.
    75  		// Since we have closed the input, it should fail rather quickly.
    76  	}
    77  	return err
    78  }
    79  
    80  // NewReader - creates new JSON reader using readCloser.
    81  func NewReader(readCloser io.ReadCloser, args *ReaderArgs) *Reader {
    82  	readCloser = &syncReadCloser{rc: readCloser}
    83  	d := jstream.NewDecoder(readCloser, 0).ObjectAsKVS()
    84  	return &Reader{
    85  		args:       args,
    86  		decoder:    d,
    87  		valueCh:    d.Stream(),
    88  		readCloser: readCloser,
    89  	}
    90  }
    91  
    92  // syncReadCloser will wrap a readcloser and make it safe to call Close while
    93  // reads are running.
    94  type syncReadCloser struct {
    95  	rc io.ReadCloser
    96  	mu sync.Mutex
    97  }
    98  
    99  func (pr *syncReadCloser) Read(p []byte) (n int, err error) {
   100  	// This ensures that Close will block until Read has completed.
   101  	// This allows another goroutine to close the reader.
   102  	pr.mu.Lock()
   103  	defer pr.mu.Unlock()
   104  	if pr.rc == nil {
   105  		return 0, io.EOF
   106  	}
   107  	return pr.rc.Read(p)
   108  }
   109  
   110  func (pr *syncReadCloser) Close() error {
   111  	pr.mu.Lock()
   112  	defer pr.mu.Unlock()
   113  	if pr.rc != nil {
   114  		err := pr.rc.Close()
   115  		pr.rc = nil
   116  		return err
   117  	}
   118  	return nil
   119  }