github.com/authzed/spicedb@v1.32.1-0.20240520085336-ebda56537386/pkg/cmd/datastore.go (about)

     1  package cmd
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/spf13/cobra"
     9  
    10  	"github.com/authzed/spicedb/internal/datastore/common"
    11  	log "github.com/authzed/spicedb/internal/logging"
    12  	"github.com/authzed/spicedb/pkg/cmd/datastore"
    13  	"github.com/authzed/spicedb/pkg/cmd/server"
    14  	"github.com/authzed/spicedb/pkg/cmd/termination"
    15  	dspkg "github.com/authzed/spicedb/pkg/datastore"
    16  )
    17  
    18  func RegisterDatastoreRootFlags(_ *cobra.Command) {
    19  }
    20  
    21  func NewDatastoreCommand(programName string) (*cobra.Command, error) {
    22  	datastoreCmd := &cobra.Command{
    23  		Use:   "datastore",
    24  		Short: "datastore operations",
    25  		Long:  "Operations against the configured datastore",
    26  	}
    27  
    28  	migrateCmd := NewMigrateCommand(programName)
    29  	RegisterMigrateFlags(migrateCmd)
    30  	datastoreCmd.AddCommand(migrateCmd)
    31  
    32  	cfg := datastore.Config{}
    33  
    34  	gcCmd := NewGCDatastoreCommand(programName, &cfg)
    35  	if err := datastore.RegisterDatastoreFlagsWithPrefix(gcCmd.Flags(), "", &cfg); err != nil {
    36  		return nil, err
    37  	}
    38  	datastoreCmd.AddCommand(gcCmd)
    39  
    40  	repairCmd := NewRepairDatastoreCommand(programName, &cfg)
    41  	if err := datastore.RegisterDatastoreFlagsWithPrefix(repairCmd.Flags(), "", &cfg); err != nil {
    42  		return nil, err
    43  	}
    44  	datastoreCmd.AddCommand(repairCmd)
    45  
    46  	headCmd := NewHeadCommand(programName)
    47  	RegisterHeadFlags(headCmd)
    48  	datastoreCmd.AddCommand(headCmd)
    49  
    50  	return datastoreCmd, nil
    51  }
    52  
    53  func NewGCDatastoreCommand(programName string, cfg *datastore.Config) *cobra.Command {
    54  	return &cobra.Command{
    55  		Use:     "gc",
    56  		Short:   "executes garbage collection",
    57  		Long:    "Executes garbage collection against the datastore",
    58  		PreRunE: server.DefaultPreRunE(programName),
    59  		RunE: termination.PublishError(func(cmd *cobra.Command, args []string) error {
    60  			ctx := context.Background()
    61  
    62  			// Disable background GC and hedging.
    63  			cfg.GCInterval = -1 * time.Hour
    64  			cfg.RequestHedgingEnabled = false
    65  
    66  			ds, err := datastore.NewDatastore(ctx, cfg.ToOption())
    67  			if err != nil {
    68  				return fmt.Errorf("failed to create datastore: %w", err)
    69  			}
    70  
    71  			gc := dspkg.UnwrapAs[common.GarbageCollector](ds)
    72  			if gc == nil {
    73  				return fmt.Errorf("datastore of type %T does not support garbage collection", ds)
    74  			}
    75  
    76  			log.Ctx(ctx).Info().
    77  				Float64("gc_window_seconds", cfg.GCWindow.Seconds()).
    78  				Float64("gc_max_operation_time_seconds", cfg.GCMaxOperationTime.Seconds()).
    79  				Msg("Running garbage collection...")
    80  			err = common.RunGarbageCollection(gc, cfg.GCWindow, cfg.GCMaxOperationTime)
    81  			if err != nil {
    82  				return err
    83  			}
    84  			log.Ctx(ctx).Info().Msg("Garbage collection completed")
    85  			return nil
    86  		}),
    87  	}
    88  }
    89  
    90  func NewRepairDatastoreCommand(programName string, cfg *datastore.Config) *cobra.Command {
    91  	return &cobra.Command{
    92  		Use:     "repair",
    93  		Short:   "executes datastore repair",
    94  		Long:    "Executes a repair operation for the datastore",
    95  		PreRunE: server.DefaultPreRunE(programName),
    96  		RunE: termination.PublishError(func(cmd *cobra.Command, args []string) error {
    97  			ctx := context.Background()
    98  
    99  			// Disable background GC and hedging.
   100  			cfg.GCInterval = -1 * time.Hour
   101  			cfg.RequestHedgingEnabled = false
   102  
   103  			ds, err := datastore.NewDatastore(ctx, cfg.ToOption())
   104  			if err != nil {
   105  				return fmt.Errorf("failed to create datastore: %w", err)
   106  			}
   107  
   108  			repairable := dspkg.UnwrapAs[dspkg.RepairableDatastore](ds)
   109  			if repairable == nil {
   110  				return fmt.Errorf("datastore of type %T does not support the repair operation", ds)
   111  			}
   112  
   113  			if len(args) == 0 {
   114  				fmt.Println()
   115  				fmt.Println("Available repair operations:")
   116  				for _, op := range repairable.RepairOperations() {
   117  					fmt.Printf("\t%s: %s\n", op.Name, op.Description)
   118  				}
   119  				return nil
   120  			}
   121  
   122  			operationName := args[0]
   123  
   124  			log.Ctx(ctx).Info().Msg("Running repair...")
   125  			err = repairable.Repair(ctx, operationName, true)
   126  			if err != nil {
   127  				return err
   128  			}
   129  
   130  			log.Ctx(ctx).Info().Msg("Datastore repair completed")
   131  			return nil
   132  		}),
   133  	}
   134  }