Add option for specifying decision ID to SDK

This allows clients to specify the decision ID for the decisions returned by the `Decision` and `Partial` functions provided by the SDK package. If not provided, a uniquely generated identifier will be generated for the decision, just as before. This option can be useful for clients that have already generated a decision ID prior to calling OPA. It would then be convenient to pass that decision ID through to OPA so decision logs use that same ID. Signed-off-by: Brian Chhun <brian.chhun@chime.com>
open-policy-agent · Jul 17, 2023 · 5cd774e · 5cd774e
1 parent 4b62709
commit 5cd774e
Show file tree

Hide file tree

Showing 2 changed files with 172 additions and 19 deletions.
diff --git a/sdk/opa.go b/sdk/opa.go
@@ -239,6 +239,7 @@ func (opa *OPA) Decision(ctx context.Context, options DecisionOptions) (*Decisio
 		Input:          &options.Input,
 		NDBuiltinCache: &options.NDBCache,
 		Metrics:        options.Metrics,
+		DecisionID:     options.DecisionID,
 	}
 
 	// Only use non-deterministic builtins cache if it's available.
@@ -294,41 +295,36 @@ type DecisionOptions struct {
 	Metrics             metrics.Metrics     // specifies the metrics to use for preparing and evaluation, optional
 	Profiler            topdown.QueryTracer // specifies the profiler to use, optional
 	Instrument          bool                // if true, instrumentation will be enabled
+	DecisionID          string              // the identifier for this decision; if not set, a globally unique identifier will be generated
 }
 
 // DecisionResult contains the output of query evaluation.
 type DecisionResult struct {
-	ID         string             // provides a globally unique identifier for this decision (which is included in the decision log.)
+	ID         string             // provides the identifier for this decision (which is included in the decision log.)
 	Result     interface{}        // provides the output of query evaluation.
 	Provenance types.ProvenanceV1 // wraps the bundle build/version information
 }
 
-func newDecisionResult() (*DecisionResult, error) {
-	id, err := uuid.New(rand.Reader)
-	if err != nil {
-		return nil, err
-	}
-	result := &DecisionResult{ID: id}
-	return result, nil
-}
-
 func (opa *OPA) executeTransaction(ctx context.Context, record *server.Info, work func(state, *DecisionResult)) (*DecisionResult, error) {
 	if record.Metrics == nil {
 		record.Metrics = metrics.New()
 	}
 	record.Metrics.Timer(metrics.SDKDecisionEval).Start()
 
-	result, err := newDecisionResult()
-	if err != nil {
-		return nil, err
+	if record.DecisionID == "" {
+		id, err := uuid.New(rand.Reader)
+		if err != nil {
+			return nil, err
+		}
+		record.DecisionID = id
 	}
 
+	result := &DecisionResult{ID: record.DecisionID}
+
 	opa.mtx.Lock()
 	s := *opa.state
 	opa.mtx.Unlock()
 
-	record.DecisionID = result.ID
-
 	if record.Timestamp.IsZero() {
 		record.Timestamp = time.Now().UTC()
 	}
@@ -375,10 +371,11 @@ func (opa *OPA) Partial(ctx context.Context, options PartialOptions) (*PartialRe
 	}
 
 	record := server.Info{
-		Timestamp: options.Now,
-		Input:     &options.Input,
-		Query:     options.Query,
-		Metrics:   options.Metrics,
+		Timestamp:  options.Now,
+		Input:      &options.Input,
+		Query:      options.Query,
+		Metrics:    options.Metrics,
+		DecisionID: options.DecisionID,
 	}
 
 	var provenance types.ProvenanceV1
@@ -448,6 +445,7 @@ type PartialOptions struct {
 	Metrics             metrics.Metrics     // specifies the metrics to use for preparing and evaluation, optional
 	Profiler            topdown.QueryTracer // specifies the profiler to use, optional
 	Instrument          bool                // if true, instrumentation will be enabled
+	DecisionID          string              // the identifier for this decision; if not set, a globally unique identifier will be generated
 }
 
 type PartialResult struct {

diff --git a/sdk/opa_test.go b/sdk/opa_test.go
@@ -735,6 +735,76 @@ main = data.foo
 
 }
 
+func TestDecisionWithConfigurableID(t *testing.T) {
+	ctx := context.Background()
+
+	server := sdktest.MustNewServer(
+		sdktest.MockBundle("/bundles/bundle.tar.gz", map[string]string{
+			"main.rego": `
+package system
+
+main = time.now_ns()
+`,
+		}),
+	)
+
+	defer server.Stop()
+
+	config := fmt.Sprintf(`{
+		"services": {
+			"test": {
+				"url": %q
+			}
+		},
+		"bundles": {
+			"test": {
+				"resource": "/bundles/bundle.tar.gz"
+			}
+		},
+		"decision_logs": {
+			"console": true
+		}
+	}`, server.URL())
+
+	testLogger := loggingtest.New()
+	opa, err := sdk.New(ctx, sdk.Options{
+		Config:        strings.NewReader(config),
+		ConsoleLogger: testLogger})
+
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer opa.Stop(ctx)
+
+	if _, err := opa.Decision(ctx, sdk.DecisionOptions{
+		Now: time.Unix(0, 1619868194450288000).UTC(),
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	if _, err := opa.Decision(ctx, sdk.DecisionOptions{
+		Now:        time.Unix(0, 1619868194450288000).UTC(),
+		DecisionID: "164031de-e511-11ec-8fea-0242ac120002",
+	}); err != nil {
+		t.Fatal(err)
+	}
+
+	entries := testLogger.Entries()
+
+	if exp, act := 2, len(entries); exp != act {
+		t.Fatalf("expected %d entries, got %d", exp, act)
+	}
+
+	if entries[0].Fields["decision_id"] == "" {
+		t.Fatalf("expected not empty decision_id")
+	}
+
+	if entries[1].Fields["decision_id"] != "164031de-e511-11ec-8fea-0242ac120002" {
+		t.Fatalf("expected %v but got %v", "164031de-e511-11ec-8fea-0242ac120002", entries[1].Fields["decision_id"])
+	}
+}
+
 func TestPartial(t *testing.T) {
 
 	ctx := context.Background()
@@ -1230,6 +1300,91 @@ allow {
 
 }
 
+func TestPartialWithConfigurableID(t *testing.T) {
+
+	ctx := context.Background()
+
+	server := sdktest.MustNewServer(
+		sdktest.MockBundle("/bundles/bundle.tar.gz", map[string]string{
+			"main.rego": `
+package test
+
+allow {
+	data.junk.x = input.y
+}
+`,
+		}),
+	)
+
+	defer server.Stop()
+
+	config := fmt.Sprintf(`{
+		"services": {
+			"test": {
+				"url": %q
+			}
+		},
+		"bundles": {
+			"test": {
+				"resource": "/bundles/bundle.tar.gz"
+			}
+		},
+		"decision_logs": {
+			"console": true
+		}
+	}`, server.URL())
+
+	testLogger := loggingtest.New()
+	opa, err := sdk.New(ctx, sdk.Options{
+		Config:        strings.NewReader(config),
+		ConsoleLogger: testLogger,
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	defer opa.Stop(ctx)
+
+	if result, err := opa.Partial(ctx, sdk.PartialOptions{
+		Input:    map[string]int{"y": 2},
+		Query:    "data.test.allow = true",
+		Unknowns: []string{"data.junk.x"},
+		Mapper:   &sdk.RawMapper{},
+		Now:      time.Unix(0, 1619868194450288000).UTC(),
+	}); err != nil {
+		t.Fatal(err)
+	} else if decision, ok := result.Result.(*rego.PartialQueries); !ok || decision.Queries[0].String() != "2 = data.junk.x" {
+		t.Fatal("expected &{[2 = data.junk.x] []} true but got:", decision, ok)
+	}
+
+	if result, err := opa.Partial(ctx, sdk.PartialOptions{
+		Input:      map[string]int{"y": 2},
+		Query:      "data.test.allow = true",
+		Unknowns:   []string{"data.junk.x"},
+		Mapper:     &sdk.RawMapper{},
+		Now:        time.Unix(0, 1619868194450288000).UTC(),
+		DecisionID: "164031de-e511-11ec-8fea-0242ac120002",
+	}); err != nil {
+		t.Fatal(err)
+	} else if decision, ok := result.Result.(*rego.PartialQueries); !ok || decision.Queries[0].String() != "2 = data.junk.x" {
+		t.Fatal("expected &{[2 = data.junk.x] []} true but got:", decision, ok)
+	}
+
+	entries := testLogger.Entries()
+
+	if exp, act := 2, len(entries); exp != act {
+		t.Fatalf("expected %d entries, got %d", exp, act)
+	}
+
+	if entries[0].Fields["decision_id"] == "" {
+		t.Fatalf("expected not empty decision_id")
+	}
+
+	if entries[1].Fields["decision_id"] != "164031de-e511-11ec-8fea-0242ac120002" {
+		t.Fatalf("expected %v but got %v", "164031de-e511-11ec-8fea-0242ac120002", entries[1].Fields["decision_id"])
+	}
+}
+
 func TestUndefinedError(t *testing.T) {
 
 	ctx := context.Background()