kv/main.go

301 lines
6.4 KiB
Go

package main
import (
"context"
"encoding/json"
"errors"
"flag"
"fmt"
"io/fs"
"os"
"path/filepath"
"sort"
"strings"
"github.com/adrg/xdg"
"github.com/peterbourgon/ff/v3"
"github.com/peterbourgon/ff/v3/ffcli"
)
var Version = "develop"
func main() {
configHome := xdg.ConfigHome
if configHome == "" {
home, err := os.UserHomeDir()
if err != nil {
home = "."
}
configHome = home
}
defaultStore := filepath.Join(configHome, "kv")
fs := flag.NewFlagSet("kv", flag.ContinueOnError)
storeFlag := fs.String("store", defaultStore, "Location of the store directory")
fs.StringVar(storeFlag, "s", *storeFlag, "--store")
a := &app{
storeLocation: storeFlag,
}
var app *ffcli.Command
app = &ffcli.Command{
Name: "kv",
ShortUsage: fmt.Sprintf("Key/Value store - version %s", Version),
ShortHelp: "kv [get|set|del|list] ...",
FlagSet: fs,
Subcommands: []*ffcli.Command{
a.get(),
a.set(),
a.del(),
a.list(),
},
Options: []ff.Option{
ff.WithEnvVarPrefix("KV"),
},
Exec: func(_ context.Context, _ []string) error {
return errors.New(app.UsageFunc(app))
},
}
if err := app.ParseAndRun(context.Background(), os.Args[1:]); err != nil {
if errors.Is(err, flag.ErrHelp) {
return
}
fmt.Println(err)
}
}
type app struct {
storeLocation *string
}
func (a *app) get() *ffcli.Command {
fs := flag.NewFlagSet("get", flag.ContinueOnError)
storeFlag := fs.String("store", "", "Use a specific store instead of aggregate")
fs.StringVar(storeFlag, "s", *storeFlag, "--store")
return &ffcli.Command{
Name: "get",
FlagSet: fs,
Exec: func(ctx context.Context, args []string) error {
if len(args) < 1 {
return errors.New("get requires at least one argument")
}
data, err := a.load()
if err != nil {
return err
}
key := strings.ToLower(strings.Join(args, " "))
val, ok := data.Val(*storeFlag, key)
if !ok {
return fmt.Errorf("no value found for %q", key)
}
fmt.Print(val)
return nil
},
}
}
func (a *app) set() *ffcli.Command {
fs := flag.NewFlagSet("set", flag.ContinueOnError)
storeFlag := fs.String("store", "kv", "Store for this key/value")
fs.StringVar(storeFlag, "s", *storeFlag, "--store")
return &ffcli.Command{
Name: "set",
FlagSet: fs,
Exec: func(ctx context.Context, args []string) error {
if len(args) < 2 {
return errors.New("set requires at least two arguments")
}
data, err := a.load()
if err != nil {
return err
}
store, ok := data[*storeFlag]
if !ok {
fi, err := os.Create(filepath.Join(*a.storeLocation, *storeFlag+".json"))
if err != nil {
return err
}
if _, err := fi.WriteString("{}"); err != nil {
return err
}
if err := fi.Close(); err != nil {
return err
}
store = make(map[string]string)
}
key := args[0]
store[key] = strings.Join(args[1:], " ")
if err := a.save(*storeFlag, store); err != nil {
return err
}
fmt.Printf("set %q\n", key)
return nil
},
}
}
func (a *app) del() *ffcli.Command {
fs := flag.NewFlagSet("del", flag.ContinueOnError)
storeFlag := fs.String("store", "kv", "Store for this key/value")
fs.StringVar(storeFlag, "s", *storeFlag, "--store")
return &ffcli.Command{
Name: "del",
FlagSet: fs,
Exec: func(ctx context.Context, args []string) error {
if len(args) < 1 {
return errors.New("del requires at least 1 argument")
}
data, err := a.load()
if err != nil {
return err
}
store, ok := data[*storeFlag]
if !ok {
return fmt.Errorf("no store found for %q", *storeFlag)
}
key := strings.ToLower(strings.Join(args, " "))
if _, ok := store[key]; !ok {
return fmt.Errorf("no value found for %q", key)
}
delete(store, key)
if err := a.save(*storeFlag, store); err != nil {
return err
}
fmt.Printf("deleted %q\n", key)
return nil
},
}
}
func (a *app) list() *ffcli.Command {
fs := flag.NewFlagSet("list", flag.ContinueOnError)
storeFlag := fs.String("store", "", "Use a specific store instead of aggregate")
fs.StringVar(storeFlag, "s", *storeFlag, "--store")
return &ffcli.Command{
Name: "list",
FlagSet: fs,
Exec: func(ctx context.Context, args []string) error {
data, err := a.load()
if err != nil {
return err
}
var keys []string
prefix := strings.ToLower(strings.Join(args, " "))
for key := range data.Map() {
if strings.HasPrefix(key, prefix) {
keys = append(keys, key)
}
}
sort.Strings(keys)
for _, key := range keys {
fmt.Println(key)
}
return nil
},
}
}
type store map[string]map[string]string
func (s store) Val(store, key string) (string, bool) {
if store != "" {
val, ok := s[store][key]
return val, ok
}
val, ok := s.Map()[key]
return val, ok
}
func (s store) Map() map[string]string {
m := s["kv"]
for sname, ss := range s {
if sname == "kv" {
continue
}
for k, v := range ss {
m[k] = v
}
}
return m
}
func load(paths ...string) store {
s := make(store)
for _, p := range paths {
func() {
var m map[string]string
fi, err := os.Open(p)
if err != nil {
fmt.Printf("could not open %q: %v\n", p, err)
return
}
defer fi.Close()
if err := json.NewDecoder(fi).Decode(&m); err != nil {
fmt.Printf("could not decode %q: %v\n", p, err)
return
}
storeName := strings.TrimSuffix(filepath.Base(p), ".json")
s[storeName] = m
}()
}
return s
}
func (a *app) load() (store, error) {
defaultConfig := filepath.Join(*a.storeLocation, "kv.json")
if _, err := os.Stat(defaultConfig); err != nil {
if !errors.Is(err, fs.ErrNotExist) {
return nil, err
}
if err := os.MkdirAll(*a.storeLocation, os.ModePerm); err != nil {
return nil, err
}
fi, err := os.Create(defaultConfig)
if err != nil {
return nil, err
}
if _, err := fi.WriteString("{}"); err != nil {
return nil, err
}
if err := fi.Close(); err != nil {
return nil, err
}
}
dirs, err := os.ReadDir(*a.storeLocation)
if err != nil {
return nil, err
}
paths := make([]string, 0, len(dirs))
for _, dir := range dirs {
paths = append(paths, filepath.Join(*a.storeLocation, dir.Name()))
}
return load(paths...), nil
}
func (a *app) save(storeName string, m map[string]string) error {
fi, err := os.Create(filepath.Join(*a.storeLocation, storeName+".json"))
if err != nil {
return err
}
defer fi.Close()
enc := json.NewEncoder(fi)
enc.SetIndent("", "\t")
return enc.Encode(&m)
}