github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/commands/offsets.go (about)

     1  package commands
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"os"
     7  
     8  	agent "github.com/observiq/carbon/agent"
     9  	"github.com/observiq/carbon/operator/helper"
    10  	"github.com/spf13/cobra"
    11  	"go.etcd.io/bbolt"
    12  )
    13  
    14  var stdout io.Writer = os.Stdout
    15  
    16  // NewOffsetsCmd returns the root command for managing offsets
    17  func NewOffsetsCmd(rootFlags *RootFlags) *cobra.Command {
    18  	offsets := &cobra.Command{
    19  		Use:   "offsets",
    20  		Short: "Manage input operator offsets",
    21  		Args:  cobra.NoArgs,
    22  		Run: func(command *cobra.Command, args []string) {
    23  			stdout.Write([]byte("No offsets subcommand specified. See `carbon offsets help` for details\n"))
    24  		},
    25  	}
    26  
    27  	offsets.AddCommand(NewOffsetsClearCmd(rootFlags))
    28  	offsets.AddCommand(NewOffsetsListCmd(rootFlags))
    29  
    30  	return offsets
    31  }
    32  
    33  // NewOffsetsClearCmd returns the command for clearing offsets
    34  func NewOffsetsClearCmd(rootFlags *RootFlags) *cobra.Command {
    35  	var all bool
    36  
    37  	offsetsClear := &cobra.Command{
    38  		Use:   "clear [flags] [operator_ids]",
    39  		Short: "Clear persisted offsets from the database",
    40  		Args:  cobra.ArbitraryArgs,
    41  		Run: func(command *cobra.Command, args []string) {
    42  			db, err := agent.OpenDatabase(rootFlags.DatabaseFile)
    43  			exitOnErr("Failed to open database", err)
    44  			defer db.Close()
    45  			defer db.Sync()
    46  
    47  			if all {
    48  				if len(args) != 0 {
    49  					stdout.Write([]byte("Providing a list of operator IDs does nothing with the --all flag\n"))
    50  				}
    51  
    52  				err := db.Update(func(tx *bbolt.Tx) error {
    53  					offsetsBucket := tx.Bucket(helper.OffsetsBucket)
    54  					if offsetsBucket != nil {
    55  						return tx.DeleteBucket(helper.OffsetsBucket)
    56  					}
    57  					return nil
    58  				})
    59  				exitOnErr("Failed to delete offsets", err)
    60  			} else {
    61  				if len(args) == 0 {
    62  					stdout.Write([]byte("Must either specify a list of operators or the --all flag\n"))
    63  					os.Exit(1)
    64  				}
    65  
    66  				for _, operatorID := range args {
    67  					err = db.Update(func(tx *bbolt.Tx) error {
    68  						offsetBucket := tx.Bucket(helper.OffsetsBucket)
    69  						if offsetBucket == nil {
    70  							return nil
    71  						}
    72  
    73  						return offsetBucket.DeleteBucket([]byte(operatorID))
    74  					})
    75  					exitOnErr("Failed to delete offsets", err)
    76  				}
    77  			}
    78  		},
    79  	}
    80  
    81  	offsetsClear.Flags().BoolVar(&all, "all", false, "clear offsets for all inputs")
    82  
    83  	return offsetsClear
    84  }
    85  
    86  // NewOffsetsListCmd returns the command for listing offsets
    87  func NewOffsetsListCmd(rootFlags *RootFlags) *cobra.Command {
    88  	offsetsList := &cobra.Command{
    89  		Use:   "list",
    90  		Short: "List operators with persisted offsets",
    91  		Args:  cobra.NoArgs,
    92  		Run: func(command *cobra.Command, args []string) {
    93  			db, err := agent.OpenDatabase(rootFlags.DatabaseFile)
    94  			exitOnErr("Failed to open database", err)
    95  			defer db.Close()
    96  
    97  			db.View(func(tx *bbolt.Tx) error {
    98  				offsetBucket := tx.Bucket(helper.OffsetsBucket)
    99  				if offsetBucket == nil {
   100  					return nil
   101  				}
   102  
   103  				return offsetBucket.ForEach(func(key, value []byte) error {
   104  					stdout.Write(append(key, '\n'))
   105  					return nil
   106  				})
   107  			})
   108  
   109  		},
   110  	}
   111  
   112  	return offsetsList
   113  }
   114  
   115  func exitOnErr(msg string, err error) {
   116  	if err != nil {
   117  		os.Stderr.WriteString(fmt.Sprintf("%s: %s\n", msg, err))
   118  		os.Exit(1)
   119  	}
   120  }