package database import ( "fmt" "github.com/strimertul/strimertul/modules" jsoniter "github.com/json-iterator/go" kv "github.com/strimertul/kilovolt/v8" "go.uber.org/zap" ) var json = jsoniter.ConfigFastest var ( // ErrUnknown is returned when a response is received that doesn't match any expected outcome. ErrUnknown = fmt.Errorf("unknown error") ) type DBModule struct { client *kv.LocalClient hub *kv.Hub logger *zap.Logger } type KvPair struct { Key string Data string } func NewDBModule(hub *kv.Hub, manager *modules.Manager) (*DBModule, error) { logger := manager.Logger(modules.ModuleDB) localClient := kv.NewLocalClient(kv.ClientOptions{}, logger) go localClient.Run() hub.AddClient(localClient) localClient.Wait() err := hub.SetAuthenticated(localClient.UID(), true) if err != nil { return nil, err } module := &DBModule{ client: localClient, hub: hub, logger: logger, } manager.Modules[modules.ModuleDB] = module return module, nil } func (mod *DBModule) Hub() *kv.Hub { return mod.hub } func (mod *DBModule) Status() modules.ModuleStatus { return modules.ModuleStatus{ Enabled: mod.hub != nil, Working: mod.client != nil, StatusString: "ok", } } func (mod *DBModule) Close() error { mod.hub.RemoveClient(mod.client) return nil } func (mod *DBModule) GetKey(key string) (string, error) { res, err := mod.makeRequest(kv.CmdReadKey, map[string]interface{}{"key": key}) if err != nil { return "", err } return res.Data.(string), nil } func (mod *DBModule) PutKey(key string, data string) error { _, err := mod.makeRequest(kv.CmdWriteKey, map[string]interface{}{"key": key, "data": data}) return err } func (mod *DBModule) Subscribe(fn kv.SubscriptionCallback, prefixes ...string) error { for _, prefix := range prefixes { _, err := mod.makeRequest(kv.CmdSubscribePrefix, map[string]interface{}{"prefix": prefix}) if err != nil { return err } mod.client.SetPrefixSubCallback(prefix, fn) } return nil } func (mod *DBModule) GetJSON(key string, dst interface{}) error { res, err := mod.GetKey(key) if err != nil { return err } return json.Unmarshal([]byte(res), dst) } func (mod *DBModule) GetAll(prefix string) (map[string]string, error) { res, err := mod.makeRequest(kv.CmdReadPrefix, map[string]interface{}{"prefix": prefix}) if err != nil { return nil, err } out := make(map[string]string) for key, value := range res.Data.(map[string]interface{}) { out[key] = value.(string) } return out, nil } func (mod *DBModule) PutJSON(key string, data interface{}) error { byt, err := json.Marshal(data) if err != nil { return err } return mod.PutKey(key, string(byt)) } func (mod *DBModule) PutJSONBulk(kvs map[string]interface{}) error { encoded := make(map[string]interface{}) for k, v := range kvs { byt, err := json.Marshal(v) if err != nil { return err } encoded[k] = string(byt) } _, chn := mod.client.MakeRequest(kv.CmdWriteBulk, encoded) _, err := getResponse(<-chn) return err } func (mod *DBModule) RemoveKey(key string) error { // TODO return mod.PutKey(key, "") } func (mod *DBModule) makeRequest(cmd string, data map[string]interface{}) (kv.Response, error) { req, chn := mod.client.MakeRequest(cmd, data) mod.hub.SendMessage(req) return getResponse(<-chn) } func getResponse(response interface{}) (kv.Response, error) { switch c := response.(type) { case kv.Response: return c, nil case kv.Error: return kv.Response{}, &KvError{c} } return kv.Response{}, ErrUnknown } type KvError struct { ErrorData kv.Error } func (kv *KvError) Error() string { return fmt.Sprintf("%s: %s", kv.ErrorData.Error, kv.ErrorData.Details) }