package controller import ( "bytes" "sort" "strconv" "strings" "time" "github.com/google/btree" "github.com/tidwall/resp" "github.com/tidwall/tile38/controller/collection" "github.com/tidwall/tile38/controller/server" "github.com/tidwall/tile38/geojson" "github.com/tidwall/tile38/geojson/geohash" ) type fvt struct { field string value float64 } type byField []fvt func (a byField) Len() int { return len(a) } func (a byField) Less(i, j int) bool { return a[i].field < a[j].field } func (a byField) Swap(i, j int) { a[i], a[j] = a[j], a[i] } func orderFields(fmap map[string]int, fields []float64) []fvt { var fv fvt fvs := make([]fvt, 0, len(fmap)) for field, idx := range fmap { if idx < len(fields) { fv.field = field fv.value = fields[idx] if fv.value != 0 { fvs = append(fvs, fv) } } } sort.Sort(byField(fvs)) return fvs } func (c *Controller) cmdGet(msg *server.Message) (string, error) { start := time.Now() vs := msg.Values[1:] var ok bool var key, id, typ, sprecision string if vs, key, ok = tokenval(vs); !ok || key == "" { return "", errInvalidNumberOfArguments } if vs, id, ok = tokenval(vs); !ok || id == "" { return "", errInvalidNumberOfArguments } withfields := false if _, peek, ok := tokenval(vs); ok && strings.ToLower(peek) == "withfields" { withfields = true vs = vs[1:] } col := c.getCol(key) if col == nil { if msg.OutputType == server.RESP { return "$-1\r\n", nil } return "", errKeyNotFound } o, fields, ok := col.Get(id) if !ok { if msg.OutputType == server.RESP { return "$-1\r\n", nil } return "", errIDNotFound } vals := make([]resp.Value, 0, 2) var buf bytes.Buffer if msg.OutputType == server.JSON { buf.WriteString(`{"ok":true`) } if vs, typ, ok = tokenval(vs); !ok || strings.ToLower(typ) == "object" { if msg.OutputType == server.JSON { buf.WriteString(`,"object":`) buf.WriteString(o.JSON()) } else { vals = append(vals, resp.StringValue(o.JSON())) } } else { switch strings.ToLower(typ) { default: return "", errInvalidArgument(typ) case "point": point := o.CalculatedPoint() if msg.OutputType == server.JSON { buf.WriteString(`,"point":`) buf.WriteString(point.ExternalJSON()) } else { if point.Z != 0 { vals = append(vals, resp.ArrayValue([]resp.Value{ resp.StringValue(strconv.FormatFloat(point.Y, 'f', -1, 64)), resp.StringValue(strconv.FormatFloat(point.X, 'f', -1, 64)), resp.StringValue(strconv.FormatFloat(point.Z, 'f', -1, 64)), })) } else { vals = append(vals, resp.ArrayValue([]resp.Value{ resp.StringValue(strconv.FormatFloat(point.Y, 'f', -1, 64)), resp.StringValue(strconv.FormatFloat(point.X, 'f', -1, 64)), })) } } case "hash": if vs, sprecision, ok = tokenval(vs); !ok || sprecision == "" { return "", errInvalidNumberOfArguments } if msg.OutputType == server.JSON { buf.WriteString(`,"hash":`) } precision, err := strconv.ParseInt(sprecision, 10, 64) if err != nil || precision < 1 || precision > 64 { return "", errInvalidArgument(sprecision) } p, err := o.Geohash(int(precision)) if err != nil { return "", err } if msg.OutputType == server.JSON { buf.WriteString(`"` + p + `"`) } else { vals = append(vals, resp.StringValue(p)) } case "bounds": bbox := o.CalculatedBBox() if msg.OutputType == server.JSON { buf.WriteString(`,"bounds":`) buf.WriteString(bbox.ExternalJSON()) } else { vals = append(vals, resp.ArrayValue([]resp.Value{ resp.ArrayValue([]resp.Value{ resp.FloatValue(bbox.Min.Y), resp.FloatValue(bbox.Min.X), }), resp.ArrayValue([]resp.Value{ resp.FloatValue(bbox.Max.Y), resp.FloatValue(bbox.Max.X), }), })) } } } if len(vs) != 0 { return "", errInvalidNumberOfArguments } if withfields { fvs := orderFields(col.FieldMap(), fields) if len(fvs) > 0 { fvals := make([]resp.Value, 0, len(fvs)*2) if msg.OutputType == server.JSON { buf.WriteString(`,"fields":{`) } for i, fv := range fvs { if msg.OutputType == server.JSON { if i > 0 { buf.WriteString(`,`) } buf.WriteString(jsonString(fv.field) + ":" + strconv.FormatFloat(fv.value, 'f', -1, 64)) } else { fvals = append(fvals, resp.StringValue(fv.field), resp.StringValue(strconv.FormatFloat(fv.value, 'f', -1, 64))) } i++ } if msg.OutputType == server.JSON { buf.WriteString(`}`) } else { vals = append(vals, resp.ArrayValue(fvals)) } } } switch msg.OutputType { case server.JSON: buf.WriteString(`,"elapsed":"` + time.Now().Sub(start).String() + "\"}") return buf.String(), nil case server.RESP: var oval resp.Value if withfields { oval = resp.ArrayValue(vals) } else { oval = vals[0] } data, err := oval.MarshalRESP() if err != nil { return "", err } return string(data), nil } return "", nil } func (c *Controller) cmdDel(msg *server.Message) (res string, d commandDetailsT, err error) { start := time.Now() vs := msg.Values[1:] var ok bool if vs, d.key, ok = tokenval(vs); !ok || d.key == "" { err = errInvalidNumberOfArguments return } if vs, d.id, ok = tokenval(vs); !ok || d.id == "" { err = errInvalidNumberOfArguments return } if len(vs) != 0 { err = errInvalidNumberOfArguments return } found := false col := c.getCol(d.key) if col != nil { d.obj, d.fields, ok = col.Remove(d.id) if ok { if col.Count() == 0 { c.deleteCol(d.key) d.revert = func() { c.setCol(d.key, col) col.ReplaceOrInsert(d.id, d.obj, nil, d.fields) } } else { d.revert = func() { col.ReplaceOrInsert(d.id, d.obj, nil, d.fields) } } found = true } } d.command = "del" d.updated = found d.timestamp = time.Now() switch msg.OutputType { case server.JSON: res = `{"ok":true,"elapsed":"` + time.Now().Sub(start).String() + "\"}" case server.RESP: if d.updated { res = ":1\r\n" } else { res = ":0\r\n" } } return } func (c *Controller) cmdDrop(msg *server.Message) (res string, d commandDetailsT, err error) { start := time.Now() vs := msg.Values[1:] var ok bool if vs, d.key, ok = tokenval(vs); !ok || d.key == "" { err = errInvalidNumberOfArguments return } if len(vs) != 0 { err = errInvalidNumberOfArguments return } col := c.getCol(d.key) if col != nil { c.deleteCol(d.key) d.revert = func() { c.setCol(d.key, col) } d.updated = true } else { d.key = "" // ignore the details d.updated = false } d.command = "drop" d.timestamp = time.Now() switch msg.OutputType { case server.JSON: res = `{"ok":true,"elapsed":"` + time.Now().Sub(start).String() + "\"}" case server.RESP: if d.updated { res = ":1\r\n" } else { res = ":0\r\n" } } return } func (c *Controller) cmdFlushDB(msg *server.Message) (res string, d commandDetailsT, err error) { start := time.Now() vs := msg.Values[1:] if len(vs) != 0 { err = errInvalidNumberOfArguments return } c.cols = btree.New(16) c.hooks = make(map[string]*Hook) c.hookcols = make(map[string]map[string]*Hook) d.command = "flushdb" d.updated = true d.timestamp = time.Now() switch msg.OutputType { case server.JSON: res = `{"ok":true,"elapsed":"` + time.Now().Sub(start).String() + "\"}" case server.RESP: res = "+OK\r\n" } return } func (c *Controller) parseSetArgs(vs []resp.Value) (d commandDetailsT, fields []string, values []float64, etype string, evs []resp.Value, err error) { var ok bool var typ string if vs, d.key, ok = tokenval(vs); !ok || d.key == "" { err = errInvalidNumberOfArguments return } if vs, d.id, ok = tokenval(vs); !ok || d.id == "" { err = errInvalidNumberOfArguments return } var arg string var nvs []resp.Value fields = make([]string, 0, 8) values = make([]float64, 0, 8) for { if nvs, arg, ok = tokenval(vs); !ok || arg == "" { err = errInvalidNumberOfArguments return } if lc(arg, "field") { vs = nvs var name string var svalue string var value float64 if vs, name, ok = tokenval(vs); !ok || name == "" { err = errInvalidNumberOfArguments return } if isReservedFieldName(name) { err = errInvalidArgument(name) return } if vs, svalue, ok = tokenval(vs); !ok || svalue == "" { err = errInvalidNumberOfArguments return } value, err = strconv.ParseFloat(svalue, 64) if err != nil { err = errInvalidArgument(svalue) return } fields = append(fields, name) values = append(values, value) continue } break } if vs, typ, ok = tokenval(vs); !ok || typ == "" { err = errInvalidNumberOfArguments return } if len(vs) == 0 { err = errInvalidNumberOfArguments return } etype = typ evs = vs switch { default: err = errInvalidArgument(typ) return case lc(typ, "point"): var slat, slon, sz string if vs, slat, ok = tokenval(vs); !ok || slat == "" { err = errInvalidNumberOfArguments return } if vs, slon, ok = tokenval(vs); !ok || slon == "" { err = errInvalidNumberOfArguments return } vs, sz, ok = tokenval(vs) if !ok || sz == "" { var sp geojson.SimplePoint sp.Y, err = strconv.ParseFloat(slat, 64) if err != nil { err = errInvalidArgument(slat) return } sp.X, err = strconv.ParseFloat(slon, 64) if err != nil { err = errInvalidArgument(slon) return } d.obj = sp } else { var sp geojson.Point sp.Coordinates.Y, err = strconv.ParseFloat(slat, 64) if err != nil { err = errInvalidArgument(slat) return } sp.Coordinates.X, err = strconv.ParseFloat(slon, 64) if err != nil { err = errInvalidArgument(slon) return } sp.Coordinates.Z, err = strconv.ParseFloat(sz, 64) if err != nil { err = errInvalidArgument(sz) return } d.obj = sp } case lc(typ, "bounds"): var sminlat, sminlon, smaxlat, smaxlon string if vs, sminlat, ok = tokenval(vs); !ok || sminlat == "" { err = errInvalidNumberOfArguments return } if vs, sminlon, ok = tokenval(vs); !ok || sminlon == "" { err = errInvalidNumberOfArguments return } if vs, smaxlat, ok = tokenval(vs); !ok || smaxlat == "" { err = errInvalidNumberOfArguments return } if vs, smaxlon, ok = tokenval(vs); !ok || smaxlon == "" { err = errInvalidNumberOfArguments return } var minlat, minlon, maxlat, maxlon float64 minlat, err = strconv.ParseFloat(sminlat, 64) if err != nil { err = errInvalidArgument(sminlat) return } minlon, err = strconv.ParseFloat(sminlon, 64) if err != nil { err = errInvalidArgument(sminlon) return } maxlat, err = strconv.ParseFloat(smaxlat, 64) if err != nil { err = errInvalidArgument(smaxlat) return } maxlon, err = strconv.ParseFloat(smaxlon, 64) if err != nil { err = errInvalidArgument(smaxlon) return } g := geojson.Polygon{ Coordinates: [][]geojson.Position{ { {X: minlon, Y: minlat, Z: 0}, {X: minlon, Y: maxlat, Z: 0}, {X: maxlon, Y: maxlat, Z: 0}, {X: maxlon, Y: minlat, Z: 0}, {X: minlon, Y: minlat, Z: 0}, }, }, } d.obj = g case lc(typ, "hash"): var sp geojson.SimplePoint var shash string if vs, shash, ok = tokenval(vs); !ok || shash == "" { err = errInvalidNumberOfArguments return } var lat, lon float64 lat, lon, err = geohash.Decode(shash) if err != nil { return } sp.X = lon sp.Y = lat d.obj = sp case lc(typ, "object"): var object string if vs, object, ok = tokenval(vs); !ok || object == "" { err = errInvalidNumberOfArguments return } d.obj, err = geojson.ObjectJSON(object) if err != nil { return } } if len(vs) != 0 { err = errInvalidNumberOfArguments } return } func (c *Controller) cmdSet(msg *server.Message) (res string, d commandDetailsT, err error) { if c.config.MaxMemory > 0 && c.outOfMemory { err = errOOM return } start := time.Now() vs := msg.Values[1:] var fields []string var values []float64 d, fields, values, _, _, err = c.parseSetArgs(vs) if err != nil { return } addedcol := false col := c.getCol(d.key) if col == nil { col = collection.New() c.setCol(d.key, col) addedcol = true } d.oldObj, d.oldFields, d.fields = col.ReplaceOrInsert(d.id, d.obj, fields, values) d.revert = func() { if addedcol { c.deleteCol(d.key) } else if d.oldObj != nil { col.ReplaceOrInsert(d.id, d.oldObj, nil, d.oldFields) } else { col.Remove(d.id) } } d.command = "set" d.updated = true // perhaps we should do a diff on the previous object? fmap := col.FieldMap() d.fmap = make(map[string]int) for key, idx := range fmap { d.fmap[key] = idx } d.timestamp = time.Now() switch msg.OutputType { case server.JSON: res = `{"ok":true,"elapsed":"` + time.Now().Sub(start).String() + "\"}" case server.RESP: res = "+OK\r\n" } return } func (c *Controller) parseFSetArgs(vs []resp.Value) (d commandDetailsT, err error) { var svalue string var ok bool if vs, d.key, ok = tokenval(vs); !ok || d.key == "" { err = errInvalidNumberOfArguments return } if vs, d.id, ok = tokenval(vs); !ok || d.id == "" { err = errInvalidNumberOfArguments return } if vs, d.field, ok = tokenval(vs); !ok || d.field == "" { err = errInvalidNumberOfArguments return } if isReservedFieldName(d.field) { err = errInvalidNumberOfArguments return } if vs, svalue, ok = tokenval(vs); !ok || svalue == "" { err = errInvalidNumberOfArguments return } if len(vs) != 0 { err = errInvalidNumberOfArguments return } d.value, err = strconv.ParseFloat(svalue, 64) if err != nil { err = errInvalidArgument(svalue) return } return } func (c *Controller) cmdFset(msg *server.Message) (res string, d commandDetailsT, err error) { start := time.Now() vs := msg.Values[1:] d, err = c.parseFSetArgs(vs) col := c.getCol(d.key) if col == nil { err = errKeyNotFound return } var ok bool d.obj, d.fields, d.updated, ok = col.SetField(d.id, d.field, d.value) if !ok { err = errIDNotFound return } d.command = "fset" d.timestamp = time.Now() fmap := col.FieldMap() d.fmap = make(map[string]int) for key, idx := range fmap { d.fmap[key] = idx } switch msg.OutputType { case server.JSON: res = `{"ok":true,"elapsed":"` + time.Now().Sub(start).String() + "\"}" case server.RESP: if d.updated { res = ":1\r\n" } else { res = ":0\r\n" } } return }