storj/cmd/tools/segment-verify/csv.go
Egon Elbre cd81c5bd58 cmd/tools/segment-verify: add csv writer
Change-Id: I9306d7a6927f4dacca9623d7bd57f8560404db3e
2022-09-15 13:28:21 +00:00

92 lines
1.7 KiB
Go

// Copyright (C) 2022 Storj Labs, Inc.
// See LICENSE for copying information.
package main
import (
"context"
"encoding/csv"
"fmt"
"io"
"os"
)
// CSVWriter writes segments to a file.
type CSVWriter struct {
header bool
file io.WriteCloser
wr *csv.Writer
}
var _ SegmentWriter = (*CSVWriter)(nil)
// NewCSVWriter creates a new segment writer that writes to the specified path.
func NewCSVWriter(path string) (*CSVWriter, error) {
f, err := os.Create(path)
if err != nil {
return nil, Error.Wrap(err)
}
return &CSVWriter{
file: f,
wr: csv.NewWriter(f),
}, nil
}
// NewCustomCSVWriter creates a new segment writer that writes to the io.Writer.
func NewCustomCSVWriter(w io.Writer) *CSVWriter {
return &CSVWriter{
file: nopCloser{w},
wr: csv.NewWriter(w),
}
}
// Close closes the writer.
func (csv *CSVWriter) Close() error {
return Error.Wrap(csv.file.Close())
}
// Write writes and flushes the segments.
func (csv *CSVWriter) Write(ctx context.Context, segments []*Segment) (err error) {
defer mon.Task()(&ctx)(&err)
if !csv.header {
csv.header = true
err := csv.wr.Write([]string{
"stream id",
"position",
"found",
"not found",
"retry",
})
if err != nil {
return Error.Wrap(err)
}
}
defer csv.wr.Flush()
for _, seg := range segments {
if ctx.Err() != nil {
return Error.Wrap(err)
}
err := csv.wr.Write([]string{
seg.StreamID.String(),
fmt.Sprint(seg.Position.Encode()),
fmt.Sprint(seg.Status.Found),
fmt.Sprint(seg.Status.NotFound),
fmt.Sprint(seg.Status.Retry),
})
if err != nil {
return Error.Wrap(err)
}
}
return nil
}
// nopCloser adds Close method to a writer.
type nopCloser struct{ io.Writer }
func (nopCloser) Close() error { return nil }