From 5ef98c74931fbdd10aa8f00f6c2e24905132139f Mon Sep 17 00:00:00 2001 From: Stephan Renatus Date: Thu, 17 Jul 2025 19:04:45 +0200 Subject: [PATCH] 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 * runtime: allow passing in custom store builder Signed-off-by: Stephan Renatus --------- Signed-off-by: Stephan Renatus Signed-off-by: Stephan Renatus --- storage/interface.go | 3 ++ v1/runtime/runtime.go | 14 ++++++-- v1/runtime/runtime_test.go | 73 ++++++++++++++++++++++++++++++++++++++ v1/storage/interface.go | 5 +++ v1/storage/storage.go | 3 ++ v1/storage/storage_test.go | 33 ++++++++++++++++- 6 files changed, 128 insertions(+), 3 deletions(-) diff --git a/storage/interface.go b/storage/interface.go index 0192c459c8..a21b5575e9 100644 --- a/storage/interface.go +++ b/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 diff --git a/v1/runtime/runtime.go b/v1/runtime/runtime.go index 206340cae1..9563768aee 100644 --- a/v1/runtime/runtime.go +++ b/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)) } diff --git a/v1/runtime/runtime_test.go b/v1/runtime/runtime_test.go index 26837ba1d4..c10a0b3af6 100644 --- a/v1/runtime/runtime_test.go +++ b/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) + } +} diff --git a/v1/storage/interface.go b/v1/storage/interface.go index 1d03567066..a783caae09 100644 --- a/v1/storage/interface.go +++ b/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 { diff --git a/v1/storage/storage.go b/v1/storage/storage.go index ecc3829940..38d51be405 100644 --- a/v1/storage/storage.go +++ b/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 diff --git a/v1/storage/storage_test.go b/v1/storage/storage_test.go index cb59613e81..597f1589be 100644 --- a/v1/storage/storage_test.go +++ b/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) + } + } +}