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/query"
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 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, 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, 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", len(rl.Repos))
84 tr.LazyPrintf("reposmap size: %d", len(rl.ReposMap))
85 tr.LazyPrintf("crashes: %d", rl.Crashes)
86 tr.LazyPrintf("stats: %+v", rl.Stats)
87 }
88 if err != nil {
89 tr.LazyPrintf("error: %v", err)
90 tr.SetError(err)
91 }
92 tr.Finish()
93 }()
94
95 q, err = s.eval(ctx, q)
96 if err != nil {
97 return nil, err
98 }
99
100 return s.Streamer.List(ctx, q, opts)
101}
102
103func (s *typeRepoSearcher) eval(ctx context.Context, q query.Q) (query.Q, error) {
104 var err error
105 q = query.Map(q, func(q query.Q) query.Q {
106 if err != nil {
107 return nil
108 }
109
110 rq, ok := q.(*query.Type)
111 if !ok || rq.Type != query.TypeRepo {
112 return q
113 }
114
115 var rl *zoekt.RepoList
116 rl, err = s.Streamer.List(ctx, rq.Child, nil)
117 if err != nil {
118 return nil
119 }
120
121 rs := &query.RepoSet{Set: make(map[string]bool, len(rl.Repos))}
122 for _, r := range rl.Repos {
123 rs.Set[r.Repository.Name] = true
124 }
125 return rs
126 })
127 return q, err
128}