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