tile38/controller/scanner.go

package controller

import (
	"bytes"
	"errors"
	"strconv"
	"sync"

	"github.com/tidwall/resp"
	"github.com/tidwall/tile38/controller/collection"
	"github.com/tidwall/tile38/controller/glob"
	"github.com/tidwall/tile38/controller/server"
	"github.com/tidwall/tile38/geojson"
)

const limitItems = 100
const capLimit = 100000

type outputT int

const (
	outputUnknown outputT = iota
	outputIDs
	outputObjects
	outputCount
	outputPoints
	outputHashes
	outputBounds
)

type scanWriter struct {
	mu             sync.Mutex
	c              *Controller
	wr             *bytes.Buffer
	msg            *server.Message
	col            *collection.Collection
	fmap           map[string]int
	farr           []string
	fvals          []float64
	output         outputT
	wheres         []whereT
	numberItems    uint64
	nofields       bool
	limit          uint64
	hitLimit       bool
	once           bool
	count          uint64
	precision      uint64
	glob           string
	globEverything bool
	globSingle     bool
	fullFields     bool
	values         []resp.Value
}

func (c *Controller) newScanWriter(
	wr *bytes.Buffer, msg *server.Message, key string, output outputT, precision uint64, globPattern string, limit uint64, wheres []whereT, nofields bool,
) (
	*scanWriter, error,
) {
	if limit == 0 {
		limit = limitItems
	} else if limit > capLimit {
		limit = capLimit
	}
	switch output {
	default:
		return nil, errors.New("invalid output type")
	case outputIDs, outputObjects, outputCount, outputBounds, outputPoints, outputHashes:
	}
	sw := &scanWriter{
		c:         c,
		wr:        wr,
		msg:       msg,
		output:    output,
		wheres:    wheres,
		precision: precision,
		nofields:  nofields,
		glob:      globPattern,
		limit:     limit,
	}
	if globPattern == "*" || globPattern == "" {
		sw.globEverything = true
	} else {
		if !glob.IsGlob(globPattern) {
			sw.globSingle = true
		}
	}
	sw.col = c.getCol(key)
	if sw.col != nil {
		sw.fmap = sw.col.FieldMap()
		sw.farr = sw.col.FieldArr()
	}
	sw.fvals = make([]float64, len(sw.farr))
	return sw, nil
}

func (sw *scanWriter) hasFieldsOutput() bool {
	switch sw.output {
	default:
		return false
	case outputObjects, outputPoints, outputHashes, outputBounds:
		return !sw.nofields
	}
}

func (sw *scanWriter) writeHead() {
	sw.mu.Lock()
	defer sw.mu.Unlock()
	switch sw.msg.OutputType {
	case server.JSON:
		if len(sw.farr) > 0 && sw.hasFieldsOutput() {
			sw.wr.WriteString(`,"fields":[`)
			for i, field := range sw.farr {
				if i > 0 {
					sw.wr.WriteByte(',')
				}
				sw.wr.WriteString(jsonString(field))
			}
			sw.wr.WriteByte(']')
		}
		switch sw.output {
		case outputIDs:
			sw.wr.WriteString(`,"ids":[`)
		case outputObjects:
			sw.wr.WriteString(`,"objects":[`)
		case outputPoints:
			sw.wr.WriteString(`,"points":[`)
		case outputBounds:
			sw.wr.WriteString(`,"bounds":[`)
		case outputHashes:
			sw.wr.WriteString(`,"hashes":[`)
		case outputCount:

		}
	case server.RESP:
	}
}

func (sw *scanWriter) writeFoot(cursor uint64) {
	sw.mu.Lock()
	defer sw.mu.Unlock()
	if !sw.hitLimit {
		cursor = 0
	}
	switch sw.msg.OutputType {
	case server.JSON:
		switch sw.output {
		default:
			sw.wr.WriteByte(']')
		case outputCount:

		}
		sw.wr.WriteString(`,"count":` + strconv.FormatUint(sw.count, 10))
		sw.wr.WriteString(`,"cursor":` + strconv.FormatUint(cursor, 10))
	case server.RESP:
		sw.wr.Reset()
		var data []byte
		var err error
		if sw.output == outputCount {
			data, err = resp.IntegerValue(int(sw.count)).MarshalRESP()
		} else {
			values := []resp.Value{
				resp.IntegerValue(int(cursor)),
				resp.ArrayValue(sw.values),
			}
			data, err = resp.ArrayValue(values).MarshalRESP()
		}
		if err != nil {
			panic("Eek this is bad. Marshal resp should not fail.")
		}
		sw.wr.Write(data)
	}
}

func (sw *scanWriter) fieldMatch(fields []float64, o geojson.Object) ([]float64, bool) {
	var z float64
	var gotz bool
	if !sw.hasFieldsOutput() || sw.fullFields {
		for _, where := range sw.wheres {
			if where.field == "z" {
				if !gotz {
					z = o.CalculatedPoint().Z
				}
				if !where.match(z) {
					return sw.fvals, false
				}
				continue
			}
			var value float64
			idx, ok := sw.fmap[where.field]
			if ok {
				if len(fields) > idx {
					value = fields[idx]
				}
			}
			if !where.match(value) {
				return sw.fvals, false
			}
		}
	} else {
		for idx := range sw.farr {
			var value float64
			if len(fields) > idx {
				value = fields[idx]
			}
			sw.fvals[idx] = value
		}
		for _, where := range sw.wheres {
			if where.field == "z" {
				if !gotz {
					z = o.CalculatedPoint().Z
				}
				if !where.match(z) {
					return sw.fvals, false
				}
				continue
			}
			var value float64
			idx, ok := sw.fmap[where.field]
			if ok {
				value = sw.fvals[idx]
			}
			if !where.match(value) {
				return sw.fvals, false
			}
		}
	}
	return sw.fvals, true
}

func (sw *scanWriter) writeObject(id string, o geojson.Object, fields []float64, noLock bool) bool {
	if !noLock {
		sw.mu.Lock()
		defer sw.mu.Unlock()
	}
	keepGoing := true
	if !sw.globEverything {
		if sw.globSingle {
			if sw.glob != id {
				return true
			}
			keepGoing = false // return current object and stop iterating
		} else {
			ok, _ := glob.Match(sw.glob, id)
			if !ok {
				return true
			}
		}
	}
	nfields, ok := sw.fieldMatch(fields, o)
	if !ok {
		return true
	}
	sw.count++
	if sw.output == outputCount {
		return true
	}

	switch sw.msg.OutputType {
	case server.JSON:
		var wr bytes.Buffer
		var jsfields string
		if sw.once {
			wr.WriteByte(',')
		} else {
			sw.once = true
		}
		if sw.hasFieldsOutput() {
			if sw.fullFields {
				if len(sw.fmap) > 0 {
					jsfields = `,"fields":{`
					var i int
					for field, idx := range sw.fmap {
						if len(fields) > idx {
							if fields[idx] != 0 {
								if i > 0 {
									jsfields += `,`
								}
								jsfields += jsonString(field) + ":" + strconv.FormatFloat(fields[idx], 'f', -1, 64)
								i++
							}
						}
					}
					jsfields += `}`
				}

			} else if len(sw.farr) > 0 {
				jsfields = `,"fields":[`
				for i, field := range nfields {
					if i > 0 {
						jsfields += ","
					}
					jsfields += strconv.FormatFloat(field, 'f', -1, 64)
				}
				jsfields += `]`
			}
		}
		if sw.output == outputIDs {
			wr.WriteString(jsonString(id))
		} else {
			wr.WriteString(`{"id":` + jsonString(id))
			switch sw.output {
			case outputObjects:
				wr.WriteString(`,"object":` + o.JSON())
			case outputPoints:
				wr.WriteString(`,"point":` + o.CalculatedPoint().ExternalJSON())
			case outputHashes:
				p, err := o.Geohash(int(sw.precision))
				if err != nil {
					p = ""
				}
				wr.WriteString(`,"hash":"` + p + `"`)
			case outputBounds:
				wr.WriteString(`,"bounds":` + o.CalculatedBBox().ExternalJSON())
			}
			wr.WriteString(jsfields)
			wr.WriteString(`}`)
		}
		sw.wr.Write(wr.Bytes())
	case server.RESP:
		vals := make([]resp.Value, 1, 3)
		vals[0] = resp.StringValue(id)
		if sw.output == outputIDs {
			sw.values = append(sw.values, vals[0])
		} else {
			switch sw.output {
			case outputObjects:
				vals = append(vals, resp.StringValue(o.String()))
			case outputPoints:
				point := o.CalculatedPoint()
				if point.Z != 0 {
					vals = append(vals, resp.ArrayValue([]resp.Value{
						resp.FloatValue(point.Y),
						resp.FloatValue(point.X),
						resp.FloatValue(point.Z),
					}))
				} else {
					vals = append(vals, resp.ArrayValue([]resp.Value{
						resp.FloatValue(point.Y),
						resp.FloatValue(point.X),
					}))
				}
			case outputHashes:
				p, err := o.Geohash(int(sw.precision))
				if err != nil {
					p = ""
				}
				vals = append(vals, resp.StringValue(p))
			case outputBounds:
				bbox := o.CalculatedBBox()
				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),
					}),
				}))
			}

			fvs := orderFields(sw.fmap, fields)
			if len(fvs) > 0 {
				fvals := make([]resp.Value, 0, len(fvs)*2)
				for i, fv := range fvs {
					fvals = append(fvals, resp.StringValue(fv.field), resp.StringValue(strconv.FormatFloat(fv.value, 'f', -1, 64)))
					i++
				}
				vals = append(vals, resp.ArrayValue(fvals))
			}

			sw.values = append(sw.values, resp.ArrayValue(vals))
		}
	}
	sw.numberItems++
	if sw.numberItems == sw.limit {
		sw.hitLimit = true
		return false
	}
	return keepGoing
}