fork of https://github.com/sourcegraph/zoekt
1package shards
2
3import (
4 "context"
5
6 "github.com/sourcegraph/zoekt"
7 "github.com/sourcegraph/zoekt/query"
8 "github.com/sourcegraph/zoekt/stream"
9 "github.com/sourcegraph/zoekt/trace"
10)
11
12// typeRepoSearcher evaluates all type:repo sub-queries before sending the query
13// to the underlying searcher. We need to evaluate type:repo queries first
14// since they need to do cross shard operations.
15type typeRepoSearcher struct {
16 zoekt.Streamer
17}
18
19func (s *typeRepoSearcher) Search(ctx context.Context, q query.Q, opts *zoekt.SearchOptions) (sr *zoekt.SearchResult, err error) {
20 tr, ctx := trace.New(ctx, "typeRepoSearcher.Search", "")
21 tr.LazyLog(q, true)
22 tr.LazyPrintf("opts: %+v", opts)
23 defer func() {
24 if sr != nil {
25 tr.LazyPrintf("num files: %d", len(sr.Files))
26 tr.LazyPrintf("stats: %+v", sr.Stats)
27 }
28 if err != nil {
29 tr.LazyPrintf("error: %v", err)
30 tr.SetError(err)
31 }
32 tr.Finish()
33 }()
34
35 q, err = s.eval(ctx, q)
36 if err != nil {
37 return nil, err
38 }
39
40 return s.Streamer.Search(ctx, q, opts)
41}
42
43func (s *typeRepoSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zoekt.SearchOptions, sender zoekt.Sender) (err error) {
44 tr, ctx := trace.New(ctx, "typeRepoSearcher.StreamSearch", "")
45 tr.LazyLog(q, true)
46 tr.LazyPrintf("opts: %+v", opts)
47 var stats zoekt.Stats
48 defer func() {
49 tr.LazyPrintf("stats: %+v", stats)
50 if err != nil {
51 tr.LazyPrintf("error: %v", err)
52 tr.SetError(err)
53 }
54 tr.Finish()
55 }()
56
57 q, err = s.eval(ctx, q)
58 if err != nil {
59 return err
60 }
61
62 return s.Streamer.StreamSearch(ctx, q, opts, stream.SenderFunc(func(event *zoekt.SearchResult) {
63 stats.Add(event.Stats)
64 sender.Send(event)
65 }))
66}
67
68func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (rl *zoekt.RepoList, err error) {
69 tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
70 tr.LazyLog(q, true)
71 tr.LazyPrintf("opts: %s", opts)
72 defer func() {
73 if rl != nil {
74 tr.LazyPrintf("repos size: %d", len(rl.Repos))
75 tr.LazyPrintf("reposmap size: %d", len(rl.ReposMap))
76 tr.LazyPrintf("crashes: %d", rl.Crashes)
77 tr.LazyPrintf("stats: %+v", rl.Stats)
78 }
79 if err != nil {
80 tr.LazyPrintf("error: %v", err)
81 tr.SetError(err)
82 }
83 tr.Finish()
84 }()
85
86 q, err = s.eval(ctx, q)
87 if err != nil {
88 return nil, err
89 }
90
91 return s.Streamer.List(ctx, q, opts)
92}
93
94func (s *typeRepoSearcher) eval(ctx context.Context, q query.Q) (query.Q, error) {
95 var err error
96 q = query.Map(q, func(q query.Q) query.Q {
97 if err != nil {
98 return nil
99 }
100
101 rq, ok := q.(*query.Type)
102 if !ok || rq.Type != query.TypeRepo {
103 return q
104 }
105
106 var rl *zoekt.RepoList
107 rl, err = s.Streamer.List(ctx, rq.Child, nil)
108 if err != nil {
109 return nil
110 }
111
112 rs := &query.RepoSet{Set: make(map[string]bool, len(rl.Repos))}
113 for _, r := range rl.Repos {
114 rs.Set[r.Repository.Name] = true
115 }
116 return rs
117 })
118 return q, err
119}