github.com/team-ide/go-dialect@v1.9.20/worker/data_source_sql.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 NewDataSourceSql(param *DataSourceParam) (res DataSource) {
    14  	res = &dataSourceSql{
    15  		DataSourceParam: param,
    16  	}
    17  	return
    18  }
    19  
    20  type dataSourceSql struct {
    21  	*DataSourceParam
    22  	saveFile *os.File
    23  	isStop   bool
    24  }
    25  
    26  func (this_ *dataSourceSql) Stop() {
    27  	this_.isStop = true
    28  }
    29  
    30  func (this_ *dataSourceSql) ReadStart() (err error) {
    31  	return
    32  }
    33  func (this_ *dataSourceSql) ReadEnd() (err error) {
    34  	return
    35  }
    36  func (this_ *dataSourceSql) Read(columnList []*dialect.ColumnModel, onRead func(data *DataSourceData) (err error)) (err error) {
    37  	defer func() {
    38  		if e := recover(); e != nil {
    39  			err = errors.New(fmt.Sprint(e))
    40  		}
    41  	}()
    42  	if this_.Path == "" {
    43  		err = errors.New("文件地址不能为空")
    44  		return
    45  	}
    46  	f, err := os.Open(this_.Path)
    47  	if err != nil {
    48  		return
    49  	}
    50  	buf := bufio.NewReader(f)
    51  	var line string
    52  	var sqlInfo string
    53  	for {
    54  		if this_.isStop {
    55  			return
    56  		}
    57  		line, err = buf.ReadString('\n')
    58  		if line != "" {
    59  			sqlInfo += line
    60  			if this_.Dia.IsSqlEnd(sqlInfo) {
    61  				sqlInfo = strings.TrimSpace(sqlInfo)
    62  				if sqlInfo != "" {
    63  					err = onRead(&DataSourceData{
    64  						HasSql: true,
    65  						Sql:    sqlInfo,
    66  					})
    67  					if err != nil {
    68  						return
    69  					}
    70  				}
    71  				sqlInfo = ""
    72  			}
    73  		}
    74  		if err != nil {
    75  			if err == io.EOF { //读取结束，会报EOF
    76  				err = nil
    77  			}
    78  			break
    79  		}
    80  	}
    81  	if err != nil {
    82  		return
    83  	}
    84  	sqlInfo = strings.TrimSpace(sqlInfo)
    85  	if sqlInfo != "" {
    86  		err = onRead(&DataSourceData{
    87  			HasSql: true,
    88  			Sql:    sqlInfo,
    89  		})
    90  		if err != nil {
    91  			return
    92  		}
    93  	}
    94  	return
    95  }
    96  
    97  func (this_ *dataSourceSql) WriteStart() (err error) {
    98  
    99  	if this_.Path == "" {
   100  		err = errors.New("文件地址不能为空")
   101  		return
   102  	}
   103  
   104  	this_.saveFile, err = os.Create(this_.Path)
   105  	if err != nil {
   106  		return
   107  	}
   108  	return
   109  }
   110  func (this_ *dataSourceSql) WriteEnd() (err error) {
   111  	if this_.saveFile != nil {
   112  		err = this_.saveFile.Close()
   113  		return
   114  	}
   115  	return
   116  }
   117  
   118  func (this_ *dataSourceSql) WriteHeader(columnList []*dialect.ColumnModel) (err error) {
   119  
   120  	return
   121  }
   122  func (this_ *dataSourceSql) Write(data *DataSourceData) (err error) {
   123  	defer func() {
   124  		if e := recover(); e != nil {
   125  			err = errors.New(fmt.Sprint(e))
   126  		}
   127  	}()
   128  	if this_.saveFile == nil {
   129  		err = this_.WriteStart()
   130  		if err != nil {
   131  			return
   132  		}
   133  	}
   134  
   135  	if this_.isStop {
   136  		return
   137  	}
   138  	if data.HasSql {
   139  		_, err = this_.saveFile.WriteString(data.Sql + ";\n")
   140  		if err != nil {
   141  			return
   142  		}
   143  	}
   144  	return
   145  }