github.com/andrewrech/ih-abstract@v0.0.0-20210322142951-2fec1c8d0f38/write.go (about)

     1  package main
     2  
     3  import (
     4  	"encoding/csv"
     5  	"log"
     6  	"os"
     7  	"strings"
     8  	"sync/atomic"
     9  )
    10  
    11  // Writer contains a file name, connection, CSV Writer, and a 'done' signal to cleanup the connection.
    12  type Writer struct {
    13  	name    string
    14  	conn    *os.File
    15  	w       *csv.Writer
    16  	counter int64
    17  	done    func()
    18  }
    19  
    20  // File creates an output CSV write file.
    21  func File(name string, h []string) (w Writer) {
    22  	f, err := os.OpenFile(name, os.O_TRUNC|os.O_CREATE|os.O_WRONLY, 0o644)
    23  	if err != nil {
    24  		log.Fatalln(err)
    25  	}
    26  
    27  	c := csv.NewWriter(f)
    28  
    29  	c.UseCRLF = false
    30  
    31  	err = c.Write(h)
    32  	if err != nil {
    33  		return Writer{}
    34  	}
    35  
    36  	w.name = name
    37  	w.conn = f
    38  	w.w = c
    39  
    40  	// done closes the connection
    41  	done := func() {
    42  		c.Flush()
    43  
    44  		err := c.Error()
    45  		if err != nil {
    46  			log.Fatalln(err)
    47  		}
    48  
    49  		f.Close()
    50  	}
    51  
    52  	w.done = done
    53  
    54  	return w
    55  }
    56  
    57  // WriteRows appends strings to a CSV file using a Writer.
    58  func WriteRows(in chan []string, name string, h []string, done chan struct{}) {
    59  	w := File(name, h)
    60  
    61  	go func() {
    62  		for l := range in {
    63  			err := w.w.Write(l)
    64  			if err != nil {
    65  				log.Fatalln(err)
    66  			}
    67  
    68  			atomic.AddInt64(&w.counter, 1)
    69  		}
    70  
    71  		w.done()
    72  		done <- struct{}{}
    73  	}()
    74  }
    75  
    76  // Write writes results to output CSV files using a common header.
    77  func Write(h []string, in map[string](chan []string)) (done chan struct{}) {
    78  	done = make(chan struct{})
    79  
    80  	nOutputFiles := len(in)
    81  	signal := make(chan struct{}, nOutputFiles)
    82  
    83  	for i, c := range in {
    84  
    85  		// filename
    86  		var fn strings.Builder
    87  		fn.WriteString(i)
    88  		fn.WriteString(".csv")
    89  
    90  		WriteRows(c, fn.String(), h, signal)
    91  	}
    92  
    93  	go func() {
    94  		for i := 0; i < nOutputFiles; i++ {
    95  			<-signal
    96  		}
    97  
    98  		close(done)
    99  	}()
   100  
   101  	return done
   102  }