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 if opts.Trace {
63 stats = &zoekt.Stats{}
64 return s.Streamer.StreamSearch(ctx, q, opts, stream.SenderFunc(func(event *zoekt.SearchResult) {
65 stats.Add(event.Stats)
66 sender.Send(event)
67 }))
68 }
69
70 return s.Streamer.StreamSearch(ctx, q, opts, sender)
71}
72
73func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (rl *zoekt.RepoList, err error) {
74 tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
75 tr.LazyLog(q, true)
76 tr.LazyPrintf("opts: %s", opts)
77 defer func() {
78 if rl != nil {
79 tr.LazyPrintf("repos size: %d", len(rl.Repos))
80 tr.LazyPrintf("crashes: %d", rl.Crashes)
81 tr.LazyPrintf("minimal size : %d", len(rl.Minimal))
82 tr.LazyPrintf("stats: %+v", rl.Stats)
83 }
84 if err != nil {
85 tr.LazyPrintf("error: %v", err)
86 tr.SetError(err)
87 }
88 tr.Finish()
89 }()
90
91 q, err = s.eval(ctx, q)
92 if err != nil {
93 return nil, err
94 }
95
96 return s.Streamer.List(ctx, q, opts)
97}
98
99func (s *typeRepoSearcher) eval(ctx context.Context, q query.Q) (query.Q, error) {
100 var err error
101 q = query.Map(q, func(q query.Q) query.Q {
102 if err != nil {
103 return nil
104 }
105
106 rq, ok := q.(*query.Type)
107 if !ok || rq.Type != query.TypeRepo {
108 return q
109 }
110
111 var rl *zoekt.RepoList
112 rl, err = s.Streamer.List(ctx, rq.Child, nil)
113 if err != nil {
114 return nil
115 }
116
117 rs := &query.RepoSet{Set: make(map[string]bool, len(rl.Repos))}
118 for _, r := range rl.Repos {
119 rs.Set[r.Repository.Name] = true
120 }
121 return rs
122 })
123 return q, err
124}