117 lines
2.4 KiB
Go
117 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"errors"
|
|
"io/ioutil"
|
|
"os"
|
|
"strings"
|
|
|
|
"github.com/sirupsen/logrus"
|
|
"gopkg.in/yaml.v3"
|
|
|
|
"github.com/blevesearch/bleve/v2"
|
|
)
|
|
|
|
type VisitDB struct {
|
|
visitsDB []Visit
|
|
visitsIdx bleve.Index
|
|
}
|
|
|
|
func NewVisitDB() VisitDB {
|
|
db := VisitDB{}
|
|
db.init()
|
|
return db
|
|
}
|
|
func (vdb *VisitDB) FindById(visitId string) (Visit, error) {
|
|
for _, visit := range vdb.visitsDB {
|
|
if visit.ID == visitId {
|
|
return visit, nil
|
|
}
|
|
}
|
|
return Visit{}, errors.New("visit not found")
|
|
}
|
|
func (vdb *VisitDB) init() {
|
|
idxPath := "visits.bleve"
|
|
if _, err := os.Stat(idxPath); err == nil {
|
|
vdb.visitsIdx, err = bleve.Open(idxPath)
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
} else if os.IsNotExist(err) {
|
|
vdb.visitsIdx, err = bleve.New(idxPath, bleve.NewIndexMapping())
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
} else {
|
|
panic(err)
|
|
}
|
|
err := vdb.loadYAMLDB("db.yaml")
|
|
if err != nil {
|
|
logrus.Fatalf("Failed to load db.yaml: %v", err)
|
|
panic(err)
|
|
}
|
|
}
|
|
|
|
func (vdb *VisitDB) loadYAMLDB(path string) error {
|
|
data, err := ioutil.ReadFile(path)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := yaml.Unmarshal(data, &vdb.visitsDB); err != nil {
|
|
return err
|
|
}
|
|
return vdb.indexVisits(vdb.visitsDB)
|
|
}
|
|
|
|
func (vdb *VisitDB) indexVisits(visits []Visit) error {
|
|
batch := vdb.visitsIdx.NewBatch()
|
|
for _, visit := range visits {
|
|
if err := batch.Index(visit.ID, visit); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return vdb.visitsIdx.Batch(batch)
|
|
}
|
|
|
|
// FindCandidates returns visits with overlapping keywords
|
|
func (vdb *VisitDB) FindCandidates(keywords []string) ([]Visit, error) {
|
|
query := bleve.NewMatchQuery(strings.Join(keywords, " "))
|
|
search := bleve.NewSearchRequest(query)
|
|
searchResults, err := vdb.visitsIdx.Search(search)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
var candidates []Visit
|
|
for _, hit := range searchResults.Hits {
|
|
for _, r := range vdb.visitsDB {
|
|
if r.ID == hit.ID {
|
|
candidates = append(candidates, r)
|
|
break
|
|
}
|
|
}
|
|
}
|
|
return candidates, nil
|
|
}
|
|
|
|
// sumProcedures calculates total price and duration
|
|
func sumProcedures(procs []Procedure) (int, int) {
|
|
totalPrice := 0
|
|
totalDuration := 0
|
|
for _, p := range procs {
|
|
totalPrice += p.Price
|
|
totalDuration += p.DurationMin
|
|
}
|
|
return totalPrice, totalDuration
|
|
}
|
|
|
|
// VisitDBAPI allows mocking VisitDB in other places
|
|
// Only public methods should be included
|
|
|
|
type VisitDBAPI interface {
|
|
FindById(visitId string) (Visit, error)
|
|
FindCandidates(keywords []string) ([]Visit, error)
|
|
}
|
|
|
|
var _ VisitDBAPI = (*VisitDB)(nil)
|