From a03c8a98ce2b5743fd1ef0934f68827ca3ee22ff Mon Sep 17 00:00:00 2001
From: Torin Sandall <torinsandall@gmail.com>
Date: Mon, 30 Aug 2021 15:57:03 -0700
Subject: [PATCH] rego: Add persistent storage example

Signed-off-by: Torin Sandall <torinsandall@gmail.com>
---
 rego/example_test.go | 91 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 91 insertions(+)

diff --git a/rego/example_test.go b/rego/example_test.go
index 1014ec8281..4c4c73cc68 100644
--- a/rego/example_test.go
+++ b/rego/example_test.go
@@ -10,9 +10,12 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"io/ioutil"
 	"os"
 	"strings"
 
+	"github.com/open-policy-agent/opa/storage/disk"
+
 	"github.com/open-policy-agent/opa/ast"
 	"github.com/open-policy-agent/opa/rego"
 	"github.com/open-policy-agent/opa/storage"
@@ -313,6 +316,94 @@ func ExampleRego_Eval_storage() {
 	// value: [dogs clouds]
 }
 
+func ExampleRego_Eval_persistent_storage() {
+
+	ctx := context.Background()
+
+	data := `{
+        "example": {
+            "users": {
+				"alice": {
+					"likes": ["dogs", "clouds"]
+				},
+				"bob": {
+					"likes": ["pizza", "cats"]
+				}
+			}
+        }
+    }`
+
+	var json map[string]interface{}
+
+	err := util.UnmarshalJSON([]byte(data), &json)
+	if err != nil {
+		// Handle error.
+	}
+
+	// Manually create a persistent storage-layer in a temporary directory.
+	rootDir, err := ioutil.TempDir("", "rego_example")
+	if err != nil {
+		panic(err)
+	}
+
+	defer os.RemoveAll(rootDir)
+
+	// Configure the store to partition data at `/example/users` so that each
+	// user's data is stored on a different row. Assuming the policy only reads
+	// data for a single user to process the policy query, OPA can avoid loading
+	// _all_ user data into memory this way.
+	store, err := disk.New(ctx, disk.Options{
+		Dir:        rootDir,
+		Partitions: []storage.Path{{"example", "user"}},
+	})
+	if err != nil {
+		// Handle error.
+	}
+
+	err = storage.WriteOne(ctx, store, storage.AddOp, storage.Path{}, json)
+	if err != nil {
+		// Handle error
+	}
+
+	// Run a query that returns the value
+	rs, err := rego.New(
+		rego.Query(`data.example.users["alice"].likes`),
+		rego.Store(store)).Eval(ctx)
+	if err != nil {
+		// Handle error.
+	}
+
+	// Inspect the result.
+	fmt.Println("value:", rs[0].Expressions[0].Value)
+
+	// Re-open the store in the same directory.
+	store.Close(ctx)
+
+	store2, err := disk.New(ctx, disk.Options{
+		Dir:        rootDir,
+		Partitions: []storage.Path{{"example", "user"}},
+	})
+	if err != nil {
+		// Handle error.
+	}
+
+	// Run the same query with a new store.
+	rs, err = rego.New(
+		rego.Query(`data.example.users["alice"].likes`),
+		rego.Store(store2)).Eval(ctx)
+	if err != nil {
+		// Handle error.
+	}
+
+	// Inspect the result and observe the same result.
+	fmt.Println("value:", rs[0].Expressions[0].Value)
+
+	// Output:
+	//
+	// value: [dogs clouds]
+	// value: [dogs clouds]
+}
+
 func ExampleRego_Eval_transactions() {
 
 	ctx := context.Background()