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