kv/main.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
		}
		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)
}