github.com/yandex/pandora@v0.5.32/components/providers/scenario/vs/vs_csv.go (about)

     1  package vs
     2  
     3  import (
     4  	"encoding/csv"
     5  	"fmt"
     6  	"io"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/spf13/afero"
    11  )
    12  
    13  type VariableSourceCsv struct {
    14  	Name            string
    15  	File            string
    16  	Fields          []string
    17  	IgnoreFirstLine bool `config:"ignore_first_line"`
    18  	Delimiter       string
    19  	fs              afero.Fs
    20  	store           []map[string]string
    21  }
    22  
    23  func (v *VariableSourceCsv) GetName() string {
    24  	return v.Name
    25  }
    26  
    27  func (v *VariableSourceCsv) GetVariables() any {
    28  	return v.store
    29  }
    30  
    31  func (v *VariableSourceCsv) Init() (err error) {
    32  	const op = "VariableSourceCsv.Init"
    33  	var file afero.File
    34  	file, err = v.fs.Open(v.File)
    35  	if err != nil {
    36  		return fmt.Errorf("%s fs.Open %w", op, err)
    37  	}
    38  	defer func() {
    39  		closeErr := file.Close()
    40  		if closeErr != nil {
    41  			if err != nil {
    42  				err = fmt.Errorf("%s multiple errors faced: %w, with close err: %s", op, err, closeErr)
    43  			} else {
    44  				err = fmt.Errorf("%s, %w", op, closeErr)
    45  			}
    46  		}
    47  	}()
    48  
    49  	store, err := readCsv(file, v.IgnoreFirstLine, v.Delimiter, v.Fields)
    50  	if err != nil {
    51  		return fmt.Errorf("%s readCsv %w", op, err)
    52  	}
    53  	v.store = store
    54  
    55  	return nil
    56  }
    57  
    58  func readCsv(file afero.File, ignoreFirstLine bool, delimiter string, fields []string) ([]map[string]string, error) {
    59  	const op = "readCsv"
    60  	reader := csv.NewReader(file)
    61  	if delimiter != "" {
    62  		reader.Comma = rune(delimiter[0])
    63  	}
    64  	for i := range fields {
    65  		fields[i] = strings.Replace(fields[i], " ", "_", -1)
    66  	}
    67  	result := make([]map[string]string, 0)
    68  	for {
    69  		record, err := reader.Read()
    70  		if err == io.EOF {
    71  			break
    72  		}
    73  		if err != nil {
    74  			return nil, fmt.Errorf("%s csv.Read %w", op, err)
    75  		}
    76  		if len(fields) == 0 {
    77  			fields = make([]string, len(record))
    78  			for i := range record {
    79  				fields[i] = strings.Replace(record[i], " ", "_", -1)
    80  			}
    81  		}
    82  		if ignoreFirstLine {
    83  			ignoreFirstLine = false
    84  			continue
    85  		}
    86  		row := make(map[string]string)
    87  		for i, field := range fields {
    88  			if field == "" {
    89  				field = strconv.Itoa(i)
    90  			}
    91  			if i >= len(record) {
    92  				row[field] = ""
    93  			} else {
    94  				row[field] = record[i]
    95  			}
    96  		}
    97  		result = append(result, row)
    98  	}
    99  	return result, nil
   100  }
   101  
   102  func NewVSCSV(cfg VariableSourceCsv, fs afero.Fs) (VariableSource, error) {
   103  	cfg.fs = fs
   104  	return &cfg, nil
   105  }
   106  
   107  var _ VariableSource = (*VariableSourceCsv)(nil)