112 lines
2.8 KiB
Go
112 lines
2.8 KiB
Go
package admin
|
|
|
|
import (
|
|
"errors"
|
|
"io"
|
|
"net/http"
|
|
|
|
"github.com/paulmach/orb"
|
|
"github.com/paulmach/orb/geojson"
|
|
"github.com/paulmach/orb/planar"
|
|
"github.com/rs/zerolog/log"
|
|
"github.com/spf13/viper"
|
|
"github.com/tidwall/rtree"
|
|
)
|
|
|
|
type AdminMemoryRTreeIndex struct {
|
|
Layers []string
|
|
Indexes map[string]*rtree.RTree
|
|
Documents map[string]map[string]*geojson.Feature
|
|
}
|
|
|
|
func NewAdminMemoryRTreeIndex(cfg *viper.Viper) (*AdminMemoryRTreeIndex, error) {
|
|
storage := &AdminMemoryRTreeIndex{
|
|
Layers: []string{},
|
|
Indexes: map[string]*rtree.RTree{},
|
|
Documents: map[string]map[string]*geojson.Feature{},
|
|
}
|
|
|
|
layers := cfg.GetStringMapString("data.layers")
|
|
|
|
// var wg sync.WaitGroup
|
|
for layer, url := range layers {
|
|
// wg.Add(1)
|
|
go func() {
|
|
// defer wg.Done()
|
|
storage.IndexLayer(layer, url)
|
|
}()
|
|
}
|
|
|
|
// wg.Wait()
|
|
|
|
return storage, nil
|
|
}
|
|
|
|
func (s *AdminMemoryRTreeIndex) IndexLayer(layer, url string) error {
|
|
log.Info().Str("layer", layer).Msg("indexing layer")
|
|
resp, err := http.Get(url)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("get request failed")
|
|
return err
|
|
}
|
|
|
|
body, err := io.ReadAll(resp.Body)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("cound not read response")
|
|
return err
|
|
}
|
|
|
|
data, err := geojson.UnmarshalFeatureCollection(body)
|
|
if err != nil {
|
|
log.Error().Err(err).Msg("could not read input data")
|
|
return err
|
|
}
|
|
s.Indexes[layer] = &rtree.RTree{}
|
|
s.Documents[layer] = map[string]*geojson.Feature{}
|
|
s.Layers = append(s.Layers, layer)
|
|
for _, feature := range data.Features {
|
|
id := feature.Properties.MustString("code")
|
|
bound := feature.Geometry.Bound()
|
|
s.Indexes[layer].Insert(bound.Min, bound.Max, id)
|
|
s.Documents[layer][id] = feature
|
|
}
|
|
log.Info().Str("layer", layer).Msg("finished indexing")
|
|
return nil
|
|
}
|
|
|
|
func (s *AdminMemoryRTreeIndex) GeoSearch(feature *geojson.Feature) (map[string]*geojson.Feature, error) {
|
|
results := map[string]*geojson.Feature{}
|
|
point := feature.Point()
|
|
for _, layer := range s.Layers {
|
|
rt := s.Indexes[layer]
|
|
rt.Search([2]float64{point.Lon(), point.Lat()}, [2]float64{point.Lon(), point.Lat()}, func(min, max [2]float64, id any) bool {
|
|
resp, found := s.Documents[layer][id.(string)]
|
|
if !found {
|
|
log.Error().Msg("document not found")
|
|
return true
|
|
}
|
|
if poly, ok := resp.Geometry.(orb.Polygon); ok {
|
|
if planar.PolygonContains(poly, point) {
|
|
results[layer] = resp
|
|
return false
|
|
}
|
|
} else if multipoly, ok := resp.Geometry.(orb.MultiPolygon); ok {
|
|
if planar.MultiPolygonContains(multipoly, point) {
|
|
results[layer] = resp
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
})
|
|
}
|
|
return results, nil
|
|
}
|
|
|
|
func (s *AdminMemoryRTreeIndex) Find(layer, id string) (*geojson.Feature, error) {
|
|
result, ok := s.Documents["layer"]["id"]
|
|
if !ok || result == nil {
|
|
return nil, errors.New("admin not found")
|
|
}
|
|
return result, nil
|
|
}
|