From c2b1c0bca9b2e09d5c113a8d4d1bbfc9918c8f2a Mon Sep 17 00:00:00 2001 From: "Harris.Chu" <1726587+HarrisChu@users.noreply.github.com> Date: Mon, 30 Mar 2026 16:17:59 +0800 Subject: [PATCH 1/2] support get multiple file data --- pkg/common/csv.go | 142 +++++++++++++++++++++++++------------ pkg/common/types.go | 10 ++- pkg/nebulagraph/client.go | 51 ++++++------- pkg/nebulagraph5/client.go | 37 +++++----- 4 files changed, 143 insertions(+), 97 deletions(-) diff --git a/pkg/common/csv.go b/pkg/common/csv.go index ae5b4e9..9da2a32 100644 --- a/pkg/common/csv.go +++ b/pkg/common/csv.go @@ -2,92 +2,140 @@ package common import ( "encoding/csv" + "fmt" "io" "os" + "sync" + "sync/atomic" ) type ( - CSVReader struct { - Path string - Delimiter string - WithHeader bool - limit int + fileReader struct { + data []Data + index atomic.Int64 } - CSVWriter struct { - Path string - Header []string - Delimiter string - DataCh <-chan []string + csvReaderImpl struct { + defaultSource string + config CsvReaderConfig + readers map[string]*fileReader + mutex sync.Mutex } ) -func NewCsvReader(path, delimiter string, withHeader bool, limit int) *CSVReader { - return &CSVReader{ - Path: path, - Delimiter: delimiter, - WithHeader: withHeader, - limit: limit, +func NewCsvReader(defaultSource string, config CsvReaderConfig) ICsvReader { + return &csvReaderImpl{ + defaultSource: defaultSource, + config: config, + readers: make(map[string]*fileReader), } } -func NewCsvWriter(path, delimiter string, header []string, dataCh <-chan []string) *CSVWriter { - return &CSVWriter{ - Path: path, - Delimiter: delimiter, - Header: header, - DataCh: dataCh, +func (r *csvReaderImpl) GetData(sourceFile string) (Data, error) { + key := sourceFile + if key == "" { + key = r.defaultSource + } + + reader := r.getOrCreateReader(key) + if reader == nil { + return nil, fmt.Errorf("failed to load csv file: %s", key) + } + + idx := int(reader.index.Add(1) - 1) + data := reader.data + if len(data) == 0 { + return nil, fmt.Errorf("no data in csv file: %s", key) } + if idx >= len(data) { + idx = idx % len(data) + } + return data[idx], nil } -// ReadForever read the csv in slice first, and send to the data channel forever. -func (c *CSVReader) ReadForever(dataCh chan<- Data) error { - lines := make([]Data, 0, c.limit) - file, err := os.Open(c.Path) +func (r *csvReaderImpl) getOrCreateReader(key string) *fileReader { + r.mutex.Lock() + defer r.mutex.Unlock() + + if reader, ok := r.readers[key]; ok { + return reader + } + + data, err := r.readCsvFile(key) if err != nil { - return err + return nil } - defer func() { - _ = file.Close() - }() + + reader := &fileReader{ + data: data, + index: atomic.Int64{}, + } + reader.index.Store(0) + r.readers[key] = reader + return reader +} + +func (r *csvReaderImpl) readCsvFile(sourceFile string) ([]Data, error) { + file, err := os.Open(sourceFile) + if err != nil { + return nil, err + } + defer file.Close() + reader := csv.NewReader(file) - comma := []rune(c.Delimiter) + comma := []rune(r.config.Delimiter) if len(comma) > 0 { reader.Comma = comma[0] } - if c.WithHeader { + + if r.config.WithHeader { _, err := reader.Read() - if err != nil { - return err + if err != nil && err != io.EOF { + return nil, err } } + + var lines []Data for { row, err := reader.Read() if err == io.EOF { break } if err != nil { - return err + return nil, err } - lines = append(lines, row) - if len(lines) == c.limit { + if r.config.Limit > 0 && len(lines) >= r.config.Limit { break } } + return lines, nil +} - go func() { - index := 0 - for { - if index == len(lines) { - index = 0 - } - dataCh <- lines[index] - index++ - } - }() +func (r *csvReaderImpl) Close() error { + r.mutex.Lock() + defer r.mutex.Unlock() + clear(r.readers) + r.readers = nil return nil +} + +type ( + CSVWriter struct { + Path string + Header []string + Delimiter string + DataCh <-chan []string + } +) +func NewCsvWriter(path, delimiter string, header []string, dataCh <-chan []string) *CSVWriter { + return &CSVWriter{ + Path: path, + Delimiter: delimiter, + Header: header, + DataCh: dataCh, + } } func (c *CSVWriter) WriteForever() error { diff --git a/pkg/common/types.go b/pkg/common/types.go index 3b547bd..11f08ce 100644 --- a/pkg/common/types.go +++ b/pkg/common/types.go @@ -26,6 +26,7 @@ type ( IGraphClient interface { IClient GetData() (Data, error) + GetFileData(sourceFile string) (Data, error) Execute(stmt string) (IGraphResponse, error) } @@ -46,8 +47,15 @@ type ( SetOption(*GraphOption) error } + CsvReaderConfig struct { + Delimiter string + WithHeader bool + Limit int + } + ICsvReader interface { - ReadForever(dataCh chan<- Data) error + GetData(sourceFile string) (Data, error) + Close() error } GraphOption struct { diff --git a/pkg/nebulagraph/client.go b/pkg/nebulagraph/client.go index d705043..ea9b00d 100644 --- a/pkg/nebulagraph/client.go +++ b/pkg/nebulagraph/client.go @@ -20,7 +20,6 @@ const EnvRetryTimeoutUs = "NEBULA_RETRY_TIMEOUT_US" type ( // GraphPool nebula connection pool GraphPool struct { - DataCh chan common.Data OutputCh chan []string initialized bool closed bool @@ -37,7 +36,6 @@ type ( GraphClient struct { Client *graph.Session Pool *GraphPool - DataCh chan common.Data logger logger } @@ -47,8 +45,6 @@ type ( ResponseTime int32 } - csvReaderStrategy int - output struct { timeStamp int64 nGQL string @@ -72,13 +68,6 @@ type ( var _ common.IGraphClient = &GraphClient{} var _ common.IGraphClientPool = &GraphPool{} -const ( - // AllInOne read csv sequentially - AllInOne csvReaderStrategy = iota - // Separate read csv concurrently - Separate -) - func formatOutput(o *output) []string { return []string{ strconv.FormatInt(o.timeStamp, 10), @@ -146,14 +135,12 @@ func (gp *GraphPool) Init() (common.IGraphClientPool, error) { if gp.graphOption.CsvPath != "" { gp.csvReader = common.NewCsvReader( gp.graphOption.CsvPath, - gp.graphOption.CsvDelimiter, - gp.graphOption.CsvWithHeader, - gp.graphOption.CsvDataLimit, + common.CsvReaderConfig{ + Delimiter: gp.graphOption.CsvDelimiter, + WithHeader: gp.graphOption.CsvWithHeader, + Limit: gp.graphOption.CsvDataLimit, + }, ) - gp.DataCh = make(chan common.Data, gp.graphOption.CsvChannelSize) - if err := gp.csvReader.ReadForever(gp.DataCh); err != nil { - return nil, err - } } return gp, nil } @@ -258,11 +245,6 @@ func (gp *GraphPool) validate(address string) ([]graph.HostAddress, error) { return hosts, nil } -// Deprecated ConfigCsvStrategy sets csv reader strategy -func (gp *GraphPool) ConfigCsvStrategy(strategy int) { - return -} - // Close closes the nebula pool func (gp *GraphPool) Close() error { gp.mutex.Lock() @@ -281,6 +263,9 @@ func (gp *GraphPool) Close() error { if gp.sessPool != nil { gp.sessPool.Close() } + if gp.csvReader != nil { + gp.csvReader.Close() + } gp.closed = true return nil @@ -302,11 +287,11 @@ func (gp *GraphPool) GetSession() (common.IGraphClient, error) { if err != nil { return nil, err } - s := &GraphClient{Client: c, Pool: gp, DataCh: gp.DataCh, logger: gp.logger} + s := &GraphClient{Client: c, Pool: gp, logger: gp.logger} gp.clients = append(gp.clients, s) return s, nil } else { - s := &GraphClient{Client: nil, Pool: gp, DataCh: gp.DataCh, logger: gp.logger} + s := &GraphClient{Client: nil, Pool: gp, logger: gp.logger} return s, nil } @@ -335,14 +320,18 @@ func (gc *GraphClient) Close() error { return nil } -// GetData get data from csv reader func (gc *GraphClient) GetData() (common.Data, error) { - if gc.DataCh != nil && len(gc.DataCh) != 0 { - if d, ok := <-gc.DataCh; ok { - return d, nil - } + if gc.Pool.csvReader != nil { + return gc.Pool.csvReader.GetData("") + } + return nil, fmt.Errorf("csv reader not initialized") +} + +func (gc *GraphClient) GetFileData(sourceFile string) (common.Data, error) { + if gc.Pool.csvReader != nil { + return gc.Pool.csvReader.GetData(sourceFile) } - return nil, fmt.Errorf("no Data at all") + return nil, fmt.Errorf("csv reader not initialized") } func (gc *GraphClient) executeRetry(stmt string) (*graph.ResultSet, error) { diff --git a/pkg/nebulagraph5/client.go b/pkg/nebulagraph5/client.go index 26b7fce..edc32a6 100644 --- a/pkg/nebulagraph5/client.go +++ b/pkg/nebulagraph5/client.go @@ -19,7 +19,6 @@ type ( // GraphPool nebula connection pool GraphPool struct { mutex sync.Mutex - DataCh chan common.Data OutputCh chan []string Version string csvStrategy csvReaderStrategy @@ -44,11 +43,9 @@ type ( graphClientGetter func(endpoint, username, password string, timeout time.Duration) (types.Client, error) GraphClientFactory struct{} - // GraphClient a wrapper for nebula client, could read data from DataCh GraphClient struct { Session types.Client Pool *GraphPool - DataCh chan common.Data username string password string since time.Time @@ -157,14 +154,12 @@ func (gp *GraphPool) Init() (common.IGraphClientPool, error) { if gp.graphOption.CsvPath != "" { gp.csvReader = common.NewCsvReader( gp.graphOption.CsvPath, - gp.graphOption.CsvDelimiter, - gp.graphOption.CsvWithHeader, - gp.graphOption.CsvDataLimit, + common.CsvReaderConfig{ + Delimiter: gp.graphOption.CsvDelimiter, + WithHeader: gp.graphOption.CsvWithHeader, + Limit: gp.graphOption.CsvDataLimit, + }, ) - gp.DataCh = make(chan common.Data, gp.graphOption.CsvChannelSize) - if err := gp.csvReader.ReadForever(gp.DataCh); err != nil { - return nil, err - } } options := []nebula.PoolOptionsFn{ @@ -230,7 +225,6 @@ func (gp *GraphPool) validate(address string) error { return nil } -// Close closes the nebula pool func (gp *GraphPool) Close() error { gp.mutex.Lock() defer gp.mutex.Unlock() @@ -238,6 +232,9 @@ func (gp *GraphPool) Close() error { client.Close() } gp.pool.Close() + if gp.csvReader != nil { + gp.csvReader.Close() + } return nil } @@ -249,7 +246,7 @@ func (gp *GraphPool) GetSession() (common.IGraphClient, error) { return nil, fmt.Errorf("GraphPool is not initialized, please call Init() first") } - s := &GraphClient{Pool: gp, DataCh: gp.DataCh, since: time.Now()} + s := &GraphClient{Pool: gp, since: time.Now()} gp.clients = append(gp.clients, s) return s, nil } @@ -282,14 +279,18 @@ func (gc *GraphClient) Close() error { return nil } -// GetData get data from csv reader func (gc *GraphClient) GetData() (common.Data, error) { - if gc.DataCh != nil && len(gc.DataCh) != 0 { - if d, ok := <-gc.DataCh; ok { - return d, nil - } + if gc.Pool.csvReader != nil { + return gc.Pool.csvReader.GetData("") + } + return nil, fmt.Errorf("csv reader not initialized") +} + +func (gc *GraphClient) GetFileData(sourceFile string) (common.Data, error) { + if gc.Pool.csvReader != nil { + return gc.Pool.csvReader.GetData(sourceFile) } - return nil, fmt.Errorf("no Data at all") + return nil, fmt.Errorf("csv reader not initialized") } // Execute executes nebula query From a8c81e179875fae5ff30ef8398b1b450ce850305 Mon Sep 17 00:00:00 2001 From: "Harris.Chu" <1726587+HarrisChu@users.noreply.github.com> Date: Mon, 30 Mar 2026 16:29:47 +0800 Subject: [PATCH 2/2] update --- .github/workflows/pr.yaml | 2 +- .github/workflows/release.yaml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 18febb6..04f46ac 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -15,7 +15,7 @@ jobs: steps: - uses: actions/checkout@v2 - - uses: actions/setup-go@v1 + - uses: actions/setup-go@v6 with: go-version: 1.24.11 diff --git a/.github/workflows/release.yaml b/.github/workflows/release.yaml index 1201c6e..3370f65 100644 --- a/.github/workflows/release.yaml +++ b/.github/workflows/release.yaml @@ -12,7 +12,7 @@ jobs: steps: - uses: actions/checkout@v2 - - uses: actions/setup-go@v1 + - uses: actions/setup-go@v6 with: go-version: 1.24.11