+5

internal/tenant/context.go

Reviewed

··· 10 10 11 11 "github.com/sourcegraph/zoekt/internal/tenant/internal/enforcement" 12 12 "github.com/sourcegraph/zoekt/internal/tenant/internal/tenanttype" 13 13 + "github.com/sourcegraph/zoekt/internal/tenant/systemtenant" 13 14 "github.com/sourcegraph/zoekt/trace" 14 15 ) 15 16 ··· 26 27 // Log logs the tenant ID to the trace. If tenant logging is enabled, it also 27 28 // logs a stack trace to a pprof profile. 28 29 func Log(ctx context.Context, tr *trace.Trace) { 30 30 + if systemtenant.Is(ctx) { 31 31 + tr.LazyPrintf("tenant: system") 32 32 + return 33 33 + } 29 34 tnt, ok := tenanttype.GetTenant(ctx) 30 35 if !ok { 31 36 if profile := pprofMissingTenant(); profile != nil {

+5 -3

internal/tenant/systemtenant/systemtenant.go

Reviewed

··· 10 10 11 11 const systemTenantKey contextKey = iota 12 12 13 13 - // UnsafeCtx is a context that allows queries across all tenants. Don't use this 14 14 - // for user requests. 15 15 - var UnsafeCtx = context.WithValue(context.Background(), systemTenantKey, systemTenantKey) 13 13 + // WithUnsafeContext taints the context to allow queries across all tenants. 14 14 + // Never use this for user requests. 15 15 + func WithUnsafeContext(ctx context.Context) context.Context { 16 16 + return context.WithValue(ctx, systemTenantKey, systemTenantKey) 17 17 + } 16 18 17 19 // Is returns true if the context has been marked to allow queries across all 18 20 // tenants.

+2 -2

internal/tenant/systemtenant/systemtenant_test.go

Reviewed

··· 7 7 "github.com/stretchr/testify/require" 8 8 ) 9 9 10 10 - func TestSystemtenantRoundtrip(t *testing.T) { 10 10 + func TestSystemTenantRoundTrip(t *testing.T) { 11 11 if Is(context.Background()) { 12 12 t.Fatal() 13 13 } 14 14 - require.True(t, Is(UnsafeCtx)) 14 14 + require.True(t, Is(WithUnsafeContext(context.Background()))) 15 15 }

+2 -2

shards/shards.go

Reviewed

··· 1083 1083 1084 1084 func mkRankedShard(s zoekt.Searcher) *rankedShard { 1085 1085 q := query.Const{Value: true} 1086 1086 - // We need to use UnsafeCtx here, otherwise we cannot return a proper 1086 1086 + // We need to use WithUnsafeContext here, otherwise we cannot return a proper 1087 1087 // rankedShard. On the user request path we use selectRepoSet which relies on 1088 1088 // rankedShard.repos being set. 1089 1089 - result, err := s.List(systemtenant.UnsafeCtx, &q, nil) 1089 1089 + result, err := s.List(systemtenant.WithUnsafeContext(context.Background()), &q, nil) 1090 1090 if err != nil { 1091 1091 log.Printf("[ERROR] mkRankedShard(%s): failed to cache repository list: %v", s, err) 1092 1092 return &rankedShard{Searcher: s}

+6 -1

web/server.go

Reviewed

··· 32 32 "time" 33 33 34 34 "github.com/grafana/regexp" 35 35 + 35 36 "github.com/sourcegraph/zoekt" 37 37 + "github.com/sourcegraph/zoekt/internal/tenant/systemtenant" 36 38 zjson "github.com/sourcegraph/zoekt/json" 37 39 "github.com/sourcegraph/zoekt/query" 38 40 ) ··· 206 208 q := &query.Const{Value: true} 207 209 opts := &zoekt.SearchOptions{ShardMaxMatchCount: 1, TotalMaxMatchCount: 1, MaxDocDisplayCount: 1} 208 210 209 209 - result, err := s.Searcher.Search(r.Context(), q, opts) 211 211 + // We need to use WithUnsafeContext here because we want to perform a full 212 212 + // search returning results. The result of this search is not used for anything 213 213 + // other than determining if the server is healthy. 214 214 + result, err := s.Searcher.Search(systemtenant.WithUnsafeContext(r.Context()), q, opts) 210 215 if err != nil { 211 216 http.Error(w, fmt.Sprintf("not ready: %v", err), http.StatusInternalServerError) 212 217 return