func init() { go func() { for { archiveJob := <-archive_chan filename := fmt.Sprintf("./db/archive/%s.omg", archiveJob.key) f, err := os.Create(filename) if err != nil { panic(err) } if _, err = f.Write(archiveJob.data); err != nil { panic(err) } f.Close() } }() store.RegisterStoreType(ArchiveType) store.DefaultDBManager.AddFuncWithSideEffects("archive", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], data.Any) if ok { if elem.EntryType != ArchiveType { str_rep, ok := db.StoreDump(args[0]) if ok { archive_chan <- ArchiveJob{args[0], str_rep} db.StoreDel(args[0]) db.ArchiveRecord[args[0]] = 0 return reply.IntReply(1) } else { return reply.NilReply } } else { return reply.IntReply(0) } } else { return reply.NilReply } }) store.DefaultDBManager.AddFuncWithSideEffects("unarchive", func(db *store.DB, args []string) string { _, ok, _ := db.StoreGet(args[0], ArchiveType) if ok { filename := fmt.Sprintf("./db/archive/%s.omg", args[0]) data, err := ioutil.ReadFile(filename) if err != nil { panic(err) } delete(db.ArchiveRecord, args[0]) db.StoreLoad(args[0], data) return reply.OKReply } else { return reply.NilReply } }) }
func init() { store.RegisterStoreType(StringType) store.DefaultDBManager.AddFuncWithSideEffects("append", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) if ok { str, _ := elem.Value.(string) str += args[1] elem.Value = str return reply.IntReply(len(str)) } else { db.StoreSet(args[0], &data.Entry{args[1], StringType}) return reply.IntReply(len(args[1])) } }) bitCount := []uint8{0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7, 4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8} store.DefaultDBManager.AddFuncWithSideEffects("bitcount", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) if ok { byte_arr, _ := elem.Value.([]byte) c := 0 for _, c := range byte_arr { c += bitCount[c] } return reply.IntReply(c) } else { return reply.IntReply(0) } }) store.DefaultDBManager.AddFuncWithSideEffects("decr", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) var num int if ok { str, _ := elem.Value.(string) num, err := strconv.Atoi(str) if err != nil { panic(err) } num -= 1 elem.Value = fmt.Sprintf("%d", num) } else { num = -1 db.StoreSet(args[0], &data.Entry{"-1", StringType}) } return reply.IntReply(num) }) store.DefaultDBManager.AddFuncWithSideEffects("decrby", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) var num int decrby, err := strconv.Atoi(args[1]) if err != nil { panic(err) } if ok { str, _ := elem.Value.(string) num, err := strconv.Atoi(str) if err != nil { panic(err) } num -= decrby elem.Value = fmt.Sprintf("%d", num) } else { str := fmt.Sprintf("-%d", decrby) db.StoreSet(args[0], &data.Entry{str, StringType}) } println(num) return reply.IntReply(num) }) store.DefaultDBManager.AddFuncWithSideEffects("incr", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) var num int if ok { str, _ := elem.Value.(string) num, err := strconv.Atoi(str) if err != nil { panic(err) } num += 1 elem.Value = fmt.Sprintf("%d", num) } else { num = 1 db.StoreSet(args[0], &data.Entry{"-1", StringType}) } return reply.IntReply(num) }) store.DefaultDBManager.AddFuncWithSideEffects("incrby", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) var num int decrby, err := strconv.Atoi(args[1]) if err != nil { panic(err) } if ok { str, _ := elem.Value.(string) num, err := strconv.Atoi(str) if err != nil { panic(err) } num += decrby elem.Value = fmt.Sprintf("%d", num) } else { str := fmt.Sprintf("%d", decrby) db.StoreSet(args[0], &data.Entry{str, StringType}) } println(num) return reply.IntReply(num) }) store.DefaultDBManager.AddFunc("get", func(db *store.DB, args []string) string { elem, ok, err := db.StoreGet(args[0], StringType) if err != nil { return reply.ErrorReply(err) } if ok { str := elem.GetString() return reply.BulkReply(str) } else { return reply.NilReply } }) store.DefaultDBManager.AddFunc("getset", func(db *store.DB, args []string) string { elem, ok, err := db.StoreGet(args[0], StringType) if err != nil { return reply.ErrorReply(err) } if ok { str := elem.GetString() elem.Value = args[1] return reply.BulkReply(str) } else { db.StoreSet(args[0], &data.Entry{args[1], StringType}) return reply.NilReply } }) store.DefaultDBManager.AddFuncWithSideEffects("psetex", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) millis, err := strconv.Atoi(args[1]) if err != nil { panic(err) } db.ExpiryRecord[args[0]] = store.Milliseconds() + uint64(millis) if ok { elem.Value = args[2] } else { db.StoreSet(args[0], &data.Entry{args[2], StringType}) } return reply.OKReply }) store.DefaultDBManager.AddFuncWithSideEffects("setex", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) seconds, err := strconv.Atoi(args[1]) if err != nil { panic(err) } db.ExpiryRecord[args[0]] = store.Milliseconds() + uint64(seconds)*1000 if ok { elem.Value = args[2] } else { db.StoreSet(args[0], &data.Entry{args[2], StringType}) } return reply.OKReply }) store.DefaultDBManager.AddFuncWithSideEffects("set", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], StringType) if ok { elem.Value = args[1] } else { db.StoreSet(args[0], &data.Entry{args[1], StringType}) } return reply.OKReply }) store.DefaultDBManager.AddFuncWithSideEffects("setnx", func(db *store.DB, args []string) string { s := args[0] _, ok, _ := db.StoreGet(s, StringType) if ok { return reply.IntReply(0) } db.StoreSet(s, &data.Entry{args[1], StringType}) return reply.IntReply(1) }) store.DefaultDBManager.AddFunc("strlen", func(db *store.DB, args []string) string { elem, ok, err := db.StoreGet(args[0], StringType) if err != nil { return reply.ErrorReply(err) } if ok { str := elem.GetString() return reply.IntReply(len(str)) } else { return reply.IntReply(0) } }) }
func init() { store.RegisterStoreType(ZSetType) store.DefaultDBManager.AddFuncWithSideEffects("zadd", func(db *store.DB, args []string) string { i, err := strconv.Atoi(args[1]) if err != nil { return "-ERR Invalid score\r\n" } elem, ok, err := db.StoreGet(args[0], ZSetType) if !ok { zs := zset.New() elem = &data.Entry{zs, ZSetType} db.StoreSet(args[0], elem) } zs, _ := elem.Value.(*zset.ZSet) zs.Insert(i, args[2]) return ":1\r\n" }) store.DefaultDBManager.AddFuncWithSideEffects("zrem", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], ZSetType) if !ok { return "$-1\r\n" } zs, _ := elem.Value.(*zset.ZSet) zs.Delete(args[1]) return ":1\r\n" }) store.DefaultDBManager.AddFunc("zcard", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], ZSetType) if !ok { return "$-1\r\n" } zs, ok := elem.Value.(*zset.ZSet) return fmt.Sprintf(":%d\r\n", zs.Card()) }) store.DefaultDBManager.AddFunc("zscore", func(db *store.DB, args []string) string { elem, ok, _ := db.StoreGet(args[0], ZSetType) if !ok { return "$-1\r\n" } zs, _ := elem.Value.(*zset.ZSet) score, ok := zs.Score(args[1]) if ok { return fmt.Sprintf(":%d\r\n", score) } else { return "$-1\r\n" } return "$-1\r\n" }) //store.DefaultDBManager.AddFunc("zrank", func (db *store.DB, args []string) string { // elem, ok, _ := db.StoreGet(args[0],ZSetType) // if !ok { // return "$-1\r\n" // } // zs, _ := elem.Value.(*zset.ZSet) // rank, ok := zs.Rank(args[1]) // if ok { // return fmt.Sprintf(":%d\r\n",rank) // } else { // return "$-1\r\n" // } // return "$-1\r\n" //}) }
func init() { gob.Register(&Post{}) gob.Register(&Feed{}) store.RegisterStoreType(FeedType) store.DefaultDBManager.AddFuncWithSideEffects("feedcap", func(db *store.DB, args []string) string { s := args[0] cap, err := strconv.Atoi(args[1]) if err != nil { return reply.ErrorReply(err) } e, ok, _ := db.StoreGet(s, FeedType) if ok { if f, ok := e.Value.(*Feed); ok { f.Capacity = cap return reply.OKReply } else { return reply.ErrorReply("Type Mismatch") } } else { f := &Feed{s, cap, 0, make([]*Post, 0)} db.StoreSet(s, &data.Entry{f, FeedType}) } return reply.OKReply }) store.DefaultDBManager.AddFunc("feedget", func(db *store.DB, args []string) string { if e, ok, _ := db.StoreGet(args[0], FeedType); ok { if f, ok := e.Value.(*Feed); ok { cap_str := fmt.Sprintf("%d", f.Capacity) var rep reply.MultiBulkWriter rep.WriteCount(2) rep.WriteString(f.Name) rep.WriteString(cap_str) return rep.String() } else { return "-ERR Type mismatch\r\n" } } else { return reply.NilReply } }) store.DefaultDBManager.AddFuncWithSideEffects("feedpost", func(db *store.DB, args []string) string { e, ok, _ := db.StoreGet(args[0], FeedType) if !ok { s := args[0] f := &Feed{s, -1, 0, make([]*Post, 0)} e = &data.Entry{f, FeedType} db.StoreSet(s, e) } if f, ok := e.Value.(*Feed); ok { timestamp := store.Milliseconds() if timestamp == f.LastTimestamp { timestamp += 1 } f.LastTimestamp = timestamp p := &Post{args[1], args[2], timestamp, 0, make(map[string]string)} f.Posts = append(f.Posts, p) return reply.OKReply } else { return reply.ErrorReply("Type mismatch") } }) store.DefaultDBManager.AddFunc("feedposts", func(db *store.DB, args []string) string { if e, ok, _ := db.StoreGet(args[0], FeedType); ok { if f, ok := e.Value.(*Feed); ok { post_count := len(f.Posts) offset := 0 count := post_count if len(args) == 3 { var err error offset, err = strconv.Atoi(args[1]) if err != nil { panic(err) } count, err = strconv.Atoi(args[2]) if err != nil { panic(err) } } if offset >= post_count { return reply.NilReply } if offset+count > post_count { count = post_count - offset } var rep reply.MultiBulkWriter rep.WriteCount(3 * post_count) for i := post_count - offset; i >= post_count-offset-count; i-- { p := f.Posts[i] timestamp_str := fmt.Sprintf("%d", p.Timestamp) rep.WriteString(p.Author) rep.WriteString(p.Body) rep.WriteString(timestamp_str) } return rep.String() } else { return reply.ErrorReply("Type mismatch") } } else { return reply.NilReply } }) store.DefaultDBManager.AddFunc("feedlen", func(db *store.DB, args []string) string { if e, ok, _ := db.StoreGet(args[0], FeedType); ok { if f, ok := e.Value.(*Feed); ok { return reply.IntReply(len(f.Posts)) } } return reply.NilReply }) //store.DefaultDBManager.AddFuncWithSideEffects("feedlike", func (db *store.DB, args []string) string { //}) }
func init() { gob.Register(&User{}) store.RegisterStoreType(UserType) store.DefaultDBManager.AddFuncWithSideEffects("useradd", func(db *store.DB, args []string) string { s := args[0] _, ok, _ := db.StoreGet(s, UserType) if ok { return reply.ErrorReply("Key already exists") } else { raw_hash, _ := scrypt.Key([]byte(args[1]), []byte(args[0]), 16384, 8, 1, 32) b64_hash := store.Base64Encode(raw_hash) final_hash := fmt.Sprintf("$%d$%d$%d$%s", 16384, 8, 1, string(b64_hash)) user := &User{args[0], final_hash, "", 0} db.StoreSet(args[0], &data.Entry{user, UserType}) return reply.OKReply } return reply.OKReply }) store.DefaultDBManager.AddFunc("userget", func(db *store.DB, args []string) string { s := args[0] e, ok, _ := db.StoreGet(s, UserType) if ok { if user, ok := e.Value.(*User); ok { score_str := fmt.Sprintf("%d", user.Score) var rep reply.MultiBulkWriter rep.WriteCount(2) rep.WriteString(user.Username) rep.WriteString(score_str) return rep.String() } else { return reply.ErrorReply("Type Mismatch") } } else { return reply.NilReply } }) store.DefaultDBManager.AddFuncWithSideEffects("userauth", func(db *store.DB, args []string) string { username := args[0] e, ok, _ := db.StoreGet(username, UserType) if ok { if user, ok := e.Value.(*User); ok { password := args[1] raw_hash, _ := scrypt.Key([]byte(password), []byte(username), 16384, 8, 1, 32) b64_hash := store.Base64Encode(raw_hash) final_hash := fmt.Sprintf("$%d$%d$%d$%s", 16384, 8, 1, string(b64_hash)) if final_hash == user.Hash { c := 64 b := make([]byte, c) n, err := io.ReadFull(rand.Reader, b) if n != len(b) || err != nil { return reply.ErrorReply("Unable to generate a secure random token") } token := store.Base64Encode(b) token_str := string(token) db.StoreSet(token_str, &data.Entry{username, 0}) if user.Token != "" { db.StoreDel(token_str) } user.Token = token_str return reply.BulkReply(token_str) } else { return reply.ErrorReply("Password is incorrect") } } else { return reply.ErrorReply("How did you manage to mess things up this good?") } } else { return reply.ErrorReply("No user exists with that username") } }) store.DefaultDBManager.AddFunc("usertoken", func(db *store.DB, args []string) string { token := args[0] e, ok, _ := db.StoreGet(token, strings.StringType) if ok { str, _ := e.Value.(string) return reply.BulkReply(str) } else { return reply.ErrorReply("Invalid or expired token") } }) }
func init() { gob.Register(make(map[string]bool)) store.RegisterStoreType(SetType) store.DefaultDBManager.AddFuncWithSideEffects("sadd", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil { return reply.ErrorReply("Type Mismatch") } var setData map[string]bool if !ok { setData = make(map[string]bool) db.StoreSet(args[0], &data.Entry{setData, SetType}) } else { setData, _ = e.Value.(map[string]bool) } l := len(args) sdata := args[1:l] count := 0 for _, member := range sdata { _, ok = setData[member] if !ok { setData[member] = true count += 1 } } return reply.IntReply(count) }) store.DefaultDBManager.AddFuncWithSideEffects("srem", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil { return reply.ErrorReply("Type Mismatch") } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { setData = make(map[string]bool) db.StoreSet(args[0], &data.Entry{setData, SetType}) } l := len(args) sdata := args[1:l] count := 0 for _, member := range sdata { _, ok = setData[member] if ok { delete(setData, member) count += 1 } } return reply.IntReply(count) }) store.DefaultDBManager.AddFuncWithSideEffects("srem", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil { return reply.ErrorReply("Type Mismatch") } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { setData = make(map[string]bool) } l := len(args) sdata := args[1:l] count := 0 for _, member := range sdata { _, ok = setData[member] if ok { delete(setData, member) count += 1 } } return reply.IntReply(count) }) store.DefaultDBManager.AddFunc("sismember", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil || !ok { return "-ERR Type Mismatch\r\n" } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { return reply.IntReply(0) } _, ok = setData[args[1]] if ok { return reply.IntReply(1) } return reply.IntReply(0) }) store.DefaultDBManager.AddFunc("scard", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil || !ok { return reply.ErrorReply("Type Mismatch") } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if ok { return reply.IntReply(len(setData)) } return reply.IntReply(0) }) store.DefaultDBManager.AddFuncWithSideEffects("spop", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil || !ok { return "-ERR Type Mismatch\r\n" } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { return reply.NilReply } for k, _ := range setData { delete(setData, k) return reply.BulkReply(k) } return reply.NilReply }) store.DefaultDBManager.AddFuncWithSideEffects("smembers", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil || !ok { return "-ERR Type Mismatch\r\n" } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { return reply.NilReply } var w reply.MultiBulkWriter w.WriteCount(len(setData)) for k, _ := range setData { w.WriteString(k) } return w.String() }) store.DefaultDBManager.AddFuncWithSideEffects("srandmember", func(db *store.DB, args []string) string { e, ok, err := db.StoreGet(args[0], SetType) if err != nil || !ok { return "-ERR Type Mismatch\r\n" } if !ok { return reply.NilReply } setData, ok := e.Value.(map[string]bool) if !ok { return reply.NilReply } if len(args) > 1 { count, err := strconv.Atoi(args[1]) if err != nil { panic(err) } var w reply.MultiBulkWriter w.WriteCount(count) for k, _ := range setData { w.WriteString(k) count -= 1 if count == 0 { break } } return w.String() } else { for k, _ := range setData { return reply.BulkReply(k) } } return reply.NilReply }) }