github.com/mondo192/jfrog-client-go@v1.0.0/utils/io/content/contentwriter.go (about)

     1  package content
     2  
     3  import (
     4  	"bytes"
     5  	"encoding/json"
     6  	"fmt"
     7  	"os"
     8  	"sync"
     9  
    10  	"github.com/mondo192/jfrog-client-go/utils"
    11  	"github.com/mondo192/jfrog-client-go/utils/errorutils"
    12  	"github.com/mondo192/jfrog-client-go/utils/log"
    13  
    14  	"github.com/mondo192/jfrog-client-go/utils/io/fileutils"
    15  )
    16  
    17  const (
    18  	jsonArrayPrefixPattern = "  \"%s\": ["
    19  	jsonArraySuffix        = "]\n"
    20  	DefaultKey             = "results"
    21  )
    22  
    23  // Write a JSON file in small chunks. Only a single JSON key can be written to the file, and array as its value.
    24  // The array's values could be any JSON value types (number, string, etc...).
    25  // Once the first 'Write" call is made, the file will stay open, waiting for the next struct to be written (thread-safe).
    26  // Finally, 'Close' will fill the end of the JSON file and the operation will be completed.
    27  type ContentWriter struct {
    28  	// arrayKey = JSON object key to be written.
    29  	arrayKey string
    30  	// The output data file path.
    31  	outputFile *os.File
    32  	// The chanel which from the output records will be pulled.
    33  	dataChannel    chan interface{}
    34  	isCompleteFile bool
    35  	errorsQueue    *utils.ErrorsQueue
    36  	runWaiter      sync.WaitGroup
    37  	once           sync.Once
    38  	empty          bool
    39  	useStdout      bool
    40  }
    41  
    42  func NewContentWriter(arrayKey string, isCompleteFile, useStdout bool) (*ContentWriter, error) {
    43  	self := ContentWriter{}
    44  	self.useStdout = useStdout
    45  	self.arrayKey = arrayKey
    46  	self.dataChannel = make(chan interface{}, utils.MaxBufferSize)
    47  	self.errorsQueue = utils.NewErrorsQueue(utils.MaxBufferSize)
    48  	self.isCompleteFile = isCompleteFile
    49  	self.empty = true
    50  	return &self, nil
    51  }
    52  
    53  func (rw *ContentWriter) SetArrayKey(arrKey string) *ContentWriter {
    54  	rw.arrayKey = arrKey
    55  	return rw
    56  }
    57  
    58  func (rw *ContentWriter) GetArrayKey() string {
    59  	return rw.arrayKey
    60  }
    61  
    62  func (rw *ContentWriter) IsEmpty() bool {
    63  	return rw.empty
    64  }
    65  
    66  func (rw *ContentWriter) GetFilePath() string {
    67  	if rw.outputFile != nil {
    68  		return rw.outputFile.Name()
    69  	}
    70  	return ""
    71  }
    72  
    73  func (rw *ContentWriter) RemoveOutputFilePath() error {
    74  	return errorutils.CheckError(os.Remove(rw.outputFile.Name()))
    75  }
    76  
    77  // Write a single item to the JSON array.
    78  func (rw *ContentWriter) Write(record interface{}) {
    79  	rw.empty = false
    80  	rw.startWritingWorker()
    81  	rw.dataChannel <- record
    82  }
    83  
    84  func (rw *ContentWriter) startWritingWorker() {
    85  	rw.once.Do(func() {
    86  		var err error
    87  		if rw.useStdout {
    88  			rw.outputFile = os.Stdout
    89  		} else {
    90  			rw.outputFile, err = fileutils.CreateTempFile()
    91  			if err != nil {
    92  				rw.errorsQueue.AddError(errorutils.CheckError(err))
    93  				return
    94  			}
    95  		}
    96  		rw.runWaiter.Add(1)
    97  		go func() {
    98  			defer rw.runWaiter.Done()
    99  			rw.run()
   100  		}()
   101  	})
   102  }
   103  
   104  // Write the data from the channel to JSON file.
   105  // The channel may block the thread, therefore should run async.
   106  func (rw *ContentWriter) run() {
   107  	var err error
   108  	if !rw.useStdout {
   109  		defer func() {
   110  			err = rw.outputFile.Close()
   111  			if err != nil {
   112  				rw.errorsQueue.AddError(errorutils.CheckError(err))
   113  			}
   114  		}()
   115  	}
   116  	openString := jsonArrayPrefixPattern
   117  	closeString := ""
   118  	if rw.isCompleteFile {
   119  		openString = "{\n" + openString
   120  	}
   121  	_, err = rw.outputFile.WriteString(fmt.Sprintf(openString, rw.arrayKey))
   122  	if err != nil {
   123  		rw.errorsQueue.AddError(errorutils.CheckError(err))
   124  		return
   125  	}
   126  	buf := bytes.NewBuffer(nil)
   127  	enc := json.NewEncoder(buf)
   128  	enc.SetIndent("    ", "  ")
   129  	recordPrefix := "\n    "
   130  	firstRecord := true
   131  	for record := range rw.dataChannel {
   132  		buf.Reset()
   133  		err = enc.Encode(record)
   134  		if err != nil {
   135  			rw.errorsQueue.AddError(errorutils.CheckError(err))
   136  			continue
   137  		}
   138  		record := recordPrefix + string(bytes.TrimRight(buf.Bytes(), "\n"))
   139  		_, err = rw.outputFile.WriteString(record)
   140  		if err != nil {
   141  			rw.errorsQueue.AddError(errorutils.CheckError(err))
   142  			continue
   143  		}
   144  		if firstRecord {
   145  			// If a record was printed, we want to print a comma and ne before each and every future record.
   146  			recordPrefix = "," + recordPrefix
   147  			// We will close the array in a new-indent line.
   148  			closeString = "\n  "
   149  			firstRecord = false
   150  		}
   151  	}
   152  	closeString = closeString + jsonArraySuffix
   153  	if rw.isCompleteFile {
   154  		closeString += "}\n"
   155  	}
   156  	_, err = rw.outputFile.WriteString(closeString)
   157  	if err != nil {
   158  		rw.errorsQueue.AddError(errorutils.CheckError(err))
   159  	}
   160  }
   161  
   162  // Finish writing to the file.
   163  func (rw *ContentWriter) Close() error {
   164  	if rw.empty {
   165  		return nil
   166  	}
   167  	close(rw.dataChannel)
   168  	rw.runWaiter.Wait()
   169  	if err := rw.GetError(); err != nil {
   170  		log.Error("Failed to close writer: " + err.Error())
   171  		return err
   172  	}
   173  	return nil
   174  }
   175  
   176  func (rw *ContentWriter) GetError() error {
   177  	return rw.errorsQueue.GetError()
   178  }