store+runtime: extension points for custom stores (#7779)

* storage: allow overriding NonEmpty

Custom store implementations can now bring their own NonEmpty() methods,
which may be more efficient than what the generic method does.

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>

* runtime: allow passing in custom store builder


Signed-off-by: Stephan Renatus <stephan@styra.com>

---------

Signed-off-by: Stephan Renatus <stephan.renatus@gmail.com>
Signed-off-by: Stephan Renatus <stephan@styra.com>

storage/interface.go

@@ -19,6 +19,9 @@ type Store = v1.Store
 // generic MakeDir functionality in storage.MakeDir
 type MakeDirer = v1.MakeDirer
 
+// NonEmptyer allows a store implemention to override NonEmpty())
+type NonEmptyer = v1.NonEmptyer
+
 // TransactionParams describes a new transaction.
 type TransactionParams = v1.TransactionParams
 

v1/runtime/runtime.go

@@ -24,6 +24,7 @@ import (
 	"time"
 
 	"github.com/fsnotify/fsnotify"
+	prometheus_sdk "github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
 	"go.opentelemetry.io/otel/propagation"
@@ -227,6 +228,9 @@ type Params struct {
 	// It can also be enabled via config, and this runtime field takes precedence.
 	DiskStorage *disk.Options
 
+	// StoreBuilder allows passing a storage backend builder
+	StoreBuilder func(_ context.Context, _ logging.Logger, _ prometheus_sdk.Registerer, config []byte, id string) (storage.Store, error)
+
 	DistributedTracingOpts tracing.Options
 
 	// Check if default Addr is set or the user has changed it.
@@ -433,12 +437,18 @@ func NewRuntime(ctx context.Context, params Params) (*Runtime, error) {
 		}
 	}
 
-	if params.DiskStorage != nil {
+	switch {
+	case params.DiskStorage != nil:
 		store, err = disk.New(ctx, logger, metrics, *params.DiskStorage)
 		if err != nil {
 			return nil, fmt.Errorf("initialize disk store: %w", err)
 		}
-	} else {
+	case params.StoreBuilder != nil:
+		store, err = params.StoreBuilder(ctx, logger, metrics, config, params.ID)
+		if err != nil {
+			return nil, fmt.Errorf("initialize store: %w", err)
+		}
+	default:
 		store = inmem.NewWithOpts(inmem.OptRoundTripOnWrite(false),
 			inmem.OptReturnASTValuesOnRead(params.ReadAstValuesFromStore))
 	}

v1/runtime/runtime_test.go

@@ -19,10 +19,12 @@ import (
 	"path/filepath"
 	"reflect"
 	"runtime"
+	"slices"
 	"strings"
 	"testing"
 	"time"
 
+	prometheus_sdk "github.com/prometheus/client_golang/prometheus"
 	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
 	"go.opentelemetry.io/otel/sdk/trace"
 	"go.opentelemetry.io/otel/sdk/trace/tracetest"
@@ -32,6 +34,7 @@ import (
 	"github.com/open-policy-agent/opa/v1/loader"
 	"github.com/open-policy-agent/opa/v1/plugins"
 	"github.com/open-policy-agent/opa/v1/plugins/discovery"
+	"github.com/open-policy-agent/opa/v1/storage/inmem"
 	"github.com/open-policy-agent/opa/v1/tracing"
 
 	"github.com/open-policy-agent/opa/internal/report"
@@ -1927,3 +1930,73 @@ func TestCacheHooksOnServer(t *testing.T) {
 		t.Log(e.Message)
 	}
 }
+
+type fakeStore struct {
+	storage.Store
+}
+
+func (f *fakeStore) Read(ctx context.Context, txn storage.Transaction, p storage.Path) (any, error) {
+	if slices.Contains(p, "foo") {
+		return map[string]any{"fake": p}, nil
+	}
+	return f.Store.Read(ctx, txn, p)
+}
+
+func TestCustomStoreBuilder(t *testing.T) {
+	ctx := context.Background()
+	testLogger := testLog.New()
+	params := NewParams()
+	params.Logger = testLogger
+	params.Addrs = &[]string{"localhost:0"}
+	params.StoreBuilder = func(_ context.Context, logger logging.Logger, registerer prometheus_sdk.Registerer, config []byte, id string) (storage.Store, error) {
+		switch {
+		case logger == nil:
+			t.Fatal("logger empty")
+		case registerer == nil:
+			t.Fatal("registerer empty")
+		case config == nil:
+			t.Fatal("config empty")
+		case id == "":
+			t.Fatal("id empty")
+		}
+		return &fakeStore{inmem.New()}, nil
+	}
+
+	rt, err := NewRuntime(ctx, params)
+	if err != nil {
+		t.Fatalf("Unexpected error %v", err)
+	}
+	go rt.StartServer(ctx)
+	if !test.Eventually(t, 5*time.Second, func() bool {
+		found := false
+		for _, e := range testLogger.Entries() {
+			found = strings.Contains(e.Message, "Server initialized.") || found
+		}
+		return found
+	}) {
+		t.Fatal("Timed out waiting for server to start")
+	}
+	host := rt.Addrs()[0]
+	r, err := http.NewRequest(http.MethodGet, "http://"+host+"/v1/data/foo/bar", nil)
+	if err != nil {
+		t.Fatalf("Unexpected error: %s", err)
+	}
+	resp, err := http.DefaultClient.Do(r)
+	if err != nil {
+		t.Fatal("expected no error, got", err)
+	}
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("status %d (want 200)", resp.StatusCode)
+	}
+
+	defer resp.Body.Close()
+	var payload struct {
+		Result map[string]any
+	}
+	if err := json.NewDecoder(resp.Body).Decode(&payload); err != nil {
+		t.Fatalf("decode response: %v", err)
+	}
+	if !reflect.DeepEqual(map[string]any{"fake": []any{"foo", "bar"}}, payload.Result) {
+		t.Errorf("unexpected result: %v", payload.Result)
+	}
+}

v1/storage/interface.go

@@ -49,6 +49,11 @@ type MakeDirer interface {
 	MakeDir(context.Context, Transaction, Path) error
 }
 
+// NonEmptyer allows a store implemention to override NonEmpty())
+type NonEmptyer interface {
+	NonEmpty(context.Context, Transaction) func([]string) (bool, error)
+}
+
 // TransactionParams describes a new transaction.
 type TransactionParams struct {
 

v1/storage/storage.go

@@ -111,6 +111,9 @@ func Txn(ctx context.Context, store Store, params TransactionParams, f func(Tran
 // path is non-empty if a Read on the path returns a value or a Read
 // on any of the path prefixes returns a non-object value.
 func NonEmpty(ctx context.Context, store Store, txn Transaction) func([]string) (bool, error) {
+	if md, ok := store.(NonEmptyer); ok {
+		return md.NonEmpty(ctx, txn)
+	}
 	return func(path []string) (bool, error) {
 		if _, err := store.Read(ctx, txn, Path(path)); err == nil {
 			return true, nil

v1/storage/storage_test.go

@@ -60,7 +60,7 @@ func TestNonEmpty(t *testing.T) {
 					t.Fatal(err)
 				}
 				if nonEmpty != tc.exp {
-					t.Errorf("Expected %v for %v on %v but got", tc.exp, tc.path, tc.content)
+					t.Errorf("Expected %v for %v on %v but got %v", tc.exp, tc.path, tc.content, nonEmpty)
 				}
 				return nil
 			})
@@ -71,3 +71,34 @@ func TestNonEmpty(t *testing.T) {
 	}
 
 }
+
+type nonEmpty struct {
+	storage.Store
+}
+
+func (*nonEmpty) NonEmpty(context.Context, storage.Transaction) func([]string) (bool, error) {
+	return func([]string) (bool, error) {
+		return true, nil
+	}
+}
+
+func TestNonEmptyer(t *testing.T) {
+	ctx := context.Background()
+	ne := &nonEmpty{inmem.New()}
+
+	for _, path := range []string{"a", "a/b/c"} {
+		err := storage.Txn(ctx, ne, storage.TransactionParams{}, func(txn storage.Transaction) error {
+			nonEmpty, err := storage.NonEmpty(ctx, ne, txn)(strings.Split(path, "/"))
+			if err != nil {
+				t.Fatal(err)
+			}
+			if nonEmpty != true {
+				t.Errorf("Expected true for %v but got false", path)
+			}
+			return nil
+		})
+		if err != nil {
+			t.Error(err)
+		}
+	}
+}