92 lines
1.7 KiB
Go
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 }
|