internal: Refactor load/store/compile implementation

This commit refactors the load/store/compile implementation that used
to live inside the runtime package. Specifically:

* Move init-time file loading logic into separate internal package
  (initload) along with store/compile logic. Add tests around
  load/store/compile that don't require the entire Runtime object.
  This also avoids duplication of the "version overwriting" logic.

* Move store/compile calls into the manager. This avoids the need for
  two compile operations on startup.

Signed-off-by: Torin Sandall <torinsandall@gmail.com>

internal/runtime/init/init.go

@@ -0,0 +1,118 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+// Package init is an internal package with helpers for data and policy loading during initialization.
+package init
+
+import (
+	"context"
+
+	"github.com/pkg/errors"
+
+	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
+	storedversion "github.com/open-policy-agent/opa/internal/version"
+	"github.com/open-policy-agent/opa/loader"
+	"github.com/open-policy-agent/opa/metrics"
+	"github.com/open-policy-agent/opa/storage"
+)
+
+// InsertAndCompileOptions contains the input for the operation.
+type InsertAndCompileOptions struct {
+	Store     storage.Store
+	Txn       storage.Transaction
+	Files     loader.Result
+	Bundles   map[string]*bundle.Bundle
+	MaxErrors int
+}
+
+// InsertAndCompileResult contains the output of the operation.
+type InsertAndCompileResult struct {
+	Compiler *ast.Compiler
+	Metrics  metrics.Metrics
+}
+
+// InsertAndCompile writes data and policy into the store and returns a compiler for the
+// store contents.
+func InsertAndCompile(ctx context.Context, opts InsertAndCompileOptions) (*InsertAndCompileResult, error) {
+
+	if len(opts.Files.Documents) > 0 {
+		if err := opts.Store.Write(ctx, opts.Txn, storage.AddOp, storage.Path{}, opts.Files.Documents); err != nil {
+			return nil, errors.Wrap(err, "storage error")
+		}
+	}
+
+	policies := make(map[string]*ast.Module, len(opts.Files.Modules))
+
+	for id, parsed := range opts.Files.Modules {
+		policies[id] = parsed.Parsed
+	}
+
+	compiler := ast.NewCompiler().SetErrorLimit(opts.MaxErrors).WithPathConflictsCheck(storage.NonEmpty(ctx, opts.Store, opts.Txn))
+	m := metrics.New()
+
+	activation := &bundle.ActivateOpts{
+		Ctx:          ctx,
+		Store:        opts.Store,
+		Txn:          opts.Txn,
+		Compiler:     compiler,
+		Metrics:      m,
+		Bundles:      opts.Bundles,
+		ExtraModules: policies,
+	}
+
+	err := bundle.Activate(activation)
+	if err != nil {
+		return nil, err
+	}
+
+	// Policies in bundles will have already been added to the store, but
+	// modules loaded outside of bundles will need to be added manually.
+	for id, parsed := range opts.Files.Modules {
+		if err := opts.Store.UpsertPolicy(ctx, opts.Txn, id, parsed.Raw); err != nil {
+			return nil, errors.Wrap(err, "storage error")
+		}
+	}
+
+	// Set the version in the store last to prevent data files from overwriting.
+	if err := storedversion.Write(ctx, opts.Store, opts.Txn); err != nil {
+		return nil, errors.Wrap(err, "storage error")
+	}
+
+	return &InsertAndCompileResult{Compiler: compiler, Metrics: m}, nil
+}
+
+// LoadPathsResult contains the output loading a set of paths.
+type LoadPathsResult struct {
+	Bundles map[string]*bundle.Bundle
+	Files   loader.Result
+}
+
+// LoadPaths reads data and policy from the given paths and returns a set of bundles or
+// raw loader file results.
+func LoadPaths(paths []string, filter loader.Filter, asBundle bool) (*LoadPathsResult, error) {
+
+	var result LoadPathsResult
+	var err error
+
+	if asBundle {
+		result.Bundles = make(map[string]*bundle.Bundle, len(paths))
+		for _, path := range paths {
+			result.Bundles[path], err = loader.NewFileLoader().AsBundle(path)
+			if err != nil {
+				return nil, err
+			}
+		}
+		return &result, nil
+	}
+
+	files, err := loader.NewFileLoader().Filtered(paths, filter)
+	if err != nil {
+		return nil, err
+	}
+
+	result.Files = *files
+
+	return &result, nil
+}

internal/runtime/init/init_test.go

@@ -0,0 +1,179 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+
+package init
+
+import (
+	"context"
+	"path/filepath"
+	"testing"
+
+	"github.com/open-policy-agent/opa/storage"
+	"github.com/open-policy-agent/opa/storage/inmem"
+	"github.com/open-policy-agent/opa/util"
+	"github.com/open-policy-agent/opa/util/test"
+	"github.com/open-policy-agent/opa/version"
+)
+
+func TestInit(t *testing.T) {
+
+	mod1 := `package a.b.c
+
+import data.a.foo
+
+p = true { foo = "bar" }
+p = true { 1 = 2 }`
+
+	mod2 := `package b.c.d
+
+import data.b.foo
+
+p = true { foo = "bar" }
+p = true { 1 = 2 }`
+
+	tests := []struct {
+		note         string
+		fs           map[string]string
+		loadParams   []string
+		expectedData map[string]string
+		expectedMods []string
+		asBundle     bool
+	}{
+		{
+			note: "load files",
+			fs: map[string]string{
+				"datafile":   `{"foo": "bar", "x": {"y": {"z": [1]}}}`,
+				"policyFile": mod1,
+			},
+			loadParams: []string{"datafile", "policyFile"},
+			expectedData: map[string]string{
+				"/foo": "bar",
+			},
+			expectedMods: []string{mod1},
+			asBundle:     false,
+		},
+		{
+			note: "load bundle",
+			fs: map[string]string{
+				"datafile":    `{"foo": "bar", "x": {"y": {"z": [1]}}}`, // Should be ignored
+				"data.json":   `{"foo": "not-bar"}`,
+				"policy.rego": mod1,
+			},
+			loadParams: []string{"/"},
+			expectedData: map[string]string{
+				"/foo": "not-bar",
+			},
+			expectedMods: []string{mod1},
+			asBundle:     true,
+		},
+		{
+			note: "load multiple bundles",
+			fs: map[string]string{
+				"/bundle1/a/data.json":   `{"foo": "bar1", "x": {"y": {"z": [1]}}}`, // Should be ignored
+				"/bundle1/a/policy.rego": mod1,
+				"/bundle1/a/.manifest":   `{"roots": ["a"]}`,
+				"/bundle2/b/data.json":   `{"foo": "bar2"}`,
+				"/bundle2/b/policy.rego": mod2,
+				"/bundle2/b/.manifest":   `{"roots": ["b"]}`,
+			},
+			loadParams: []string{"bundle1", "bundle2"},
+			expectedData: map[string]string{
+				"/a/foo": "bar1",
+				"/b/foo": "bar2",
+			},
+			expectedMods: []string{mod1, mod2},
+			asBundle:     true,
+		},
+		{
+			note: "preserve OPA version",
+			fs: map[string]string{
+				"/root/system/version/data.json": `{"version": "XYZ"}`, // Should be overwritten
+			},
+			loadParams: []string{"root"},
+			expectedData: map[string]string{
+				"/system/version/version": version.Version,
+			},
+			asBundle: true,
+		},
+	}
+
+	ctx := context.Background()
+
+	for _, tc := range tests {
+		t.Run(tc.note, func(t *testing.T) {
+			test.WithTempFS(tc.fs, func(rootDir string) {
+
+				paths := []string{}
+
+				for _, fileName := range tc.loadParams {
+					paths = append(paths, filepath.Join(rootDir, fileName))
+				}
+
+				// Create a new store and perform a file load/insert sequence.
+				store := inmem.New()
+
+				err := storage.Txn(ctx, store, storage.WriteParams, func(txn storage.Transaction) error {
+
+					loaded, err := LoadPaths(paths, nil, tc.asBundle)
+					if err != nil {
+						return err
+					}
+
+					_, err = InsertAndCompile(ctx, InsertAndCompileOptions{
+						Store:     store,
+						Txn:       txn,
+						Files:     loaded.Files,
+						Bundles:   loaded.Bundles,
+						MaxErrors: -1,
+					})
+
+					return err
+				})
+
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				// Verify the loading was successful as expected.
+				txn := storage.NewTransactionOrDie(ctx, store)
+
+				for storePath, expected := range tc.expectedData {
+					node, err := store.Read(ctx, txn, storage.MustParsePath(storePath))
+					if util.Compare(node, expected) != 0 || err != nil {
+						t.Fatalf("Expected %v but got %v (err: %v)", expected, node, err)
+					}
+				}
+
+				ids, err := store.ListPolicies(ctx, txn)
+				if err != nil {
+					t.Fatal(err)
+				}
+
+				if len(tc.expectedMods) != len(ids) {
+					t.Fatalf("Expected %d modules, got %d", len(tc.expectedMods), len(ids))
+				}
+
+				actualMods := map[string]struct{}{}
+				for _, id := range ids {
+					result, err := store.GetPolicy(ctx, txn, id)
+					if err != nil {
+						t.Fatalf("Unexpected error: %s", err)
+					}
+					actualMods[string(result)] = struct{}{}
+				}
+
+				for _, expectedMod := range tc.expectedMods {
+					if _, found := actualMods[expectedMod]; !found {
+						t.Fatalf("Expected %v but got: %v", expectedMod, actualMods)
+					}
+				}
+
+				_, err = store.Read(ctx, txn, storage.MustParsePath("/system/version"))
+				if err != nil {
+					t.Fatal(err)
+				}
+			})
+		})
+	}
+}

plugins/plugins.go

@@ -11,7 +11,10 @@ import (
 	"sync"
 
 	"github.com/open-policy-agent/opa/ast"
+	"github.com/open-policy-agent/opa/bundle"
 	"github.com/open-policy-agent/opa/config"
+	initload "github.com/open-policy-agent/opa/internal/runtime/init"
+	"github.com/open-policy-agent/opa/loader"
 	"github.com/open-policy-agent/opa/plugins/rest"
 	"github.com/open-policy-agent/opa/storage"
 	"github.com/open-policy-agent/opa/util"
@@ -129,6 +132,10 @@ type Manager struct {
 	mtx                   sync.Mutex
 	pluginStatus          map[string]*Status
 	pluginStatusListeners map[string]StatusListener
+	initBundles           map[string]*bundle.Bundle
+	initFiles             loader.Result
+	maxErrors             int
+	initialized           bool
 }
 
 type managerContextKey string
@@ -165,6 +172,27 @@ func Info(term *ast.Term) func(*Manager) {
 	}
 }
 
+// InitBundles provides the initial set of bundles to load.
+func InitBundles(b map[string]*bundle.Bundle) func(*Manager) {
+	return func(m *Manager) {
+		m.initBundles = b
+	}
+}
+
+// InitFiles provides the initial set of other data/policy files to load.
+func InitFiles(f loader.Result) func(*Manager) {
+	return func(m *Manager) {
+		m.initFiles = f
+	}
+}
+
+// MaxErrors sets the error limit for the manager's shared compiler.
+func MaxErrors(n int) func(*Manager) {
+	return func(m *Manager) {
+		m.maxErrors = n
+	}
+}
+
 // New creates a new Manager using config.
 func New(raw []byte, id string, store storage.Store, opts ...func(*Manager)) (*Manager, error) {
 
@@ -185,6 +213,7 @@ func New(raw []byte, id string, store storage.Store, opts ...func(*Manager)) (*M
 		services:              services,
 		pluginStatus:          map[string]*Status{},
 		pluginStatusListeners: map[string]StatusListener{},
+		maxErrors:             -1,
 	}
 
 	for _, f := range opts {
@@ -194,6 +223,46 @@ func New(raw []byte, id string, store storage.Store, opts ...func(*Manager)) (*M
 	return m, nil
 }
 
+// Init returns an error if the manager could not initialize itself. Init() should
+// be called before Start(). Init() is idempotent.
+func (m *Manager) Init(ctx context.Context) error {
+
+	if m.initialized {
+		return nil
+	}
+
+	params := storage.TransactionParams{
+		Write:   true,
+		Context: storage.NewContext(),
+	}
+
+	err := storage.Txn(ctx, m.Store, params, func(txn storage.Transaction) error {
+
+		result, err := initload.InsertAndCompile(ctx, initload.InsertAndCompileOptions{
+			Store:     m.Store,
+			Txn:       txn,
+			Files:     m.initFiles,
+			Bundles:   m.initBundles,
+			MaxErrors: m.maxErrors,
+		})
+
+		if err != nil {
+			return err
+		}
+
+		SetCompilerOnContext(params.Context, result.Compiler)
+		_, err = m.Store.Register(ctx, txn, storage.TriggerConfig{OnCommit: m.onCommit})
+		return err
+	})
+
+	if err != nil {
+		return err
+	}
+
+	m.initialized = true
+	return nil
+}
+
 // Labels returns the set of labels from the configuration.
 func (m *Manager) Labels() map[string]string {
 	m.mtx.Lock()
@@ -259,27 +328,17 @@ func (m *Manager) RegisterCompilerTrigger(f func(txn storage.Transaction)) {
 	m.registeredTriggers = append(m.registeredTriggers, f)
 }
 
-// Start starts the manager.
+// Start starts the manager. Init() should be called once before Start().
 func (m *Manager) Start(ctx context.Context) error {
 
 	if m == nil {
 		return nil
 	}
 
-	if err := storage.Txn(ctx, m.Store, storage.WriteParams, func(txn storage.Transaction) error {
-
-		c, err := loadCompilerFromStore(ctx, m.Store, txn)
-		if err != nil {
+	if !m.initialized {
+		if err := m.Init(ctx); err != nil {
 			return err
 		}
-
-		m.setCompiler(c)
-
-		_, err = m.Store.Register(ctx, txn, storage.TriggerConfig{OnCommit: m.onCommit})
-		return err
-
-	}); err != nil {
-		return err
 	}
 
 	var toStart []Plugin
@@ -402,20 +461,18 @@ func (m *Manager) copyPluginStatus() map[string]*Status {
 
 func (m *Manager) onCommit(ctx context.Context, txn storage.Transaction, event storage.TriggerEvent) {
 
-	if event.PolicyChanged() {
+	compiler := GetCompilerOnContext(event.Context)
 
-		var compiler *ast.Compiler
-
-		// If the context does not contain the compiler fallback to loading the
-		// compiler from the store. Currently the bundle plugin sets the
-		// compiler on the context but the server does not (nor would users
-		// implementing their own policy loading.)
-		if compiler = GetCompilerOnContext(event.Context); compiler == nil {
-			compiler, _ = loadCompilerFromStore(ctx, m.Store, txn)
-		}
+	// If the context does not contain the compiler fallback to loading the
+	// compiler from the store. Currently the bundle plugin sets the
+	// compiler on the context but the server does not (nor would users
+	// implementing their own policy loading.)
+	if compiler == nil && event.PolicyChanged() {
+		compiler, _ = loadCompilerFromStore(ctx, m.Store, txn)
+	}
 
+	if compiler != nil {
 		m.setCompiler(compiler)
-
 		for _, f := range m.registeredTriggers {
 			f(txn)
 		}

plugins/plugins_test.go

@@ -6,9 +6,11 @@ package plugins
 
 import (
 	"context"
+	"fmt"
 	"reflect"
 	"testing"
 
+	"github.com/open-policy-agent/opa/internal/storage/mock"
 	"github.com/open-policy-agent/opa/storage/inmem"
 )
 
@@ -114,3 +116,92 @@ func (p *testPlugin) Stop(ctx context.Context) {
 func (p *testPlugin) Reconfigure(ctx context.Context, config interface{}) {
 	p.m.UpdatePluginStatus("p1", &Status{State: StateNotReady})
 }
+
+func TestPluginManagerLazyInitBeforePluginStart(t *testing.T) {
+
+	m, err := New([]byte(`{"plugins": {"someplugin": {"enabled": true}}}`), "test", inmem.New())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	mock := &mockForInitStartOrdering{Manager: m}
+
+	m.Register("someplugin", mock)
+
+	if err := m.Start(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	if !mock.Started {
+		t.Fatal("expected plugin to be started")
+	}
+
+}
+
+func TestPluginManagerInitBeforePluginStart(t *testing.T) {
+
+	m, err := New([]byte(`{"plugins": {"someplugin": {}}}`), "test", inmem.New())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if err := m.Init(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	mock := &mockForInitStartOrdering{Manager: m}
+
+	m.Register("someplugin", mock)
+
+	if err := m.Start(context.Background()); err != nil {
+		t.Fatal(err)
+	}
+
+	if !mock.Started {
+		t.Fatal("expected plugin to be started")
+	}
+
+}
+
+func TestPluginManagerInitIdempotence(t *testing.T) {
+
+	mockStore := mock.New()
+
+	m, err := New([]byte(`{"plugins": {"someplugin": {}}}`), "test", mockStore)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ctx := context.Background()
+
+	if err := m.Init(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	exp := len(mockStore.Transactions)
+
+	if err := m.Init(ctx); err != nil {
+		t.Fatal(err)
+	}
+
+	if len(mockStore.Transactions) != exp {
+		t.Fatal("expected num txns to be:", exp, "but got:", len(mockStore.Transactions))
+	}
+
+}
+
+type mockForInitStartOrdering struct {
+	Manager *Manager
+	Started bool
+}
+
+func (m *mockForInitStartOrdering) Start(ctx context.Context) error {
+	m.Started = true
+	if m.Manager.initialized {
+		return nil
+	}
+	return fmt.Errorf("expected manager to be initialized")
+}
+
+func (m *mockForInitStartOrdering) Stop(ctx context.Context)                            { return }
+func (m *mockForInitStartOrdering) Reconfigure(ctx context.Context, config interface{}) { return }

runtime/runtime.go

@@ -18,8 +18,6 @@ import (
 	"syscall"
 	"time"
 
-	"github.com/open-policy-agent/opa/bundle"
-
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"gopkg.in/fsnotify.v1"
@@ -27,8 +25,8 @@ import (
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/internal/prometheus"
 	"github.com/open-policy-agent/opa/internal/runtime"
+	initload "github.com/open-policy-agent/opa/internal/runtime/init"
 	"github.com/open-policy-agent/opa/internal/uuid"
-	storedversion "github.com/open-policy-agent/opa/internal/version"
 	"github.com/open-policy-agent/opa/loader"
 	"github.com/open-policy-agent/opa/metrics"
 	"github.com/open-policy-agent/opa/plugins"
@@ -176,7 +174,8 @@ type Runtime struct {
 	metrics *prometheus.Provider
 }
 
-// NewRuntime returns a new Runtime object initialized with params.
+// NewRuntime returns a new Runtime object initialized with params. Clients must
+// call StartServer() or StartREPL() to start the runtime in either mode.
 func NewRuntime(ctx context.Context, params Params) (*Runtime, error) {
 
 	if params.ID == "" {
@@ -187,55 +186,28 @@ func NewRuntime(ctx context.Context, params Params) (*Runtime, error) {
 		}
 	}
 
-	loaded, err := loadPaths(params.Paths, params.Filter, params.BundleMode)
+	config, err := loadConfig(params)
+	if err != nil {
+		return nil, errors.Wrap(err, "config error")
+	}
+
+	loaded, err := initload.LoadPaths(params.Paths, params.Filter, params.BundleMode)
 	if err != nil {
 		return nil, errors.Wrap(err, "load error")
 	}
 
-	// TOOD(tsandall): All of this storage setup could be done by the plugin manager.
-	// This would avoid the need to parse and recompile modules provided on startup.
-	store := inmem.New()
-
-	txn, err := store.NewTransaction(ctx, storage.WriteParams)
+	info, err := runtime.Term(runtime.Params{Config: config})
 	if err != nil {
 		return nil, err
 	}
 
-	if len(loaded.Documents) > 0 {
-		if err := store.Write(ctx, txn, storage.AddOp, storage.Path{}, loaded.Documents); err != nil {
-			return nil, errors.Wrap(err, "storage error")
-		}
-	}
-
-	if err := compileAndStoreInputs(ctx, store, txn, loaded, params.ErrorLimit); err != nil {
-		store.Abort(ctx, txn)
-		return nil, errors.Wrap(err, "compile error")
-	}
-
-	// Write the version *after* any data loaded from files or bundles to
-	// avoid it being deleted.
-	if err := storedversion.Write(ctx, store, txn); err != nil {
-		store.Abort(ctx, txn)
-		return nil, errors.Wrap(err, "storage error")
-	}
-
-	if err := store.Commit(ctx, txn); err != nil {
-		return nil, errors.Wrap(err, "storage error")
-	}
-
-	bs, err := loadConfig(params)
+	manager, err := plugins.New(config, params.ID, inmem.New(), plugins.Info(info), plugins.InitBundles(loaded.Bundles), plugins.InitFiles(loaded.Files), plugins.MaxErrors(params.ErrorLimit))
 	if err != nil {
 		return nil, errors.Wrap(err, "config error")
 	}
 
-	info, err := runtime.Term(runtime.Params{Config: bs})
-	if err != nil {
-		return nil, err
-	}
-
-	manager, err := plugins.New(bs, params.ID, store, plugins.Info(info))
-	if err != nil {
-		return nil, errors.Wrap(err, "config error")
+	if err := manager.Init(ctx); err != nil {
+		return nil, errors.Wrap(err, "initialization error")
 	}
 
 	metrics := prometheus.New(metrics.New(), errorLogger)
@@ -248,7 +220,7 @@ func NewRuntime(ctx context.Context, params Params) (*Runtime, error) {
 	manager.Register("discovery", disco)
 
 	rt := &Runtime{
-		Store:   store,
+		Store:   manager.Store,
 		Params:  params,
 		Manager: manager,
 		info:    info,
@@ -432,7 +404,7 @@ func (rt *Runtime) readWatcher(ctx context.Context, watcher *fsnotify.Watcher, p
 
 func (rt *Runtime) processWatcherUpdate(ctx context.Context, paths []string, removed string) error {
 
-	loaded, err := loadPaths(paths, rt.Params.Filter, rt.Params.BundleMode)
+	loaded, err := initload.LoadPaths(paths, rt.Params.Filter, rt.Params.BundleMode)
 	if err != nil {
 		return err
 	}
@@ -440,11 +412,6 @@ func (rt *Runtime) processWatcherUpdate(ctx context.Context, paths []string, rem
 	removed = loader.CleanPath(removed)
 
 	return storage.Txn(ctx, rt.Store, storage.WriteParams, func(txn storage.Transaction) error {
-		if len(loaded.Documents) > 0 {
-			if err := rt.Store.Write(ctx, txn, storage.AddOp, storage.Path{}, loaded.Documents); err != nil {
-				return err
-			}
-		}
 
 		if !rt.Params.BundleMode {
 			ids, err := rt.Store.ListPolicies(ctx, txn)
@@ -456,7 +423,7 @@ func (rt *Runtime) processWatcherUpdate(ctx context.Context, paths []string, rem
 					if err := rt.Store.DeletePolicy(ctx, txn, id); err != nil {
 						return err
 					}
-				} else if _, exists := loaded.Modules[id]; !exists {
+				} else if _, exists := loaded.Files.Modules[id]; !exists {
 					// This branch get hit in two cases.
 					// 1. Another piece of code has access to the store and inserts
 					//    a policy out-of-band.
@@ -470,7 +437,7 @@ func (rt *Runtime) processWatcherUpdate(ctx context.Context, paths []string, rem
 					if err != nil {
 						return err
 					}
-					loaded.Modules[id] = &loader.RegoFile{
+					loaded.Files.Modules[id] = &loader.RegoFile{
 						Name:   id,
 						Raw:    bs,
 						Parsed: module,
@@ -478,12 +445,15 @@ func (rt *Runtime) processWatcherUpdate(ctx context.Context, paths []string, rem
 				}
 			}
 		}
-		if err := compileAndStoreInputs(ctx, rt.Store, txn, loaded, -1); err != nil {
-			return err
-		}
 
-		// re-add the version as it might have been overwritten from loading data files
-		if err := storedversion.Write(ctx, rt.Store, txn); err != nil {
+		_, err := initload.InsertAndCompile(ctx, initload.InsertAndCompileOptions{
+			Store:     rt.Store,
+			Txn:       txn,
+			Files:     loaded.Files,
+			Bundles:   loaded.Bundles,
+			MaxErrors: -1,
+		})
+		if err != nil {
 			return err
 		}
 
@@ -512,71 +482,6 @@ func (rt *Runtime) gracefulServerShutdown(s *server.Server) error {
 	return nil
 }
 
-type loadResult struct {
-	loader.Result
-	Bundles map[string]*bundle.Bundle
-}
-
-func loadPaths(paths []string, filter loader.Filter, asBundle bool) (*loadResult, error) {
-	result := &loadResult{}
-	var err error
-
-	if asBundle {
-		result.Bundles = make(map[string]*bundle.Bundle, len(paths))
-		for _, path := range paths {
-			result.Bundles[path], err = loader.NewFileLoader().AsBundle(path)
-			if err != nil {
-				return nil, err
-			}
-		}
-	} else {
-		loaded, err := loader.NewFileLoader().Filtered(paths, filter)
-		if err != nil {
-			return nil, err
-		}
-		result.Modules = loaded.Modules
-		result.Documents = loaded.Documents
-	}
-
-	return result, nil
-}
-
-func compileAndStoreInputs(ctx context.Context, store storage.Store, txn storage.Transaction, loaded *loadResult, errorLimit int) error {
-
-	policies := make(map[string]*ast.Module, len(loaded.Modules))
-
-	for id, parsed := range loaded.Modules {
-		policies[id] = parsed.Parsed
-	}
-
-	c := ast.NewCompiler().SetErrorLimit(errorLimit).WithPathConflictsCheck(storage.NonEmpty(ctx, store, txn))
-
-	opts := &bundle.ActivateOpts{
-		Ctx:          ctx,
-		Store:        store,
-		Txn:          txn,
-		Compiler:     c,
-		Metrics:      metrics.New(),
-		Bundles:      loaded.Bundles,
-		ExtraModules: policies,
-	}
-
-	err := bundle.Activate(opts)
-	if err != nil {
-		return err
-	}
-
-	// Policies in bundles will have already been added to the store, but
-	// modules loaded outside of bundles will need to be added manually.
-	for id, parsed := range loaded.Modules {
-		if err := store.UpsertPolicy(ctx, txn, id, parsed.Raw); err != nil {
-			return err
-		}
-	}
-
-	return nil
-}
-
 func getWatcher(rootPaths []string) (*fsnotify.Watcher, error) {
 
 	watchPaths, err := getWatchPaths(rootPaths)

runtime/runtime_test.go

@@ -25,143 +25,6 @@ import (
 	"github.com/open-policy-agent/opa/util/test"
 )
 
-func TestInit(t *testing.T) {
-
-	mod1 := `package a.b.c
-
-import data.a.foo
-
-p = true { foo = "bar" }
-p = true { 1 = 2 }`
-
-	mod2 := `package b.c.d
-
-import data.b.foo
-
-p = true { foo = "bar" }
-p = true { 1 = 2 }`
-
-	tests := []struct {
-		note         string
-		fs           map[string]string
-		loadParams   []string
-		expectedData map[string]string
-		expectedMods []string
-		asBundle     bool
-	}{
-		{
-			note: "load files",
-			fs: map[string]string{
-				"datafile":   `{"foo": "bar", "x": {"y": {"z": [1]}}}`,
-				"policyFile": mod1,
-			},
-			loadParams: []string{"datafile", "policyFile"},
-			expectedData: map[string]string{
-				"/foo": "bar",
-			},
-			expectedMods: []string{mod1},
-			asBundle:     false,
-		},
-		{
-			note: "load bundle",
-			fs: map[string]string{
-				"datafile":    `{"foo": "bar", "x": {"y": {"z": [1]}}}`, // Should be ignored
-				"data.json":   `{"foo": "not-bar"}`,
-				"policy.rego": mod1,
-			},
-			loadParams: []string{"/"},
-			expectedData: map[string]string{
-				"/foo": "not-bar",
-			},
-			expectedMods: []string{mod1},
-			asBundle:     true,
-		},
-		{
-			note: "load multiple bundles",
-			fs: map[string]string{
-				"/bundle1/a/data.json":   `{"foo": "bar1", "x": {"y": {"z": [1]}}}`, // Should be ignored
-				"/bundle1/a/policy.rego": mod1,
-				"/bundle1/a/.manifest":   `{"roots": ["a"]}`,
-				"/bundle2/b/data.json":   `{"foo": "bar2"}`,
-				"/bundle2/b/policy.rego": mod2,
-				"/bundle2/b/.manifest":   `{"roots": ["b"]}`,
-			},
-			loadParams: []string{"bundle1", "bundle2"},
-			expectedData: map[string]string{
-				"/a/foo": "bar1",
-				"/b/foo": "bar2",
-			},
-			expectedMods: []string{mod1, mod2},
-			asBundle:     true,
-		},
-	}
-
-	for _, tc := range tests {
-		t.Run(tc.note, func(t *testing.T) {
-			test.WithTempFS(tc.fs, func(rootDir string) {
-				params := NewParams()
-				for _, fileName := range tc.loadParams {
-					params.Paths = append(params.Paths, filepath.Join(rootDir, fileName))
-				}
-
-				params.BundleMode = tc.asBundle
-
-				testInitRuntime(t, params, tc.expectedData, tc.expectedMods)
-			})
-		})
-	}
-}
-
-func testInitRuntime(t *testing.T, params Params, expectedStoreData map[string]string, expectedMods []string) {
-	t.Helper()
-
-	ctx := context.Background()
-
-	rt, err := NewRuntime(ctx, params)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	txn := storage.NewTransactionOrDie(ctx, rt.Store)
-
-	for storePath, expected := range expectedStoreData {
-		node, err := rt.Store.Read(ctx, txn, storage.MustParsePath(storePath))
-		if util.Compare(node, expected) != 0 || err != nil {
-			t.Errorf("Expected %v but got %v (err: %v)", expected, node, err)
-			return
-		}
-	}
-
-	ids, err := rt.Store.ListPolicies(ctx, txn)
-	if err != nil {
-		t.Fatal(err)
-	}
-
-	if len(expectedMods) != len(ids) {
-		t.Fatalf("Expected %d modules, got %d", len(expectedMods), len(ids))
-	}
-
-	actualMods := map[string]struct{}{}
-	for _, id := range ids {
-		result, err := rt.Store.GetPolicy(ctx, txn, id)
-		if err != nil {
-			t.Fatalf("Unexpected error: %s", err)
-		}
-		actualMods[string(result)] = struct{}{}
-	}
-
-	for _, expectedMod := range expectedMods {
-		if _, found := actualMods[expectedMod]; !found {
-			t.Fatalf("Expected %v but got: %v", expectedMod, actualMods)
-		}
-	}
-
-	_, err = rt.Store.Read(ctx, txn, storage.MustParsePath("/system/version"))
-	if err != nil {
-		t.Fatal(err)
-	}
-}
-
 func TestWatchPaths(t *testing.T) {
 
 	fs := map[string]string{