fork of https://github.com/sourcegraph/zoekt
1package search
2
3import (
4 "context"
5
6 "github.com/sourcegraph/zoekt"
7 "github.com/sourcegraph/zoekt/internal/tenant"
8 "github.com/sourcegraph/zoekt/internal/trace"
9 "github.com/sourcegraph/zoekt/query"
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 tenant.Log(ctx, tr)
24 defer func() {
25 if sr != nil {
26 tr.LazyPrintf("num files: %d", len(sr.Files))
27 tr.LazyPrintf("stats: %+v", sr.Stats)
28 }
29 if err != nil {
30 tr.LazyPrintf("error: %v", err)
31 tr.SetError(err)
32 }
33 tr.Finish()
34 }()
35
36 q, err = s.eval(ctx, tr, q)
37 if err != nil {
38 return nil, err
39 }
40
41 return s.Streamer.Search(ctx, q, opts)
42}
43
44func (s *typeRepoSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zoekt.SearchOptions, sender zoekt.Sender) (err error) {
45 tr, ctx := trace.New(ctx, "typeRepoSearcher.StreamSearch", "")
46 tr.LazyLog(q, true)
47 tr.LazyPrintf("opts: %+v", opts)
48 tenant.Log(ctx, tr)
49 var stats zoekt.Stats
50 defer func() {
51 tr.LazyPrintf("stats: %+v", stats)
52 if err != nil {
53 tr.LazyPrintf("error: %v", err)
54 tr.SetError(err)
55 }
56 tr.Finish()
57 }()
58
59 q, err = s.eval(ctx, tr, q)
60 if err != nil {
61 return err
62 }
63
64 return s.Streamer.StreamSearch(ctx, q, opts, zoekt.SenderFunc(func(event *zoekt.SearchResult) {
65 stats.Add(event.Stats)
66 sender.Send(event)
67 }))
68}
69
70func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (rl *zoekt.RepoList, err error) {
71 tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
72 tr.LazyLog(q, true)
73 tr.LazyPrintf("opts: %s", opts)
74 tenant.Log(ctx, tr)
75 defer func() {
76 if rl != nil {
77 tr.LazyPrintf("repos.size=%d reposmap.size=%d crashes=%d stats=%+v", len(rl.Repos), len(rl.ReposMap), rl.Crashes, 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, tr, 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, tr *trace.Trace, 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 tr.LazyPrintf("evaluating sub-expression %s", rq)
107
108 var rl *zoekt.RepoList
109 rl, err = s.Streamer.List(ctx, rq.Child, nil)
110 if err != nil {
111 return nil
112 }
113
114 rs := &query.RepoSet{Set: make(map[string]bool, len(rl.Repos))}
115 for _, r := range rl.Repos {
116 rs.Set[r.Repository.Name] = true
117 }
118
119 tr.LazyPrintf("replaced sub-expression with %s", rs)
120
121 return rs
122 })
123 return q, err
124}