github.com/cobalt77/jfrog-client-go@v0.14.5/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/cobalt77/jfrog-client-go/utils"
    11  	"github.com/cobalt77/jfrog-client-go/utils/errorutils"
    12  	"github.com/cobalt77/jfrog-client-go/utils/log"
    13  
    14  	"github.com/cobalt77/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 rw.outputFile.Close()
   110  	}
   111  	openString := jsonArrayPrefixPattern
   112  	closeString := ""
   113  	if rw.isCompleteFile {
   114  		openString = "{\n" + openString
   115  	}
   116  	_, err = rw.outputFile.WriteString(fmt.Sprintf(openString, rw.arrayKey))
   117  	if err != nil {
   118  		rw.errorsQueue.AddError(errorutils.CheckError(err))
   119  		return
   120  	}
   121  	buf := bytes.NewBuffer(nil)
   122  	enc := json.NewEncoder(buf)
   123  	enc.SetIndent("    ", "  ")
   124  	recordPrefix := "\n    "
   125  	firstRecord := true
   126  	for record := range rw.dataChannel {
   127  		buf.Reset()
   128  		err = enc.Encode(record)
   129  		if err != nil {
   130  			rw.errorsQueue.AddError(errorutils.CheckError(err))
   131  			continue
   132  		}
   133  		record := recordPrefix + string(bytes.TrimRight(buf.Bytes(), "\n"))
   134  		_, err = rw.outputFile.WriteString(record)
   135  		if err != nil {
   136  			rw.errorsQueue.AddError(errorutils.CheckError(err))
   137  			continue
   138  		}
   139  		if firstRecord {
   140  			// If a record was printed, we want to print a comma and ne before each and every future record.
   141  			recordPrefix = "," + recordPrefix
   142  			// We will close the array in a new-indent line.
   143  			closeString = "\n  "
   144  			firstRecord = false
   145  		}
   146  	}
   147  	closeString = closeString + jsonArraySuffix
   148  	if rw.isCompleteFile {
   149  		closeString += "}\n"
   150  	}
   151  	_, err = rw.outputFile.WriteString(closeString)
   152  	if err != nil {
   153  		rw.errorsQueue.AddError(errorutils.CheckError(err))
   154  	}
   155  	return
   156  }
   157  
   158  // Finish writing to the file.
   159  func (rw *ContentWriter) Close() error {
   160  	if rw.empty {
   161  		return nil
   162  	}
   163  	close(rw.dataChannel)
   164  	rw.runWaiter.Wait()
   165  	if err := rw.GetError(); err != nil {
   166  		log.Error("Failed to close writer: " + err.Error())
   167  		return err
   168  	}
   169  	return nil
   170  }
   171  
   172  func (rw *ContentWriter) GetError() error {
   173  	return rw.errorsQueue.GetError()
   174  }