feat(extract): introduce --output flag

ccamel · ccamel · commit dbc87a1bc583 · 2025-01-16T13:22:25.000+01:00
diff --git a/cmd/delegators.go b/cmd/delegators.go
@@ -3,21 +3,31 @@ package cmd
 import (
 	"github.com/axone-protocol/cosmos-extractor/pkg/delegators"
 	"github.com/spf13/cobra"
+	"github.com/teambenny/goetl"
 )
 
 var extractDelegatorsCmd = &cobra.Command{
-	Use:   "delegators [source] [dest]",
-	Short: "Extract all delegators into a CSV file",
-	Args:  cobra.ExactArgs(2),
+	Use:   "delegators [source]",
+	Short: "Extract all delegators",
+	Args:  cobra.ExactArgs(1),
 	RunE: func(cmd *cobra.Command, args []string) error {
 		chainName, _ := cmd.Flags().GetString(flagChainName)
-		pipeline, err := delegators.Pipeline(chainName, args[0], args[1], logger)
+		src := args[0]
+
+		read, err := delegators.NewDelegatorsReader(chainName, src, logger)
 		if err != nil {
 			return err
 		}
 
-		err = <-pipeline.Run()
+		write, err := newCSVWriter(cmd, args)
+		if err != nil {
+			return err
+		}
+
+		pipeline := goetl.NewPipeline(read, write)
+		pipeline.Name = "Delegators"
 
+		err = <-pipeline.Run()
 		return err
 	},
 }
diff --git a/cmd/extract.go b/cmd/extract.go
@@ -1,11 +1,16 @@
 package cmd
 
 import (
+	"os"
+
+	"github.com/axone-protocol/cosmos-extractor/pkg/csv"
 	"github.com/spf13/cobra"
+	"github.com/teambenny/goetl"
 )
 
 const (
 	flagChainName = "chain-name"
+	flagOutput    = "output"
 )
 
 var extractCmd = &cobra.Command{
@@ -17,4 +22,19 @@ func init() {
 	rootCmd.AddCommand(extractCmd)
 
 	extractCmd.PersistentFlags().StringP(flagChainName, "n", "cosmos", "Name of the chain")
+	extractCmd.PersistentFlags().StringP(flagOutput, "o", "", "Output file (defaults to stdout)")
+}
+
+func newCSVWriter(cmd *cobra.Command, _ []string) (goetl.Processor, error) {
+	options := []csv.Option{
+		csv.WithWriterHeader(),
+	}
+	output, _ := cmd.Flags().GetString(flagOutput)
+	if output != "" {
+		options = append(options, csv.WithFile(output))
+	} else {
+		options = append(options, csv.WithWriter(os.Stdout))
+	}
+
+	return csv.NewCSVWriter(options...)
 }
diff --git a/cmd/infos.go b/cmd/infos.go
@@ -3,21 +3,31 @@ package cmd
 import (
 	"github.com/axone-protocol/cosmos-extractor/pkg/infos"
 	"github.com/spf13/cobra"
+	"github.com/teambenny/goetl"
 )
 
 var extractInfosCmd = &cobra.Command{
 	Use:   "infos [source] [dest]",
-	Short: "Extract chain informations into a CSV file",
-	Args:  cobra.ExactArgs(2),
+	Short: "Extract chain informations",
+	Args:  cobra.ExactArgs(1),
 	RunE: func(cmd *cobra.Command, args []string) error {
 		chainName, _ := cmd.Flags().GetString(flagChainName)
-		pipeline, err := infos.Pipeline(chainName, args[0], args[1], logger)
+		src := args[0]
+
+		read, err := infos.NewInfoReader(chainName, src, logger)
 		if err != nil {
 			return err
 		}
 
-		err = <-pipeline.Run()
+		write, err := newCSVWriter(cmd, args)
+		if err != nil {
+			return err
+		}
+
+		pipeline := goetl.NewPipeline(read, write)
+		pipeline.Name = "Delegators"
 
+		err = <-pipeline.Run()
 		return err
 	},
 }
diff --git a/cmd/root.go b/cmd/root.go
@@ -18,8 +18,7 @@ const (
 // rootCmd represents the base command when called without any subcommands.
 var rootCmd = &cobra.Command{
 	Use:   "cosmos-extractor",
-	Short: "A template fo Golang projects",
-	Long:  "A template fo Golang projects",
+	Short: "Tool for extracting diverse data from Cosmos chain snapshots",
 	PersistentPreRunE: func(cmd *cobra.Command, _ []string) error {
 		logLevel := cmd.Flag(flagLogLevel).Value.String()
 		filterFn, err := log.ParseLogLevel(logLevel)
diff --git a/cmd/version.go b/cmd/version.go
@@ -11,7 +11,7 @@ import (
 
 const (
 	flagLong   = "long"
-	flagOutput = "output"
+	flagFormat = "format"
 )
 
 // NewVersionCommand returns a CLI command to interactively print the application binary version information.
@@ -31,8 +31,8 @@ var versionCmd = &cobra.Command{
 			err error
 		)
 
-		output, _ := cmd.Flags().GetString(flagOutput)
-		switch strings.ToLower(output) {
+		format, _ := cmd.Flags().GetString(flagFormat)
+		switch strings.ToLower(format) {
 		case "json":
 			bz, err = json.Marshal(verInfo)
 
@@ -53,5 +53,5 @@ func init() {
 	rootCmd.AddCommand(versionCmd)
 
 	versionCmd.Flags().Bool(flagLong, false, "Print long version information")
-	versionCmd.Flags().StringP(flagOutput, "o", "text", "Output format (text|json)")
+	versionCmd.Flags().StringP(flagFormat, "f", "text", "Output format (text|json)")
 }
diff --git a/pkg/csv/writer.go b/pkg/csv/writer.go
@@ -0,0 +1,138 @@
+package csv
+
+import (
+	"bufio"
+	"fmt"
+	"io"
+	"os"
+
+	"github.com/teambenny/goetl"
+	"github.com/teambenny/goetl/etldata"
+	"github.com/teambenny/goetl/etlutil"
+	"github.com/teambenny/goetl/processors"
+)
+
+type Option func(*writer) error
+
+func WithWriterHeader() Option {
+	return func(w *writer) error {
+		w.delegated.Parameters.WriteHeader = true
+
+		return nil
+	}
+}
+
+func WithFile(file string) Option {
+	return WithFileAndFlag(file, os.O_RDWR|os.O_CREATE)
+}
+
+func WithFileAndFlag(file string, flag int) Option {
+	return func(w *writer) error {
+		file, err := os.OpenFile(file, flag, os.ModePerm)
+		if err != nil {
+			return err
+		}
+		bw := bufio.NewWriter(file)
+
+		delegated := processors.NewCSVWriter(bw)
+
+		if w.delegated != nil {
+			copyCSVWriterParameters(w.delegated.Parameters, &delegated.Parameters)
+		} else {
+			delegated.Parameters.WriteHeader = false
+		}
+
+		w.finalizer = func() error {
+			if err := bw.Flush(); err != nil {
+				return err
+			}
+			return file.Close()
+		}
+		w.delegated = delegated
+		w.processorName = fmt.Sprintf("CSVWriter<%s>", file.Name())
+
+		return nil
+	}
+}
+
+func WithWriter(w io.Writer) Option {
+	return func(cw *writer) error {
+		delegated := processors.NewCSVWriter(w)
+
+		if cw.delegated != nil {
+			copyCSVWriterParameters(cw.delegated.Parameters, &delegated.Parameters)
+		} else {
+			delegated.Parameters.WriteHeader = false
+		}
+
+		if closer, ok := w.(flusher); ok {
+			cw.finalizer = closer.Flush
+		} else {
+			cw.finalizer = func() error { return nil }
+		}
+
+		cw.delegated = delegated
+
+		if stringer, ok := w.(fmt.Stringer); ok {
+			cw.processorName = fmt.Sprintf("CSVWriter<%s>", stringer.String())
+		} else {
+			cw.processorName = "CSVWriter<custom>"
+		}
+
+		return nil
+	}
+}
+
+type flusher interface {
+	Flush() error
+}
+
+type writer struct {
+	finalizer     func() error
+	delegated     *processors.CSVWriter
+	processorName string
+}
+
+func NewCSVWriter(options ...Option) (goetl.Processor, error) {
+	delegated := defaultCSVWriter()
+	writer := &writer{
+		finalizer:     func() error { return nil },
+		delegated:     delegated,
+		processorName: "CSVWriter<null>",
+	}
+
+	for _, option := range options {
+		err := option(writer)
+		if err != nil {
+			return nil, err
+		}
+	}
+
+	return writer, nil
+}
+
+func (w *writer) ProcessData(d etldata.Payload, outputChan chan etldata.Payload, killChan chan error) {
+	w.delegated.ProcessData(d, outputChan, killChan)
+}
+
+func (w *writer) Finish(outputChan chan etldata.Payload, killChan chan error) {
+	w.delegated.Finish(outputChan, killChan)
+	killChan <- w.finalizer()
+}
+
+func (w *writer) String() string {
+	return w.processorName
+}
+
+func defaultCSVWriter() *processors.CSVWriter {
+	csvW := processors.NewCSVWriter(io.Discard)
+	csvW.Parameters.WriteHeader = false
+
+	return csvW
+}
+
+func copyCSVWriterParameters(from etlutil.CSVParameters, to *etlutil.CSVParameters) {
+	to.WriteHeader = from.WriteHeader
+	to.HeaderWritten = from.HeaderWritten
+	to.SendUpstream = from.SendUpstream
+}
diff --git a/pkg/delegators/pipeline.go b/pkg/delegators/pipeline.go
diff --git a/pkg/delegators/reader.go b/pkg/delegators/reader.go
@@ -1,6 +1,7 @@
 package delegators
 
 import (
+	"errors"
 	"fmt"
 	"io"
 	"strings"
@@ -15,6 +16,7 @@ import (
 
 	sdk "github.com/cosmos/cosmos-sdk/types"
 	"github.com/cosmos/cosmos-sdk/types/bech32"
+	stakingtypes "github.com/cosmos/cosmos-sdk/x/staking/types"
 )
 
 type delegatorsReader struct {
@@ -68,19 +70,18 @@ func (r *delegatorsReader) ProcessData(_ etldata.Payload, outputChan chan etldat
 
 	keepers.Bank.IterateAllBalances(ctx, func(addr sdk.AccAddress, _ sdk.Coin) (stop bool) {
 		for _, val := range validators {
-			if config.GetBech32AccountAddrPrefix() != prefix {
-				config.SetBech32PrefixForValidator(
-					fmt.Sprintf("%svaloper", prefix),
-					fmt.Sprintf("%svaloperpub", prefix),
-				)
-			}
-
 			valAddr, err := sdk.ValAddressFromBech32(val.OperatorAddress)
 			etlutil.KillPipelineIfErr(err, killChan)
 
 			delegation, err := keepers.Staking.GetDelegation(ctx, addr, valAddr)
 			if err != nil {
-				continue
+				if errors.Is(err, stakingtypes.ErrNoDelegation) {
+					continue
+				}
+
+				r.logger.Error(err.Error())
+				killChan <- err
+				return true
 			}
 
 			payload := Delegation{
diff --git a/pkg/infos/pipeline.go b/pkg/infos/pipeline.go
diff --git a/pkg/processors/csv_writer.go b/pkg/processors/csv_writer.go
