diff --git a/lib/cmd_export.go b/lib/cmd_export.go index 0a29d0c..adf6df8 100644 --- a/lib/cmd_export.go +++ b/lib/cmd_export.go @@ -1,8 +1,6 @@ package lib import ( - "encoding/csv" - "encoding/json" "errors" "fmt" "os" @@ -78,7 +76,7 @@ func CmdExport(f CmdExportFlags, args []string, printHelp func()) error { defer outFile.Close() } - // validate format. + // infer format from extension if not specified. if f.Format == "" { if strings.HasSuffix(f.Out, ".csv") { f.Format = "csv" @@ -90,9 +88,6 @@ func CmdExport(f CmdExportFlags, args []string, printHelp func()) error { f.Format = "csv" } } - if f.Format != "csv" && f.Format != "tsv" && f.Format != "json" { - return errors.New("format must be \"csv\" or \"tsv\" or \"json\"") - } // open tree. db, err := maxminddb.Open(args[0]) @@ -101,61 +96,17 @@ func CmdExport(f CmdExportFlags, args []string, printHelp func()) error { } defer db.Close() - if f.Format == "tsv" || f.Format == "csv" { - // export. - hdrWritten := false - var wr writer - if f.Format == "csv" { - csvwr := csv.NewWriter(outFile) - wr = csvwr - } else { - tsvwr := NewTsvWriter(outFile) - wr = tsvwr - } - for result := range db.Networks() { - record := make(map[string]interface{}) - if err := result.Decode(&record); err != nil { - return fmt.Errorf("failed to get record for next subnet: %w", err) - } - subnet := result.Prefix() - - recordStr := mapInterfaceToStr(record) - if !hdrWritten { - hdrWritten = true - - if !f.NoHdr { - hdr := append([]string{"range"}, sortedMapKeys(recordStr)...) - if err := wr.Write(hdr); err != nil { - return fmt.Errorf( - "failed to write header %v: %w", - hdr, err, - ) - } - } - } - - line := append( - []string{subnet.String()}, - sortedMapValsByKeys(recordStr)..., - ) - if err := wr.Write(line); err != nil { - return fmt.Errorf("failed to write line %v: %w", line, err) - } - } - wr.Flush() - if err := wr.Error(); err != nil { - return fmt.Errorf("writer had failure: %w", err) - } - } else if f.Format == "json" { - enc := json.NewEncoder(outFile) - for result := range db.Networks() { - record := make(map[string]interface{}) - if err := result.Decode(&record); err != nil { - return fmt.Errorf("failed to get record for next subnet: %w", err) - } - record["range"] = result.Prefix().String() - enc.Encode(record) - } + var exp exporter + switch f.Format { + case "csv": + exp = newCSVExporter(outFile, f.NoHdr) + case "tsv": + exp = newTSVExporter(outFile, f.NoHdr) + case "json": + exp = newJSONExporter(outFile) + default: + return errors.New("format must be \"csv\" or \"tsv\" or \"json\"") } - return nil + + return exportNetworks(db, exp) } diff --git a/lib/exporter.go b/lib/exporter.go new file mode 100644 index 0000000..7522c0b --- /dev/null +++ b/lib/exporter.go @@ -0,0 +1,26 @@ +package lib + +import ( + "fmt" + + "github.com/oschwald/maxminddb-golang/v2" +) + +// exporter defines the interface for exporting MMDB records. +type exporter interface { + WriteRecord(result maxminddb.Result) error + Flush() error +} + +// exportNetworks iterates over all networks in the database and writes them using the exporter. +func exportNetworks(db *maxminddb.Reader, exp exporter) error { + for result := range db.Networks() { + if err := result.Err(); err != nil { + return fmt.Errorf("failed networks traversal: %w", err) + } + if err := exp.WriteRecord(result); err != nil { + return err + } + } + return exp.Flush() +} diff --git a/lib/exporter_csv.go b/lib/exporter_csv.go new file mode 100644 index 0000000..f8bc2aa --- /dev/null +++ b/lib/exporter_csv.go @@ -0,0 +1,70 @@ +package lib + +import ( + "encoding/csv" + "fmt" + "io" + + "github.com/oschwald/maxminddb-golang/v2" +) + +// csvExporter exports records in CSV format. +type csvExporter struct { + wr *csv.Writer + cache map[uintptr]map[string]string + hdrKeys []string + noHdr bool +} + +func newCSVExporter(w io.Writer, noHdr bool) *csvExporter { + return &csvExporter{ + wr: csv.NewWriter(w), + cache: make(map[uintptr]map[string]string), + noHdr: noHdr, + } +} + +func (e *csvExporter) WriteRecord(result maxminddb.Result) error { + offset := result.Offset() + prefix := result.Prefix() + + var recordStr map[string]string + if cached, ok := e.cache[offset]; ok { + recordStr = cached + } else { + record := make(map[string]any) + if err := result.Decode(&record); err != nil { + return fmt.Errorf("failed to decode record: %w", err) + } + recordStr = mapInterfaceToStr(record) + e.cache[offset] = recordStr + } + + // Write header on first record. + if e.hdrKeys == nil { + e.hdrKeys = sortedMapKeys(recordStr) + if !e.noHdr { + hdr := append([]string{"range"}, e.hdrKeys...) + if err := e.wr.Write(hdr); err != nil { + return fmt.Errorf("failed to write header %v: %w", hdr, err) + } + } + } + + // Build values in header key order. + vals := make([]string, len(e.hdrKeys)) + for i, k := range e.hdrKeys { + vals[i] = recordStr[k] + } + + line := append([]string{prefix.String()}, vals...) + if err := e.wr.Write(line); err != nil { + return fmt.Errorf("failed to write line %v: %w", line, err) + } + return nil +} + +func (e *csvExporter) Flush() error { + e.wr.Flush() + return e.wr.Error() +} diff --git a/lib/exporter_json.go b/lib/exporter_json.go new file mode 100644 index 0000000..a709a3a --- /dev/null +++ b/lib/exporter_json.go @@ -0,0 +1,56 @@ +package lib + +import ( + "bufio" + "bytes" + "encoding/json" + "fmt" + "io" + + "github.com/oschwald/maxminddb-golang/v2" +) + +const rangePlaceholder = "__RANGE__" + +// jsonExporter exports records in JSON Lines format. +type jsonExporter struct { + bw *bufio.Writer + cache map[uintptr][]byte +} + +func newJSONExporter(w io.Writer) *jsonExporter { + return &jsonExporter{ + bw: bufio.NewWriter(w), + cache: make(map[uintptr][]byte), + } +} + +func (e *jsonExporter) WriteRecord(result maxminddb.Result) error { + offset := result.Offset() + prefix := result.Prefix() + + cached, ok := e.cache[offset] + if !ok { + record := make(map[string]any) + if err := result.Decode(&record); err != nil { + return fmt.Errorf("failed to decode record: %w", err) + } + record["range"] = rangePlaceholder + + encoded, err := json.Marshal(record) + if err != nil { + return fmt.Errorf("failed to encode record: %w", err) + } + cached = encoded + e.cache[offset] = cached + } + + line := bytes.Replace(cached, []byte(rangePlaceholder), []byte(prefix.String()), 1) + e.bw.Write(line) + e.bw.WriteByte('\n') + return nil +} + +func (e *jsonExporter) Flush() error { + return e.bw.Flush() +} diff --git a/lib/exporter_tsv.go b/lib/exporter_tsv.go new file mode 100644 index 0000000..a4e12a0 --- /dev/null +++ b/lib/exporter_tsv.go @@ -0,0 +1,69 @@ +package lib + +import ( + "fmt" + "io" + + "github.com/oschwald/maxminddb-golang/v2" +) + +// tsvExporter exports records in TSV format. +type tsvExporter struct { + wr *TsvWriter + cache map[uintptr]map[string]string + hdrKeys []string + noHdr bool +} + +func newTSVExporter(w io.Writer, noHdr bool) *tsvExporter { + return &tsvExporter{ + wr: NewTsvWriter(w), + cache: make(map[uintptr]map[string]string), + noHdr: noHdr, + } +} + +func (e *tsvExporter) WriteRecord(result maxminddb.Result) error { + offset := result.Offset() + prefix := result.Prefix() + + var recordStr map[string]string + if cached, ok := e.cache[offset]; ok { + recordStr = cached + } else { + record := make(map[string]any) + if err := result.Decode(&record); err != nil { + return fmt.Errorf("failed to decode record: %w", err) + } + recordStr = mapInterfaceToStr(record) + e.cache[offset] = recordStr + } + + // Write header on first record. + if e.hdrKeys == nil { + e.hdrKeys = sortedMapKeys(recordStr) + if !e.noHdr { + hdr := append([]string{"range"}, e.hdrKeys...) + if err := e.wr.Write(hdr); err != nil { + return fmt.Errorf("failed to write header %v: %w", hdr, err) + } + } + } + + // Build values in header key order. + vals := make([]string, len(e.hdrKeys)) + for i, k := range e.hdrKeys { + vals[i] = recordStr[k] + } + + line := append([]string{prefix.String()}, vals...) + if err := e.wr.Write(line); err != nil { + return fmt.Errorf("failed to write line %v: %w", line, err) + } + return nil +} + +func (e *tsvExporter) Flush() error { + e.wr.Flush() + return e.wr.Error() +} diff --git a/lib/writer_tsv.go b/lib/writer_tsv.go index 3b3e718..11d29b4 100644 --- a/lib/writer_tsv.go +++ b/lib/writer_tsv.go @@ -1,27 +1,31 @@ package lib import ( - "fmt" + "bufio" "io" "strings" ) type TsvWriter struct { - w io.Writer + bw *bufio.Writer } func NewTsvWriter(w io.Writer) *TsvWriter { return &TsvWriter{ - w: w, + bw: bufio.NewWriter(w), } } func (w *TsvWriter) Write(record []string) error { - _, err := fmt.Fprintln(w.w, strings.Join(record, "\t")) - return err + _, err := w.bw.WriteString(strings.Join(record, "\t")) + if err != nil { + return err + } + return w.bw.WriteByte('\n') } func (w *TsvWriter) Flush() { + w.bw.Flush() } func (w *TsvWriter) Error() error {