github.com/team-ide/go-dialect@v1.9.20/worker/data_source_csv.go (about)

     1  package worker
     2  
     3  import (
     4  	"bufio"
     5  	"errors"
     6  	"fmt"
     7  	"github.com/team-ide/go-dialect/dialect"
     8  	"io"
     9  	"os"
    10  	"strings"
    11  )
    12  
    13  func NewDataSourceCsv(param *DataSourceParam) (res DataSource) {
    14  	res = &dataSourceCsv{
    15  		DataSourceParam: param,
    16  	}
    17  	return
    18  }
    19  
    20  type dataSourceCsv struct {
    21  	*DataSourceParam
    22  	saveFile      *os.File
    23  	isStop        bool
    24  	headerWritten bool
    25  }
    26  
    27  func (this_ *dataSourceCsv) Stop() {
    28  	this_.isStop = true
    29  }
    30  
    31  func (this_ *dataSourceCsv) ReadStart() (err error) {
    32  	return
    33  }
    34  func (this_ *dataSourceCsv) ReadEnd() (err error) {
    35  	return
    36  }
    37  func (this_ *dataSourceCsv) Read(columnList []*dialect.ColumnModel, onRead func(data *DataSourceData) (err error)) (err error) {
    38  	defer func() {
    39  		if e := recover(); e != nil {
    40  			err = errors.New(fmt.Sprint(e))
    41  		}
    42  	}()
    43  	if this_.Path == "" {
    44  		err = errors.New("文件地址不能为空")
    45  		return
    46  	}
    47  	f, err := os.Open(this_.Path)
    48  	if err != nil {
    49  		return
    50  	}
    51  	buf := bufio.NewReader(f)
    52  	var line string
    53  	var rowInfo string
    54  	separator := this_.GetCsvSeparator()
    55  	columnLength := len(columnList)
    56  	if columnLength == 0 {
    57  		err = errors.New("column is null")
    58  		return
    59  	}
    60  	for {
    61  		if this_.isStop {
    62  			return
    63  		}
    64  		line, err = buf.ReadString('\n')
    65  		if line != "" {
    66  			if rowInfo == "" {
    67  				rowInfo = line
    68  			} else {
    69  				rowInfo += line
    70  			}
    71  
    72  			ss := strings.Split(strings.TrimSpace(rowInfo), separator)
    73  			if len(ss) > columnLength {
    74  				err = errors.New("row [" + rowInfo + "] can not to column names [" + strings.Join(GetColumnNames(columnList), ",") + "]")
    75  				return
    76  			}
    77  			if len(ss) == columnLength {
    78  				rowInfo = strings.TrimSpace(rowInfo)
    79  				if rowInfo != "" {
    80  					rowInfo = strings.ReplaceAll(rowInfo, this_.GetLinefeed(), "\n")
    81  					err = readRow(this_.Dia, rowInfo, separator, columnList, onRead)
    82  					if err != nil {
    83  						return
    84  					}
    85  				}
    86  				rowInfo = ""
    87  			}
    88  		}
    89  		if err != nil {
    90  			if err == io.EOF { //读取结束，会报EOF
    91  				err = nil
    92  			}
    93  			break
    94  		}
    95  	}
    96  	if err != nil {
    97  		return
    98  	}
    99  	rowInfo = strings.TrimSpace(rowInfo)
   100  	if rowInfo != "" {
   101  		err = readRow(this_.Dia, rowInfo, separator, columnList, onRead)
   102  		if err != nil {
   103  			return
   104  		}
   105  	}
   106  	return
   107  }
   108  
   109  func (this_ *dataSourceCsv) WriteStart() (err error) {
   110  
   111  	if this_.Path == "" {
   112  		err = errors.New("文件地址不能为空")
   113  		return
   114  	}
   115  
   116  	this_.saveFile, err = os.Create(this_.Path)
   117  	if err != nil {
   118  		return
   119  	}
   120  	return
   121  }
   122  func (this_ *dataSourceCsv) WriteEnd() (err error) {
   123  	if this_.saveFile != nil {
   124  		err = this_.saveFile.Close()
   125  		return
   126  	}
   127  	return
   128  }
   129  
   130  func (this_ *dataSourceCsv) WriteHeader(columnList []*dialect.ColumnModel) (err error) {
   131  	if this_.headerWritten {
   132  		return
   133  	}
   134  	this_.headerWritten = true
   135  
   136  	if this_.saveFile == nil {
   137  		err = this_.WriteStart()
   138  		if err != nil {
   139  			return
   140  		}
   141  	}
   142  
   143  	if this_.isStop {
   144  		return
   145  	}
   146  
   147  	var valueList []string
   148  	for _, column := range columnList {
   149  		str := column.ColumnName
   150  		str = strings.ReplaceAll(str, "\r\n", this_.GetLinefeed())
   151  		str = strings.ReplaceAll(str, "\n", this_.GetLinefeed())
   152  		str = strings.ReplaceAll(str, "\r", this_.GetLinefeed())
   153  		valueList = append(valueList, str)
   154  	}
   155  
   156  	_, err = this_.saveFile.WriteString(strings.Join(valueList, this_.GetCsvSeparator()) + "\n")
   157  	if err != nil {
   158  		return
   159  	}
   160  	return
   161  }
   162  func (this_ *dataSourceCsv) Write(data *DataSourceData) (err error) {
   163  	defer func() {
   164  		if e := recover(); e != nil {
   165  			err = errors.New(fmt.Sprint(e))
   166  		}
   167  	}()
   168  	if this_.saveFile == nil {
   169  		err = this_.WriteStart()
   170  		if err != nil {
   171  			return
   172  		}
   173  	}
   174  
   175  	if this_.isStop {
   176  		return
   177  	}
   178  	columnList := data.ColumnList
   179  	if data.Data == nil || columnList == nil {
   180  		return
   181  	}
   182  	var valueList []string
   183  	for _, column := range data.ColumnList {
   184  		str := dialect.GetStringValue(data.Data[column.ColumnName])
   185  		str = strings.ReplaceAll(str, "\r\n", this_.GetLinefeed())
   186  		str = strings.ReplaceAll(str, "\n", this_.GetLinefeed())
   187  		str = strings.ReplaceAll(str, "\r", this_.GetLinefeed())
   188  		valueList = append(valueList, str)
   189  	}
   190  
   191  	_, err = this_.saveFile.WriteString(strings.Join(valueList, this_.GetCsvSeparator()) + "\n")
   192  	if err != nil {
   193  		return
   194  	}
   195  
   196  	return
   197  }