storj/cmd/tools/segment-verify/csv.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 }