eintopf.info@v0.13.16/service/search/search.go

// Copyright (C) 2022 The Eintopf authors
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as
// published by the Free Software Foundation, either version 3 of the
// License, or (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package search

import (
	"context"
	"encoding/json"
	"fmt"
	"log"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"

	"eintopf.info/internal/cache"
	"github.com/blevesearch/bleve/v2"
	"github.com/blevesearch/bleve/v2/analysis/lang/de"
	"github.com/blevesearch/bleve/v2/mapping"
	"github.com/blevesearch/bleve/v2/search/query"
)

// Indexable defines a document, that can be indexed.
// Any model, that implements this interface can be indexed and therefore is
// searchable.
// An indexable document can be uniquely identifyed by combining its identifier
// and type.
type Indexable interface {
	// Identifier returns an id which should uniquely identify the object for
	// its type.
	Identifier() string

	// Type returns the type of the object.
	Type() string

	// QueryText returns the string to index for a text search.
	QueryText() string

	// SearchFields returns a map of additional fields to be indexed. Those
	// fields can be used for filtering or aggregations.
	SearchFields() map[string]interface{}
}

// Service defines a search service.
//
// -go:generate go run github.com/petergtz/pegomock/pegomock generate eintopf.info/service/search Service --output=../../internal/mock/search_service.go --package=mock --mock-name=SearchService
type Service interface {
	// Index takes one or many indexable document and adds them to the search index.
	Index(docs ...Indexable) error

	// Delete deletes the document with the given type and id.
	// Note: both type and id have to be provided in order to uniquely identify
	// the document.
	Delete(docType, id string) error

	// Search performes a full text search on all indexed documents.
	// Performes a wildcard match for empty queries.
	// In addition to the search query a set of additional search options can be
	// provided.
	Search(ctx context.Context, opts *Options) (*Result, error)

	// LastModified returns the last time, the store was updated. This can be
	// used to invalidate a client side cache.
	LastModified() time.Time

	// Stops the search service.
	Stop()
}

// Options defines a set of optional search options.
type Options struct {
	// Query is the search query used for a text search.
	Query string `json:"query"`

	// Sort is the field, that should be sorted by.
	// When left empty, the default sorting is used.
	Sort string `json:"sort"`

	// SortDescending defines the sort order.
	SortDescending bool `json:"sortAscending"`

	// Page is current page.
	Page int `json:"page"`

	// PageSize defines the number of hits returned per page.
	//
	// PageSize is infinite when set to 0.
	PageSize int `json:"pageSize"`

	// Filters is a list of filters, that reduce the search result. All filters
	// are combined with AND logic in addition with the search query.
	Filters []Filter `json:"filter"`

	// Aggregations is a map of aggregations, to perform aggregations on fields.
	// The provided map key can be used to identify the corresponding bucket in
	// the result.
	Aggregations map[string]Aggregation `json:"aggregations"`
}

// CacheKey returns a string uniquely identifying the Options object.
func (o *Options) CacheKey() string {
	key := o.Query + o.Sort + strconv.FormatBool(o.SortDescending) + strconv.Itoa(o.Page) + strconv.Itoa(o.Page)
	for _, f := range o.Filters {
		key += f.CacheKey()
	}
	for k, a := range o.Aggregations {
		key += k + a.CacheKey()
	}
	return key
}

// Result contains a search result.
type Result struct {
	// Hits are the search hits for the current pagination.
	Hits []Hit `json:"hits"`

	// Total is the total number of search hits.
	// It is independet of the current pagination.
	Total uint64 `json:"total"`

	// Buckets is a set of aggregation buckets.
	// The map key corresponds to aggregation name.
	Buckets map[string]Bucket `json:"buckets"`
}

// Hit is a single search hit.
type Hit struct {
	// ID is the unique identifier of the stored entity. It might not match the
	// id of the entity in case it was indexed multiple times.
	ID string `json:"id"`
	// Type is the type of the document.
	Type string `json:"type"`
	// Raw contains the raw data of the document in the form of a json string.
	Raw string `json:"raw"`
}

// Unmarshal unmarshals the raw data into v using json.Unmarshal.
func (h *Hit) Unmarshal(v interface{}) error {
	return json.Unmarshal([]byte(h.Raw), v)
}

// NewService returns a new search service.
// Takes the path to the index directory. If the index already exits, the
// existing index is used. Otherwise a new one will be created.
func NewService(indexPath string, searchTimeout time.Duration, resultCacheSize int, bucketCacheSize int, tz *time.Location) (Service, error) {
	var index bleve.Index

	// Check if an index already exists at the specified index path. Use the
	// existing index, if it exists.
	if _, err := os.Stat(indexPath); !os.IsNotExist(err) {
		index, err = bleve.Open(indexPath)
		if err != nil {
			return nil, err
		}
	} else {
		// Create a simple index mapping, containing a single document mapping.
		mapping := bleve.NewIndexMapping()
		mapping.DefaultType = "doc"
		mapping.AddDocumentMapping("doc", documentMapping())

		index, err = bleve.New(indexPath, mapping)
		if err != nil {
			return nil, err
		}
	}

	return &service{
		tz:              tz,
		lastStoreUpdate: time.Now(),
		index:           index,
		searchTimeout:   searchTimeout,
		cache:           cache.NewFavoritesCache[*Result](resultCacheSize),
		bucketCache:     cache.NewFavoritesCache[Bucket](bucketCacheSize),
	}, nil
}

// service is an implementation of the Service interface using bleve search. It
// is internal to hide implementation details.
type service struct {
	tz *time.Location

	// lastUpdate stores the time of the last store update.
	lastStoreUpdate  time.Time
	mLastStoreUpdate sync.Mutex

	// index is the bleve index. The index is kept simple, by only storing one
	// data structure (document).
	index bleve.Index

	// searchTimeout is the maximum duration a search request may take.
	// When the request takes longer, it gets canceled.
	searchTimeout time.Duration

	// cache stores search result.
	//
	// The cache gets cleared, when an index or delete operation is performed
	cache *cache.Favorites[*Result]
	// bucketCache stores buckets resulting from a unique aggregation. The cache
	// key is retrieved from the CacheKey method on the aggregation.
	//
	// The cache gets cleared, when an index or delete operation is performed.
	bucketCache *cache.Favorites[Bucket]
}

// document is the internal data structure, that gets indexed into bleve.
// This document provides a structure for the index data, while beeing generic
// enough to allow multiple data types, text quering and filters and
// aggregations on specified fields.
type document struct {
	// Type is the external type of the document.
	Type string `json:"type"`
	// Raw stores the actual document in a json marshaled form.
	Raw string `json:"raw"`
	// Query is a special field allowing text queries.
	Query string `json:"query"`
	// Fields holds a set of fields, that can be used for filtering or
	// aggregations.
	Fields map[string]interface{} `json:"fields"`
}

// documentMapping returns the bleve document mapping for the document data
// structure.
func documentMapping() *mapping.DocumentMapping {
	queryMapping := bleve.NewTextFieldMapping()
	queryMapping.Analyzer = de.AnalyzerName

	m := bleve.NewDocumentMapping()
	m.AddFieldMappingsAt("type", bleve.NewKeywordFieldMapping())
	m.AddFieldMappingsAt("query", queryMapping)
	m.AddFieldMappingsAt("raw", &mapping.FieldMapping{
		Type:               "text",
		Store:              true,
		Index:              false,
		IncludeTermVectors: false,
		IncludeInAll:       false,
		DocValues:          false,
	})
	m.AddSubDocumentMapping("fields", bleve.NewDocumentMapping())

	return m
}

func fieldKey(field string) string {
	if field == "type" {
		return field
	}
	return "fields." + field
}

func uniqueID(docType string, id string) string {
	return docType + "_" + id
}

// Index takes an indexable document and converts it into the internal document
// data structure, which then gets indexed into the bleve index.
//
// Returns an error if the document cannot be marshaled into a json string.
// Returns an error if the bleve index operation failed.
func (s *service) Index(docs ...Indexable) error {
	batch := s.index.NewBatch()
	for _, doc := range docs {
		raw, err := json.Marshal(doc)
		if err != nil {
			return err
		}
		batch.Index(uniqueID(doc.Type(), doc.Identifier()), document{
			Type:   doc.Type(),
			Query:  doc.QueryText(),
			Raw:    string(raw),
			Fields: doc.SearchFields(),
		})
	}
	err := s.index.Batch(batch)
	if err != nil {
		return err
	}

	s.cache.Clear()
	s.bucketCache.Clear()
	s.updateLastStoreUpdate()

	return nil
}

// Delete removes the document from the bleve index.
func (s *service) Delete(docType string, id string) error {
	err := s.index.Delete(uniqueID(docType, id))
	if err != nil {
		return err
	}

	s.cache.Clear()
	s.bucketCache.Clear()
	s.updateLastStoreUpdate()

	return nil
}

func (s *service) updateLastStoreUpdate() {
	s.mLastStoreUpdate.Lock()
	s.lastStoreUpdate = time.Now()
	s.mLastStoreUpdate.Unlock()
}

// Search perfomes a search request on the bleve index.
func (s *service) Search(ctx context.Context, opts *Options) (*Result, error) {
	ctx, cancel := context.WithTimeout(ctx, s.searchTimeout)
	defer cancel()

	if opts == nil {
		opts = &Options{}
	}
	if result, ok := s.cache.Get(opts.CacheKey()); ok {
		return result, nil
	}

	sRequest := buildSearchRequest(opts.Query, opts.Filters)
	if opts.PageSize > 0 {
		// If the page size is larger than 0, paginate the result according to
		// the page and page size.
		sRequest.From = opts.Page * opts.PageSize
		sRequest.Size = opts.PageSize
	} else {
		sRequest.Size = 100000
	}

	// Only retrieve the "raw" and "type" fields.
	sRequest.Fields = []string{"raw", "type"}

	if opts.Sort != "" {
		sort := fieldKey(opts.Sort)
		if opts.SortDescending {
			sort = fmt.Sprintf("-%s", sort)
		}
		sRequest.SortBy([]string{sort})
	}

	type searchResult struct {
		err   error
		hits  []Hit
		total uint64
	}
	sChan := make(chan searchResult)
	go func() {
		result, err := s.index.SearchInContext(ctx, sRequest)
		if err != nil {
			sChan <- searchResult{err: err}
			return
		}
		hits := make([]Hit, 0, len(result.Hits))
		for _, hit := range result.Hits {
			typ := hit.Fields["type"].(string)
			hits = append(hits, Hit{
				ID:   strings.TrimPrefix(hit.ID, fmt.Sprintf("%s_", typ)),
				Type: typ,
				Raw:  hit.Fields["raw"].(string),
			})
		}
		sChan <- searchResult{hits: hits, total: result.Total}
	}()

	// Perform a seperate search request per aggregation. This enables
	// independet aggregations.
	buckets := make(map[string]Bucket, len(opts.Aggregations))
	bucketsM := sync.Mutex{}
	bucketsWG := sync.WaitGroup{}

	for name, aggregation := range opts.Aggregations {
		bucketsWG.Add(1)
		go func(name string, aggregation Aggregation) {
			defer bucketsWG.Done()

			// Check if the bucket resulting in this aggregation is cached.
			cacheKey := opts.Query + aggregation.CacheKey()
			if bucket, ok := s.bucketCache.Get(cacheKey); ok && bucket != nil {
				bucketsM.Lock()
				defer bucketsM.Unlock()

				buckets[name] = bucket
				return
			}
			bucket, err := s.aggregate(ctx, opts.Query, aggregation)
			if err != nil {
				// Log the error instead of returning it. This makes sure the search
				// doesn't fail, if a search request has an invalid aggregation.
				log.Printf("aggregate: %s: %s\n", name, err)
			}
			s.bucketCache.Set(cacheKey, bucket)

			bucketsM.Lock()
			defer bucketsM.Unlock()
			buckets[name] = bucket
		}(name, aggregation)
	}

	sResult := <-sChan
	if sResult.err != nil {
		return nil, sResult.err
	}
	bucketsWG.Wait()

	result := &Result{Hits: sResult.hits, Total: sResult.total, Buckets: buckets}

	s.cache.Set(opts.CacheKey(), result)

	return result, nil
}

func (s *service) aggregate(ctx context.Context, queryString string, aggregation Aggregation) (Bucket, error) {
	search := buildSearchRequest(queryString, aggregation.Filters)
	search.Fields = []string{fieldKey(aggregation.Field)}
	search.Size = 10000
	result, err := s.index.SearchInContext(ctx, search)
	if err != nil {
		return nil, fmt.Errorf("aggregate: %s", err)
	}

	var aggregator aggregator
	switch aggregation.Type {
	case TermsAggregation:
		aggregator = &termsAggregator{terms: make(map[string]int)}
	case DateRangeAggregation:
		aggregator = &dateRangeAggregator{
			min: time.Date(9999, 0, 0, 0, 0, 0, 0, s.tz),
			max: time.Unix(0, 0),
		}
	default:
		return nil, fmt.Errorf("invalid aggregation type: %s", aggregation.Type)
	}
	for _, hit := range result.Hits {
		field, ok := hit.Fields[fieldKey(aggregation.Field)]
		if !ok {
			continue
		}
		err := aggregator.aggregate(field)
		if err != nil {
			return nil, fmt.Errorf("aggregagte(%s): %s", field, err)
		}
	}
	return aggregator.bucket(), nil
}

// buildSearchRequest builds a bleve.SearchRequest from a query string and a set
// of filters.
// All filters and the query string query are combined with AND logic.
// For the query string a match query gets created. If the query string is
// empty, a wildcard query gets created.
func buildSearchRequest(queryString string, filters []Filter) *bleve.SearchRequest {
	query := buildQueryStringQuery(queryString)

	if len(filters) > 0 {
		boolQuery := bleve.NewBooleanQuery()
		boolQuery.AddMust(query)
		for _, filter := range filters {
			if f := filter.filterQuery(); f != nil {
				boolQuery.AddMust(f)
			}
		}
		query = boolQuery
	}
	return bleve.NewSearchRequest(query)
}

func buildQueryStringQuery(queryString string) query.Query {
	switch queryString {
	case "":
		wildcardQuery := bleve.NewWildcardQuery("*")
		wildcardQuery.SetField("query")
		return wildcardQuery
	default:
		matchQuery := bleve.NewMatchQuery(queryString)
		matchQuery.SetField("query")
		return matchQuery
	}
}

func (s *service) LastModified() time.Time {
	return s.lastStoreUpdate
}

func (s *service) Stop() {
	s.index.Close()
}