Skip to content

Commit

Permalink
Refine: read large CSV file line by line
Browse files Browse the repository at this point in the history
  • Loading branch information
Loyalsoldier committed Aug 10, 2024
1 parent 4b52a4d commit 596da97
Showing 1 changed file with 38 additions and 20 deletions.
58 changes: 38 additions & 20 deletions plugin/maxmind/country_csv.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,8 +3,8 @@ package maxmind
import (
"encoding/csv"
"encoding/json"
"errors"
"fmt"
"io"
"os"
"path/filepath"
"strings"
Expand Down Expand Up @@ -159,53 +159,71 @@ func (g *geoLite2CountryCSV) getCountryCode() (map[string]string, error) {

ccMap := make(map[string]string)
for _, line := range lines[1:] {
if len(line) < 5 {
return nil, fmt.Errorf("❌ [type %s | action %s] invalid record: %v", typeCountryCSV, g.Action, line)
}

id := strings.TrimSpace(line[0])
countryCode := strings.TrimSpace(line[4])
if id == "" || countryCode == "" {
continue
}
ccMap[id] = strings.ToUpper(countryCode)
}

if len(ccMap) == 0 {
return nil, fmt.Errorf("❌ [type %s | action %s] invalid country code data", typeCountryCSV, g.Action)
}

return ccMap, nil
}

func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entries map[string]*lib.Entry) error {
if len(ccMap) == 0 {
return errors.New("country code list must be specified")
return fmt.Errorf("❌ [type %s | action %s] invalid country code data", typeCountryCSV, g.Action)
}
if entries == nil {
entries = make(map[string]*lib.Entry, 300)
}

// Filter want list
wantList := make(map[string]bool)
for _, want := range g.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList[want] = true
}
}

fReader, err := os.Open(file)
if err != nil {
return err
}
defer fReader.Close()

reader := csv.NewReader(fReader)
lines, err := reader.ReadAll()
if err != nil {
return err
}
reader.Read() // skip header

// Filter want list
wantList := make(map[string]bool)
for _, want := range g.Want {
if want = strings.ToUpper(strings.TrimSpace(want)); want != "" {
wantList[want] = true
for {
record, err := reader.Read()
if err == io.EOF {
break
}
if err != nil {
return err
}

if len(record) < 4 {
return fmt.Errorf("❌ [type %s | action %s] invalid record: %v", typeCountryCSV, g.Action, record)
}
}

for _, line := range lines[1:] {
ccID := ""
switch {
case strings.TrimSpace(line[1]) != "":
ccID = strings.TrimSpace(line[1])
case strings.TrimSpace(line[2]) != "":
ccID = strings.TrimSpace(line[2])
case strings.TrimSpace(line[3]) != "":
ccID = strings.TrimSpace(line[3])
case strings.TrimSpace(record[1]) != "":
ccID = strings.TrimSpace(record[1])
case strings.TrimSpace(record[2]) != "":
ccID = strings.TrimSpace(record[2])
case strings.TrimSpace(record[3]) != "":
ccID = strings.TrimSpace(record[3])
default:
continue
}
Expand All @@ -214,7 +232,7 @@ func (g *geoLite2CountryCSV) process(file string, ccMap map[string]string, entri
if len(wantList) > 0 && !wantList[countryCode] {
continue
}
cidrStr := strings.ToLower(strings.TrimSpace(line[0]))
cidrStr := strings.ToLower(strings.TrimSpace(record[0]))
entry, found := entries[countryCode]
if !found {
entry = lib.NewEntry(countryCode)
Expand Down

0 comments on commit 596da97

Please sign in to comment.