vetrag/db.go

107 lines
2.2 KiB
Go

package main
import (
"errors"
"io/ioutil"
"os"
"strings"
"github.com/sirupsen/logrus"
"gopkg.in/yaml.v3"
"github.com/blevesearch/bleve/v2"
)
type ReasonDB struct {
reasonsDB []Reason
reasonsIdx bleve.Index
}
func NewReasonDB() ReasonDB {
db := ReasonDB{}
db.init()
return db
}
func (rdb *ReasonDB) findById(reasonId string) (Reason, error) {
for _, reason := range rdb.reasonsDB {
if reason.ID == reasonId {
return reason, nil
}
}
return Reason{}, errors.New("reason not found")
}
func (rdb *ReasonDB) init() {
idxPath := "reasons.bleve"
if _, err := os.Stat(idxPath); err == nil {
rdb.reasonsIdx, err = bleve.Open(idxPath)
if err != nil {
panic(err)
}
} else if os.IsNotExist(err) {
rdb.reasonsIdx, err = bleve.New(idxPath, bleve.NewIndexMapping())
if err != nil {
panic(err)
}
} else {
panic(err)
}
err := rdb.loadYAMLDB("db.yaml")
if err != nil {
logrus.Fatalf("Failed to load db.yaml: %v", err)
panic(err)
}
}
func (rdb *ReasonDB) loadYAMLDB(path string) error {
data, err := ioutil.ReadFile(path)
if err != nil {
return err
}
if err := yaml.Unmarshal(data, &rdb.reasonsDB); err != nil {
return err
}
return rdb.indexReasons(rdb.reasonsDB)
}
func (rdb *ReasonDB) indexReasons(reasons []Reason) error {
batch := rdb.reasonsIdx.NewBatch()
for _, reason := range reasons {
if err := batch.Index(reason.ID, reason); err != nil {
return err
}
}
return rdb.reasonsIdx.Batch(batch)
}
// findCandidates returns reasons with overlapping keywords
func (rdb *ReasonDB) findCandidates(keywords []string) ([]Reason, error) {
query := bleve.NewMatchQuery(strings.Join(keywords, " "))
search := bleve.NewSearchRequest(query)
searchResults, err := rdb.reasonsIdx.Search(search)
if err != nil {
return nil, err
}
var candidates []Reason
for _, hit := range searchResults.Hits {
for _, r := range rdb.reasonsDB {
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
}