fork of https://github.com/sourcegraph/zoekt
Add RepoIds filter to /api/search based on new RepoIds query primitive (#521)
This is a new parameter that can be passed in the request body to filter
searches to a list of repositories by `repoid`. This makes use of a new
`RepoIds` query filter.
author
Dylan
committer
GitHub
date
(Jan 30, 2023, 3:50 PM +0200)
commit
e42e25e7
e42e25e7453727853ace1df54844d3dcb6885ce2
parent
989b9b1d
989b9b1d7d48296153d3b2785ae219319fa14556
+198
-7
+4
eval.go
+4
eval.go
···
83
83
return d.simplifyMultiRepo(q, func(repo *Repository) bool {
84
84
return r.Set[repo.Name]
85
85
})
86
+
case *query.RepoIDs:
87
+
return d.simplifyMultiRepo(q, func(repo *Repository) bool {
88
+
return r.Repos.Contains(repo.ID)
89
+
})
86
90
case *query.Language:
87
91
_, has := d.metaData.LanguageMap[r.Language]
88
92
if !has && d.metaData.IndexFeatureVersion < 12 {
+24
eval_test.go
+24
eval_test.go
···
212
212
}
213
213
}
214
214
215
+
func TestSimplifyRepoIDs(t *testing.T) {
216
+
d := compoundReposShard(t, "foo", "bar")
217
+
all := &query.RepoIDs{Repos: roaring.BitmapOf(hash("foo"), hash("bar"))}
218
+
some := &query.RepoIDs{Repos: roaring.BitmapOf(hash("foo"), hash("banana"))}
219
+
none := &query.RepoIDs{Repos: roaring.BitmapOf(hash("banana"))}
220
+
221
+
tr := cmp.Transformer("", func(b *roaring.Bitmap) []uint32 { return b.ToArray() })
222
+
223
+
got := d.simplify(all)
224
+
if d := cmp.Diff(&query.Const{Value: true}, got, tr); d != "" {
225
+
t.Fatalf("-want, +got:\n%s", d)
226
+
}
227
+
228
+
got = d.simplify(some)
229
+
if d := cmp.Diff(some, got, tr); d != "" {
230
+
t.Fatalf("-want, +got:\n%s", d)
231
+
}
232
+
233
+
got = d.simplify(none)
234
+
if d := cmp.Diff(&query.Const{Value: false}, got); d != "" {
235
+
t.Fatalf("-want, +got:\n%s", d)
236
+
}
237
+
}
238
+
215
239
func TestSimplifyRepo(t *testing.T) {
216
240
re := func(pat string) *query.Repo {
217
241
t.Helper()
+10
-5
json/json.go
+10
-5
json/json.go
···
27
27
}
28
28
29
29
type jsonSearchArgs struct {
30
-
Q string
31
-
Opts *zoekt.SearchOptions
30
+
Q string
31
+
RepoIDs *[]uint32
32
+
Opts *zoekt.SearchOptions
32
33
}
33
34
34
35
type jsonSearchReply struct {
···
67
68
searchArgs.Opts = &zoekt.SearchOptions{}
68
69
}
69
70
70
-
query, err := query.Parse(searchArgs.Q)
71
+
q, err := query.Parse(searchArgs.Q)
71
72
if err != nil {
72
73
jsonError(w, http.StatusBadRequest, err.Error())
73
74
return
74
75
}
75
76
77
+
if searchArgs.RepoIDs != nil {
78
+
q = query.NewAnd(q, query.NewRepoIDs(*searchArgs.RepoIDs...))
79
+
}
80
+
76
81
// Set a timeout if the user hasn't specified one.
77
82
if searchArgs.Opts.MaxWallTime == 0 {
78
83
var cancel context.CancelFunc
···
80
85
defer cancel()
81
86
}
82
87
83
-
if err := CalculateDefaultSearchLimits(ctx, query, s.Searcher, searchArgs.Opts); err != nil {
88
+
if err := CalculateDefaultSearchLimits(ctx, q, s.Searcher, searchArgs.Opts); err != nil {
84
89
jsonError(w, http.StatusInternalServerError, err.Error())
85
90
return
86
91
}
87
92
88
-
searchResult, err := s.Searcher.Search(ctx, query, searchArgs.Opts)
93
+
searchResult, err := s.Searcher.Search(ctx, q, searchArgs.Opts)
89
94
if err != nil {
90
95
jsonError(w, http.StatusInternalServerError, err.Error())
91
96
return
+74
json/json_test.go
+74
json/json_test.go
···
87
87
}
88
88
}
89
89
90
+
func TestClientServerWithRepoIDsProvided(t *testing.T) {
91
+
searchQuery := "hello"
92
+
expectedSearch := mustParse(searchQuery)
93
+
expectedSearch = query.NewAnd(expectedSearch, query.NewRepoIDs(1, 3, 5, 7))
94
+
mock := &mockSearcher.MockSearcher{
95
+
WantSearch: expectedSearch,
96
+
SearchResult: &zoekt.SearchResult{
97
+
Files: []zoekt.FileMatch{
98
+
{FileName: "bin.go"},
99
+
},
100
+
},
101
+
}
102
+
103
+
ts := httptest.NewServer(zjson.JSONServer(mock))
104
+
defer ts.Close()
105
+
106
+
searchBody := "{\"Q\":\"hello\",\"RepoIDs\":[1,3,5,7]}"
107
+
108
+
r, err := http.Post(ts.URL+"/search", "application/json", bytes.NewBufferString(searchBody))
109
+
if err != nil {
110
+
t.Fatal(err)
111
+
}
112
+
if r.StatusCode != 200 {
113
+
body, _ := io.ReadAll(r.Body)
114
+
t.Fatalf("Got status code %d, err %s", r.StatusCode, string(body))
115
+
}
116
+
117
+
var searchResult struct{ Result *zoekt.SearchResult }
118
+
err = json.NewDecoder(r.Body).Decode(&searchResult)
119
+
if err != nil {
120
+
t.Fatal(err)
121
+
}
122
+
if !reflect.DeepEqual(searchResult.Result, mock.SearchResult) {
123
+
t.Fatalf("\na %+v\nb %+v", searchResult.Result, mock.SearchResult)
124
+
}
125
+
}
126
+
127
+
func TestClientServerWithEmptyRepoIDsProvided(t *testing.T) {
128
+
searchQuery := "hello"
129
+
expectedSearch := mustParse(searchQuery)
130
+
expectedSearch = query.NewAnd(expectedSearch, query.NewRepoIDs())
131
+
mock := &mockSearcher.MockSearcher{
132
+
WantSearch: expectedSearch,
133
+
SearchResult: &zoekt.SearchResult{
134
+
Files: []zoekt.FileMatch{
135
+
{FileName: "bin.go"},
136
+
},
137
+
},
138
+
}
139
+
140
+
ts := httptest.NewServer(zjson.JSONServer(mock))
141
+
defer ts.Close()
142
+
143
+
searchBody := "{\"Q\":\"hello\",\"RepoIDs\":[]}"
144
+
145
+
r, err := http.Post(ts.URL+"/search", "application/json", bytes.NewBufferString(searchBody))
146
+
if err != nil {
147
+
t.Fatal(err)
148
+
}
149
+
if r.StatusCode != 200 {
150
+
body, _ := io.ReadAll(r.Body)
151
+
t.Fatalf("Got status code %d, err %s", r.StatusCode, string(body))
152
+
}
153
+
154
+
var searchResult struct{ Result *zoekt.SearchResult }
155
+
err = json.NewDecoder(r.Body).Decode(&searchResult)
156
+
if err != nil {
157
+
t.Fatal(err)
158
+
}
159
+
if !reflect.DeepEqual(searchResult.Result, mock.SearchResult) {
160
+
t.Fatalf("\na %+v\nb %+v", searchResult.Result, mock.SearchResult)
161
+
}
162
+
}
163
+
90
164
func TestProgressNotEncodedInSearch(t *testing.T) {
91
165
searchQuery := "hello"
92
166
mock := &mockSearcher.MockSearcher{
+15
matchtree.go
+15
matchtree.go
···
965
965
},
966
966
}, nil
967
967
968
+
case *query.RepoIDs:
969
+
reposWant := make([]bool, len(d.repoMetaData))
970
+
for repoIdx, r := range d.repoMetaData {
971
+
if s.Repos.Contains(r.ID) {
972
+
reposWant[repoIdx] = true
973
+
}
974
+
}
975
+
return &docMatchTree{
976
+
reason: "RepoIDs",
977
+
numDocs: d.numDocs(),
978
+
predicate: func(docID uint32) bool {
979
+
return reposWant[d.repos[docID]]
980
+
},
981
+
}, nil
982
+
968
983
case *query.Repo:
969
984
reposWant := make([]bool, len(d.repoMetaData))
970
985
for repoIdx, r := range d.repoMetaData {
+23
matchtree_test.go
+23
matchtree_test.go
···
306
306
t.Fatalf("expect %d documents, but got at least 1 more", len(want))
307
307
}
308
308
}
309
+
310
+
func TestRepoIDs(t *testing.T) {
311
+
d := &indexData{
312
+
repoMetaData: []Repository{{Name: "r0", ID: 0}, {Name: "r1", ID: 1}, {Name: "r2", ID: 2}, {Name: "r3", ID: 3}},
313
+
fileBranchMasks: []uint64{1, 1, 1, 1, 1, 1},
314
+
repos: []uint16{0, 0, 1, 2, 3, 3},
315
+
}
316
+
mt, err := d.newMatchTree(&query.RepoIDs{Repos: roaring.BitmapOf(1, 3, 99)})
317
+
if err != nil {
318
+
t.Fatal(err)
319
+
}
320
+
want := []uint32{2, 4, 5}
321
+
for i := 0; i < len(want); i++ {
322
+
nextDoc := mt.nextDoc()
323
+
if nextDoc != want[i] {
324
+
t.Fatalf("want %d, got %d", want[i], nextDoc)
325
+
}
326
+
mt.prepare(nextDoc)
327
+
}
328
+
if mt.nextDoc() != maxUInt32 {
329
+
t.Fatalf("expected %d document, but got at least 1 more", len(want))
330
+
}
331
+
}
+27
query/query.go
+27
query/query.go
···
231
231
return sb.String()
232
232
}
233
233
234
+
// NewRepoIDs is a helper for creating a RepoIDs which
235
+
// searches only the matched repos.
236
+
func NewRepoIDs(ids ...uint32) *RepoIDs {
237
+
return &RepoIDs{Repos: roaring.BitmapOf(ids...)}
238
+
}
239
+
240
+
func (q *RepoIDs) String() string {
241
+
var sb strings.Builder
242
+
243
+
sb.WriteString("(repoids ")
244
+
245
+
if size := q.Repos.GetCardinality(); size > 1 {
246
+
sb.WriteString("count:" + strconv.FormatUint(size, 10))
247
+
} else {
248
+
sb.WriteString("repoid=" + q.Repos.String())
249
+
}
250
+
251
+
sb.WriteString(")")
252
+
return sb.String()
253
+
}
254
+
234
255
// MarshalBinary implements a specialized encoder for BranchesRepos.
235
256
func (q BranchesRepos) MarshalBinary() ([]byte, error) {
236
257
return branchesReposEncode(q.List)
···
247
268
type BranchRepos struct {
248
269
Branch string
249
270
Repos *roaring.Bitmap
271
+
}
272
+
273
+
// Similar to BranchRepos but will be used to match only by repoid and
274
+
// therefore matches all branches
275
+
type RepoIDs struct {
276
+
Repos *roaring.Bitmap
250
277
}
251
278
252
279
// RepoSet is a list of repos to match. It is a Sourcegraph addition and only
+1
rpc/rpc.go
+1
rpc/rpc.go
+9
shards/shards.go
+9
shards/shards.go
···
388
388
hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool {
389
389
return setQuery.Set[repo.Name]
390
390
})
391
+
case *query.RepoIDs:
392
+
setSize = int(setQuery.Repos.GetCardinality())
393
+
hasRepos = hasReposForPredicate(func(repo *zoekt.Repository) bool {
394
+
return setQuery.Repos.Contains(repo.ID)
395
+
})
391
396
case *query.BranchesRepos:
392
397
for _, br := range setQuery.List {
393
398
setSize += int(br.Repos.GetCardinality())
···
442
447
// (content baz). This work can be done now once, rather than per shard.
443
448
switch c := c.(type) {
444
449
case *query.RepoSet:
450
+
and.Children[i] = &query.Const{Value: true}
451
+
return filtered, query.Simplify(and)
452
+
453
+
case *query.RepoIDs:
445
454
and.Children[i] = &query.Const{Value: true}
446
455
return filtered, query.Simplify(and)
447
456
+11
-2
shards/shards_test.go
+11
-2
shards/shards_test.go
···
226
226
}
227
227
}
228
228
229
-
func TestFilteringShardsByRepoSet(t *testing.T) {
229
+
func TestFilteringShardsByRepoSetOrBranchesReposOrRepoIDs(t *testing.T) {
230
230
ss := newShardedSearcher(1)
231
231
232
232
repoSetNames := []string{}
233
+
repoIDs := []uint32{}
233
234
n := 10 * runtime.GOMAXPROCS(0)
234
235
for i := 0; i < n; i++ {
235
236
shardName := fmt.Sprintf("shard%d", i)
236
237
repoName := fmt.Sprintf("repository%.3d", i)
238
+
repoID := hash(repoName)
237
239
238
240
if i%3 == 0 {
239
241
repoSetNames = append(repoSetNames, repoName)
242
+
repoIDs = append(repoIDs, repoID)
240
243
}
241
244
242
245
ss.replace(map[string]zoekt.Searcher{
243
246
shardName: &rankSearcher{
244
-
repo: &zoekt.Repository{ID: hash(repoName), Name: repoName},
247
+
repo: &zoekt.Repository{ID: repoID, Name: repoName},
245
248
rank: uint16(n - i),
246
249
},
247
250
})
···
266
269
set := query.NewRepoSet(repoSetNames...)
267
270
sub := &query.Substring{Pattern: "bla"}
268
271
272
+
repoIDsQuery := query.NewRepoIDs(repoIDs...)
273
+
269
274
queries := []query.Q{
270
275
query.NewAnd(set, sub),
271
276
// Test with the same reposet again
···
274
279
query.NewAnd(branchesRepos, sub),
275
280
// Test with the same repoBranches with IDs again
276
281
query.NewAnd(branchesRepos, sub),
282
+
283
+
query.NewAnd(repoIDsQuery, sub),
284
+
// Test with the same repoIDs again
285
+
query.NewAnd(repoIDsQuery, sub),
277
286
}
278
287
279
288
for _, q := range queries {