fork of https://github.com/sourcegraph/zoekt
1package shards
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 if tenant.EnforceTenant() {
24 tenant.Log(ctx, tr)
25 }
26 defer func() {
27 if sr != nil {
28 tr.LazyPrintf("num files: %d", len(sr.Files))
29 tr.LazyPrintf("stats: %+v", sr.Stats)
30 }
31 if err != nil {
32 tr.LazyPrintf("error: %v", err)
33 tr.SetError(err)
34 }
35 tr.Finish()
36 }()
37
38 q, err = s.eval(ctx, tr, q)
39 if err != nil {
40 return nil, err
41 }
42
43 return s.Streamer.Search(ctx, q, opts)
44}
45
46func (s *typeRepoSearcher) StreamSearch(ctx context.Context, q query.Q, opts *zoekt.SearchOptions, sender zoekt.Sender) (err error) {
47 tr, ctx := trace.New(ctx, "typeRepoSearcher.StreamSearch", "")
48 tr.LazyLog(q, true)
49 tr.LazyPrintf("opts: %+v", opts)
50 if tenant.EnforceTenant() {
51 tenant.Log(ctx, tr)
52 }
53 var stats zoekt.Stats
54 defer func() {
55 tr.LazyPrintf("stats: %+v", stats)
56 if err != nil {
57 tr.LazyPrintf("error: %v", err)
58 tr.SetError(err)
59 }
60 tr.Finish()
61 }()
62
63 q, err = s.eval(ctx, tr, q)
64 if err != nil {
65 return err
66 }
67
68 return s.Streamer.StreamSearch(ctx, q, opts, zoekt.SenderFunc(func(event *zoekt.SearchResult) {
69 stats.Add(event.Stats)
70 sender.Send(event)
71 }))
72}
73
74func (s *typeRepoSearcher) List(ctx context.Context, q query.Q, opts *zoekt.ListOptions) (rl *zoekt.RepoList, err error) {
75 tr, ctx := trace.New(ctx, "typeRepoSearcher.List", "")
76 tr.LazyLog(q, true)
77 tr.LazyPrintf("opts: %s", opts)
78 if tenant.EnforceTenant() {
79 tenant.Log(ctx, tr)
80 }
81 defer func() {
82 if rl != nil {
83 tr.LazyPrintf("repos.size=%d reposmap.size=%d crashes=%d stats=%+v", len(rl.Repos), len(rl.ReposMap), rl.Crashes, rl.Stats)
84 }
85 if err != nil {
86 tr.LazyPrintf("error: %v", err)
87 tr.SetError(err)
88 }
89 tr.Finish()
90 }()
91
92 q, err = s.eval(ctx, tr, q)
93 if err != nil {
94 return nil, err
95 }
96
97 return s.Streamer.List(ctx, q, opts)
98}
99
100func (s *typeRepoSearcher) eval(ctx context.Context, tr *trace.Trace, q query.Q) (query.Q, error) {
101 var err error
102 q = query.Map(q, func(q query.Q) query.Q {
103 if err != nil {
104 return nil
105 }
106
107 rq, ok := q.(*query.Type)
108 if !ok || rq.Type != query.TypeRepo {
109 return q
110 }
111
112 tr.LazyPrintf("evaluating sub-expression %s", rq)
113
114 var rl *zoekt.RepoList
115 rl, err = s.Streamer.List(ctx, rq.Child, nil)
116 if err != nil {
117 return nil
118 }
119
120 rs := &query.RepoSet{Set: make(map[string]bool, len(rl.Repos))}
121 for _, r := range rl.Repos {
122 rs.Set[r.Repository.Name] = true
123 }
124
125 tr.LazyPrintf("replaced sub-expression with %s", rs)
126
127 return rs
128 })
129 return q, err
130}