Experiment/Draft: Search engine

go.mod

@@ -47,13 +47,30 @@ require (
 )
 
 require (
+	github.com/blevesearch/bleve/v2 v2.2.2
 	github.com/vearutop/statigz v1.1.6
+	github.com/vektah/dataloaden v0.2.1-0.20190515034641-a19b9a6e7c9e
 	github.com/vektah/gqlparser/v2 v2.0.1
 )
 
 require (
+	github.com/RoaringBitmap/roaring v0.9.4 // indirect
 	github.com/agnivade/levenshtein v1.1.0 // indirect
 	github.com/antchfx/xpath v1.1.6 // indirect
+	github.com/bits-and-blooms/bitset v1.2.0 // indirect
+	github.com/blevesearch/bleve_index_api v1.0.1 // indirect
+	github.com/blevesearch/go-porterstemmer v1.0.3 // indirect
+	github.com/blevesearch/mmap-go v1.0.3 // indirect
+	github.com/blevesearch/scorch_segment_api/v2 v2.1.0 // indirect
+	github.com/blevesearch/segment v0.9.0 // indirect
+	github.com/blevesearch/snowballstem v0.9.0 // indirect
+	github.com/blevesearch/upsidedown_store_api v1.0.1 // indirect
+	github.com/blevesearch/vellum v1.0.7 // indirect
+	github.com/blevesearch/zapx/v11 v11.3.1 // indirect
+	github.com/blevesearch/zapx/v12 v12.3.1 // indirect
+	github.com/blevesearch/zapx/v13 v13.3.1 // indirect
+	github.com/blevesearch/zapx/v14 v14.3.1 // indirect
+	github.com/blevesearch/zapx/v15 v15.3.1 // indirect
 	github.com/chromedp/sysutil v1.0.0 // indirect
 	github.com/cpuguy83/go-md2man/v2 v2.0.0 // indirect
 	github.com/davecgh/go-spew v1.1.1 // indirect
@@ -62,6 +79,8 @@ require (
 	github.com/gobwas/pool v0.2.1 // indirect
 	github.com/gobwas/ws v1.1.0-rc.5 // indirect
 	github.com/golang/groupcache v0.0.0-20200121045136-8c9f03a8e57e // indirect
+	github.com/golang/protobuf v1.5.2 // indirect
+	github.com/golang/snappy v0.0.3 // indirect
 	github.com/hashicorp/errwrap v1.0.0 // indirect
 	github.com/hashicorp/go-multierror v1.1.0 // indirect
 	github.com/hashicorp/golang-lru v0.5.1 // indirect
@@ -75,6 +94,7 @@ require (
 	github.com/mitchellh/mapstructure v1.4.2 // indirect
 	github.com/modern-go/concurrent v0.0.0-20180306012644-bacd9c7ef1dd // indirect
 	github.com/modern-go/reflect2 v1.0.1 // indirect
+	github.com/mschoch/smat v0.2.0 // indirect
 	github.com/nfnt/resize v0.0.0-20180221191011-83c6a9932646 // indirect
 	github.com/pelletier/go-toml v1.9.4 // indirect
 	github.com/pkg/errors v0.9.1 // indirect
@@ -85,14 +105,16 @@ require (
 	github.com/spf13/cast v1.4.1 // indirect
 	github.com/spf13/cobra v1.0.0 // indirect
 	github.com/spf13/jwalterweatherman v1.1.0 // indirect
+	github.com/steveyen/gtreap v0.1.0 // indirect
 	github.com/stretchr/objx v0.2.0 // indirect
 	github.com/subosito/gotenv v1.2.0 // indirect
 	github.com/tidwall/match v1.1.1 // indirect
 	github.com/urfave/cli/v2 v2.1.1 // indirect
-	github.com/vektah/dataloaden v0.2.1-0.20190515034641-a19b9a6e7c9e // indirect
+	go.etcd.io/bbolt v1.3.5 // indirect
 	go.uber.org/atomic v1.7.0 // indirect
 	golang.org/x/mod v0.4.2 // indirect
 	golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1 // indirect
+	google.golang.org/protobuf v1.27.1 // indirect
 	gopkg.in/ini.v1 v1.63.2 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20210107192922-496545a6307b // indirect
 )

graphql/schema/schema.graphql

@@ -64,6 +64,9 @@ type Query {
 
   logs: [LogEntry!]!
 
+  # Search
+  search(query: String!, type: SearchType): SearchResultItemConnection!
+
   # Scrapers
 
   """List available scrapers"""

graphql/schema/types/performer.graphql

@@ -7,7 +7,7 @@ enum GenderEnum {
   NON_BINARY
 }
 
-type Performer {
+type Performer implements SearchResultItem {
   id: ID!
   checksum: String!
   name: String

graphql/schema/types/scene.graphql

@@ -25,7 +25,7 @@ type SceneMovie {
   scene_index: Int
 }
 
-type Scene {
+type Scene implements SearchResultItem {
   id: ID!
   checksum: String
   oshash: String

graphql/schema/types/search.graphql

@@ -0,0 +1,31 @@
+enum SearchType {
+    SEARCH_SCENE
+    SEARCH_PERFORMER
+    SEARCH_TAG
+    SEARCH_STUDIO
+}
+
+type SearchResultItemConnection {
+    edges: [SearchItemEdge]
+
+    took: Float!
+    maxScore: Float!
+    total: Int!
+
+    status: SearchResultStatus
+}
+
+type SearchResultStatus {
+    successful: Int!
+    failed: Int!
+    total: Int!
+}
+
+type SearchItemEdge {
+    node: SearchResultItem
+    score: Float!
+}
+
+interface SearchResultItem {
+    id: ID!
+}

graphql/schema/types/studio.graphql

@@ -1,4 +1,4 @@
-type Studio {
+type Studio implements SearchResultItem {
   id: ID!
   checksum: String!
   name: String!

graphql/schema/types/tag.graphql

@@ -1,4 +1,4 @@
-type Tag {
+type Tag implements SearchResultItem {
   id: ID!
   name: String!
   aliases: [String!]!

pkg/api/resolver.go

@@ -11,6 +11,7 @@ import (
 	"github.com/stashapp/stash/pkg/models"
 	"github.com/stashapp/stash/pkg/plugin"
 	"github.com/stashapp/stash/pkg/scraper"
+	"github.com/stashapp/stash/pkg/search"
 )
 
 var (
@@ -32,6 +33,7 @@ type hookExecutor interface {
 type Resolver struct {
 	txnManager   models.TransactionManager
 	hookExecutor hookExecutor
+	searchEngine *search.Engine
 }
 
 func (r *Resolver) scraperCache() *scraper.Cache {

pkg/api/resolver_query_search.go

@@ -0,0 +1,60 @@
+package api
+
+import (
+	"context"
+	"errors"
+	"fmt"
+
+	"github.com/stashapp/stash/pkg/models"
+	"github.com/stashapp/stash/pkg/search"
+	"github.com/stashapp/stash/pkg/search/documents"
+)
+
+var ErrUnknownType = errors.New("unknown item type")
+
+func (r *queryResolver) Search(ctx context.Context, query string, ty *models.SearchType) (*models.SearchResultItemConnection, error) {
+	s, err := r.searchEngine.Search(ctx, query, ty)
+	if err != nil {
+		return nil, err
+	}
+
+	var edges []*models.SearchItemEdge
+	for _, item := range s.Items {
+		h, err := r.hydrate(ctx, item)
+		if err != nil {
+			edges = append(edges, nil)
+			continue
+		}
+
+		edges = append(edges, &models.SearchItemEdge{
+			Score: item.Score,
+			Node:  h,
+		})
+	}
+
+	res := models.SearchResultItemConnection{
+		Edges:    edges,
+		Took:     s.Took.Seconds(),
+		MaxScore: s.MaxScore,
+		Total:    int(s.Total),
+
+		Status: s.Status,
+	}
+
+	return &res, nil
+}
+
+func (r *queryResolver) hydrate(ctx context.Context, item search.Item) (models.SearchResultItem, error) {
+	switch item.Type {
+	case documents.TypeScene:
+		return r.FindScene(ctx, &item.ID, nil)
+	case documents.TypePerformer:
+		return r.FindPerformer(ctx, item.ID)
+	case documents.TypeTag:
+		return r.FindTag(ctx, item.ID)
+	case documents.TypeStudio:
+		return r.FindStudio(ctx, item.ID)
+	default:
+		return nil, fmt.Errorf("%w: %v", ErrUnknownType, item.Type)
+	}
+}

pkg/api/server.go

@@ -69,9 +69,11 @@ func Start(uiBox embed.FS, loginUIBox embed.FS) {
 
 	txnManager := manager.GetInstance().TxnManager
 	pluginCache := manager.GetInstance().PluginCache
+	searchEngine := manager.GetInstance().Search
 	resolver := &Resolver{
 		txnManager:   txnManager,
 		hookExecutor: pluginCache,
+		searchEngine: searchEngine,
 	}
 
 	gqlSrv := gqlHandler.New(models.NewExecutableSchema(models.Config{Resolvers: resolver}))

pkg/event/dispatcher.go

@@ -0,0 +1,133 @@
+// package event dispatches change-events in stash
+package event
+
+import (
+	"context"
+	"fmt"
+	"sync"
+)
+
+// ChangeType defines what has changed
+type ChangeType int
+
+const (
+	SceneMarker ChangeType = iota
+	Scene
+	Image
+	Gallery
+	Movie
+	Performer
+	Studio
+	Tag
+)
+
+func (ct ChangeType) String() string {
+	switch ct {
+	case SceneMarker:
+		return "SceneMarker"
+	case Scene:
+		return "Scene"
+	case Image:
+		return "Image"
+	case Gallery:
+		return "Gallery"
+	case Movie:
+		return "Movie"
+	case Performer:
+		return "Performer"
+	case Studio:
+		return "Studio"
+	case Tag:
+		return "Tag"
+	default:
+		panic("Unhandled ChangeType case")
+	}
+}
+
+// Change represents a pair of a Type and the ID which has changed
+type Change struct {
+	Type ChangeType
+	ID   int
+}
+
+func (c Change) String() string {
+	return fmt.Sprintf("%v(%v)", c.Type, c.ID)
+}
+
+// Dispatcher represents a single event dispatcher
+type Dispatcher struct {
+	incoming chan Change
+	mu       sync.Mutex
+	chans    []chan Change
+}
+
+// NewDispatcher creates a new even dispatcher
+func NewDispatcher() *Dispatcher {
+	incoming := make(chan Change, 1)
+	return &Dispatcher{
+		incoming: incoming,
+	}
+}
+
+// Start starts the dispatcher goroutine under the given context
+func (d *Dispatcher) Start(ctx context.Context) {
+	go func() {
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case change := <-d.incoming:
+				d.broadcast(change)
+			}
+		}
+	}()
+}
+
+// Register registers chan for receiving events. It is up to the caller to ensure
+// the channel doesn't block. I.e., events must be lifted from chan quickly.
+func (d *Dispatcher) Register(c chan Change) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	d.chans = append(d.chans, c)
+}
+
+// Unregister removes a channel for dispatches
+func (d *Dispatcher) Unregister(c chan Change) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	// Find the index at which the chan occur
+	idx := -1
+	for i := range d.chans {
+		if d.chans[i] == c {
+			idx = i
+			break
+		}
+	}
+
+	// If already unregistered, ignore
+	if idx == -1 {
+		return
+	}
+
+	// Swap chan to last element, cut it off
+	last := len(d.chans) - 1
+	d.chans[idx] = d.chans[last]
+	d.chans = d.chans[:last]
+}
+
+// Publish broadcasts a change
+func (d *Dispatcher) Publish(c Change) {
+	d.incoming <- c
+}
+
+// broadcast fans-out a change to all channels registered
+func (d *Dispatcher) broadcast(change Change) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	for _, ch := range d.chans {
+		ch <- change
+	}
+}