geography/handlers/admin/bleve.go

package admin

import (
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net/http"
	"sync"

	"github.com/blevesearch/bleve/v2"
	"github.com/blevesearch/bleve/v2/geo"
	"github.com/paulmach/orb/geojson"
	"github.com/rs/zerolog/log"
	"github.com/spf13/viper"
)

type AdminBleveIndex struct {
	Index bleve.Index
	Data  map[string]*geojson.Feature
}

func NewAdminBleveIndex(cfg *viper.Viper) (*AdminBleveIndex, error) {
	file := cfg.GetString("storage.index.bleve.file")
	if file == "" {
		return nil, errors.New("issue in config : storage.index.beleve.file missing")
	}
	index, err := bleve.Open(file)
	if err != nil {
		log.Info().Msg("Bleve index does not exist : creating")
		mapping := bleve.NewIndexMapping()
		docMapping := bleve.NewDocumentMapping()
		docMapping.AddFieldMappingsAt("geometry", bleve.NewGeoShapeFieldMapping())
		// docMapping.AddFieldMappingsAt("properties.nom", bleve.NewTextFieldMapping())
		mapping.DefaultMapping = docMapping
		index, err := bleve.New(file, mapping)
		if err != nil {
			return nil, fmt.Errorf("issue creating index : %v", err)
		}
		indexstorage := &AdminBleveIndex{
			Index: index,
			Data:  map[string]*geojson.Feature{},
		}
		err = indexstorage.IndexData(cfg.GetStringMapString("data.layers"))
		return indexstorage, nil
	}

	return &AdminBleveIndex{
		Index: index,
		Data:  map[string]*geojson.Feature{},
	}, nil
}

func (b *AdminBleveIndex) IndexData(datalayers map[string]string) error {
	log.Info().Msg("indexing data")
	var wg sync.WaitGroup
	var errs chan error
	for layer, url := range datalayers {
		wg.Add(1)
		go func(errs chan error) {
			defer wg.Done()
			resp, err := http.Get(url)
			if err != nil {
				log.Error().Err(err).Msg("get request failed")
				return
			}

			body, err := io.ReadAll(resp.Body)
			if err != nil {
				log.Error().Err(err).Msg("cound not read response")
				return
			}

			data, err := geojson.UnmarshalFeatureCollection(body)
			if err != nil {
				log.Error().Err(err).Msg("could not read input data")
				return
			}
			for _, feature := range data.Features {
				if err = b.Store(layer, feature); err != nil {
					log.Error().Err(err).Msg("could not store data")
				}
			}
		}(errs)

	}

	wg.Wait()

	if nb, err := b.Index.DocCount(); err == nil {
		log.Info().Uint64("number of documents", nb).Msg("indexed all data")
	}

	return nil
}

func (b *AdminBleveIndex) Store(layer string, feature *geojson.Feature) error {
	id := fmt.Sprintf("%s/%s", layer, feature.Properties.MustString("code"))
	log.Debug().Str("id", id).Str("name", feature.Properties.MustString("nom", "")).Msg("store data")
	f, err := feature.MarshalJSON()
	if err != nil {
		log.Error().Err(err).Msg("geojson to json bytes issue")
		return err
	}
	var doc map[string]any

	if err = json.Unmarshal(f, &doc); err != nil {
		log.Error().Err(err).Msg("json reading issue")
		return err
	}
	err = b.Index.Index(id, doc)
	if err != nil {
		log.Error().Err(err).Msg("error adding to index")
	}

	b.Data[id] = feature
	return nil
}

func (b *AdminBleveIndex) GeoSearch(f *geojson.Feature) (map[string][]*geojson.Feature, error) {
	coordinates := []float64{f.Point().Lon(), f.Point().Lat()}
	query, err := bleve.NewGeoShapeQuery(
		[][][][]float64{{{coordinates}}},
		geo.PointType,
		"contains",
	)
	if err != nil {
		return nil, fmt.Errorf("could not create query : %v", err)
	}
	query.SetField("geometry")
	search := bleve.NewSearchRequest(query)
	log.Debug().Any("query", query).Msg("query")
	res, err := b.Index.Search(search)
	if err != nil {
		return nil, fmt.Errorf("issue in search : %v", err)
	}
	log.Debug().Any("result", res).Msg("got result")

	for _, r := range res.Hits {
		log.Debug().Any("region", b.Data[r.ID].Properties.MustString("nom")).Str("type", b.Data[r.ID].Geometry.GeoJSONType()).Msg("hit")
	}

	return map[string][]*geojson.Feature{}, nil
}

func (s *AdminBleveIndex) Find(layer, id string) (*geojson.Feature, error) {
	result, ok := s.Data[fmt.Sprintf("%s/%s", layer, id)]
	if !ok || result == nil {
		return nil, errors.New("admin not found")
	}
	return result, nil
}