func setupBenchmark(nodes int, pods int) *topdown.QueryParams {
// policy compilation
c := ast.NewCompiler()
modules := map[string]*ast.Module{
"test": ast.MustParseModule(policy),
}
if c.Compile(modules); c.Failed() {
panic(c.Errors)
}
// storage setup
store := storage.New(storage.InMemoryConfig())
// parameter setup
ctx := context.Background()
request := ast.ObjectTerm(ast.Item(ast.StringTerm("pod"), ast.MustParseTerm(requestedPod)))
path := ast.MustParseRef("data.opa.test.scheduler.fit")
txn := storage.NewTransactionOrDie(ctx, store)
params := topdown.NewQueryParams(ctx, c, store, txn, request.Value, path)
// data setup
setupNodes(ctx, store, txn, nodes)
setupRCs(ctx, store, txn, 1)
setupPods(ctx, store, txn, pods, nodes)
return params
}
func setup(t *testing.T, filename string) *topdown.QueryParams {
// policy compilation
c := ast.NewCompiler()
modules := map[string]*ast.Module{
"test": ast.MustParseModule(policy),
}
if c.Compile(modules); c.Failed() {
t.Fatal("unexpected error:", c.Errors)
}
// storage setup
store := storage.New(storage.Config{
Builtin: loadDataStore(filename),
})
// parameter setup
ctx := context.Background()
request := ast.ObjectTerm(ast.Item(ast.StringTerm("pod"), ast.MustParseTerm(requestedPod)))
path := ast.MustParseRef("data.opa.test.scheduler.fit")
txn := storage.NewTransactionOrDie(ctx, store)
params := topdown.NewQueryParams(ctx, c, store, txn, request.Value, path)
return params
}
// Ref returns a ref that represents p rooted at head.
func (p Path) Ref(head *ast.Term) (ref ast.Ref) {
ref = make(ast.Ref, len(p)+1)
ref[0] = head
for i := range p {
idx, err := strconv.ParseInt(p[i], 10, 64)
if err == nil {
ref[i+1] = ast.IntNumberTerm(int(idx))
} else {
ref[i+1] = ast.StringTerm(p[i])
}
}
return ref
}
func stringPathToRef(s string) (r ast.Ref) {
if len(s) == 0 {
return r
}
p := strings.Split(s, "/")
for _, x := range p {
if x == "" {
continue
}
i, err := strconv.Atoi(x)
if err != nil {
r = append(r, ast.StringTerm(x))
} else {
r = append(r, ast.IntNumberTerm(i))
}
}
return r
}
func iterStorage(ctx context.Context, store Store, txn Transaction, nonGround, ground ast.Ref, bindings *ast.ValueMap, iter func(*ast.ValueMap, interface{})) error {
if len(nonGround) == 0 {
path, err := NewPathForRef(ground)
if err != nil {
return err
}
node, err := store.Read(ctx, txn, path)
if err != nil {
if IsNotFound(err) {
return nil
}
return err
}
iter(bindings, node)
return nil
}
head := nonGround[0]
tail := nonGround[1:]
headVar, isVar := head.Value.(ast.Var)
if !isVar || len(ground) == 0 {
ground = append(ground, head)
return iterStorage(ctx, store, txn, tail, ground, bindings, iter)
}
path, err := NewPathForRef(ground)
if err != nil {
return err
}
node, err := store.Read(ctx, txn, path)
if err != nil {
if IsNotFound(err) {
return nil
}
return err
}
switch node := node.(type) {
case map[string]interface{}:
for key := range node {
ground = append(ground, ast.StringTerm(key))
cpy := bindings.Copy()
cpy.Put(headVar, ast.String(key))
err := iterStorage(ctx, store, txn, tail, ground, cpy, iter)
if err != nil {
return err
}
ground = ground[:len(ground)-1]
}
case []interface{}:
for i := range node {
idx := ast.IntNumberTerm(i)
ground = append(ground, idx)
cpy := bindings.Copy()
cpy.Put(headVar, idx.Value)
err := iterStorage(ctx, store, txn, tail, ground, cpy, iter)
if err != nil {
return err
}
ground = ground[:len(ground)-1]
}
}
return nil
}
