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 }