github.com/blong14/gache@v0.0.0-20240124023949-89416fd8bbfa/internal/proxy/reader.go (about)

     1  package proxy
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	gdb "github.com/blong14/gache/internal/db"
     8  	gfile "github.com/blong14/gache/internal/io/file"
     9  )
    10  
    11  type waiter struct {
    12  	sync.WaitGroup
    13  	chns []chan gdb.QueryResponse
    14  }
    15  
    16  func (w *waiter) Add(ch chan gdb.QueryResponse) {
    17  	w.chns = append(w.chns, ch)
    18  	w.WaitGroup.Add(1)
    19  }
    20  
    21  func (w *waiter) Wait(ctx context.Context) {
    22  	for _, ch := range w.chns {
    23  		go func(done chan gdb.QueryResponse) {
    24  			defer w.WaitGroup.Done()
    25  			select {
    26  			case <-ctx.Done():
    27  			case <-done:
    28  			}
    29  		}(ch)
    30  	}
    31  	w.WaitGroup.Wait()
    32  }
    33  
    34  type CSVReader struct {
    35  	worker Actor
    36  	waiter *waiter
    37  }
    38  
    39  func NewCSVReader(worker Actor) *CSVReader {
    40  	return &CSVReader{
    41  		worker: worker,
    42  		waiter: &waiter{chns: make([]chan gdb.QueryResponse, 0)},
    43  	}
    44  }
    45  
    46  func (f *CSVReader) Read(ctx context.Context, query *gdb.Query) {
    47  	if query.Header.Inst != gdb.Load {
    48  		if query != nil {
    49  			query.Done(gdb.QueryResponse{Success: false})
    50  		}
    51  		return
    52  	}
    53  	reader := gfile.ScanCSV(string(query.Header.FileName))
    54  	defer reader.Close()
    55  	reader.Init()
    56  	for reader.Scan() {
    57  		var rows []gdb.KeyValue
    58  		for _, r := range reader.Rows() {
    59  			rows = append(rows, gdb.KeyValue{
    60  				Key:   []byte(r[0]),
    61  				Value: []byte(r[1]),
    62  			})
    63  		}
    64  		q, done := gdb.NewBatchSetValueQuery(ctx, query.Header.TableName, rows)
    65  		f.waiter.Add(done)
    66  		f.worker.Send(ctx, q)
    67  	}
    68  	f.waiter.Wait(ctx)
    69  	success := false
    70  	if err := reader.Err(); err == nil {
    71  		success = true
    72  	}
    73  	query.Done(
    74  		gdb.QueryResponse{
    75  			Success: success,
    76  			Value:   []byte("done"),
    77  		},
    78  	)
    79  }