Add ability to specify namespace for optimized files

Currently the namespace for partially evaluated files
in an optimized bundle cannot be modified. As a result
if more than one optimized bundle is loaded in OPA, a root
conflict error would occur as the optimized bundles have a
root called "partial" automatially added to their
manifest. This change allows the namespace to be configured via
the build command.

Fixes: #5933

Signed-off-by: Ashutosh Narkar <anarkar4387@gmail.com>

cmd/build.go

@@ -42,6 +42,7 @@ type buildParams struct {
 	claimsFile         string
 	excludeVerifyFiles []string
 	plugin             string
+	ns                 string
 }
 
 func newBuildParams() buildParams {
@@ -102,7 +103,8 @@ The -O flag controls the optimization level. By default, optimization is disable
 When optimization is enabled the 'build' command generates a bundle that is semantically
 equivalent to the input files however the structure of the files in the bundle may have
 been changed by rewriting, inlining, pruning, etc. Higher optimization levels may result
-in longer build times.
+in longer build times. The --partial-namespace flag can used in conjunction with the -O flag
+to specify the namespace for the partially evaluated files in the optimized bundle.
 
 The 'build' command supports targets (specified by -t):
 
@@ -233,6 +235,7 @@ against OPA v0.22.0:
 	buildCommand.Flags().VarP(&buildParams.entrypoints, "entrypoint", "e", "set slash separated entrypoint path")
 	buildCommand.Flags().VarP(&buildParams.revision, "revision", "r", "set output bundle revision")
 	buildCommand.Flags().StringVarP(&buildParams.outputFile, "output", "o", "bundle.tar.gz", "set the output filename")
+	buildCommand.Flags().StringVar(&buildParams.ns, "partial-namespace", "partial", "set the namespace to use for partially evaluated files in an optimized bundle")
 
 	addBundleModeFlag(buildCommand.Flags(), &buildParams.bundleMode, false)
 	addIgnoreFlag(buildCommand.Flags(), &buildParams.ignore)
@@ -294,7 +297,8 @@ func dobuild(params buildParams, args []string) error {
 		WithPaths(args...).
 		WithFilter(buildCommandLoaderFilter(params.bundleMode, params.ignore)).
 		WithBundleVerificationConfig(bvc).
-		WithBundleSigningConfig(bsc)
+		WithBundleSigningConfig(bsc).
+		WithPartialNamespace(params.ns)
 
 	if params.revision.isSet {
 		compiler = compiler.WithRevision(*params.revision.v)

compile/compile.go

@@ -81,6 +81,7 @@ type Compiler struct {
 	keyID                        string                     // represents the name of the default key used to verify a signed bundle
 	metadata                     *map[string]interface{}    // represents additional data included in .manifest file
 	fsys                         fs.FS                      // file system to use when loading paths
+	ns                           string
 }
 
 // New returns a new compiler instance that can be invoked.
@@ -228,6 +229,12 @@ func (c *Compiler) WithFS(fsys fs.FS) *Compiler {
 	return c
 }
 
+// WithPartialNamespace sets the namespace to use for partial evaluation results
+func (c *Compiler) WithPartialNamespace(ns string) *Compiler {
+	c.ns = ns
+	return c
+}
+
 func addEntrypointsFromAnnotations(c *Compiler, ar []*ast.AnnotationsRef) error {
 	for _, ref := range ar {
 		var entrypoint ast.Ref
@@ -490,6 +497,10 @@ func (c *Compiler) optimize(ctx context.Context) error {
 		WithShallowInlining(c.optimizationLevel <= 1).
 		WithEnablePrintStatements(c.enablePrintStatements)
 
+	if c.ns != "" {
+		o = o.WithPartialNamespace(c.ns)
+	}
+
 	err := o.Do(ctx)
 	if err != nil {
 		return err
@@ -844,6 +855,11 @@ func (o *optimizer) WithShallowInlining(yes bool) *optimizer {
 	return o
 }
 
+func (o *optimizer) WithPartialNamespace(ns string) *optimizer {
+	o.nsprefix = ns
+	return o
+}
+
 func (o *optimizer) Do(ctx context.Context) error {
 
 	// NOTE(tsandall): if there are multiple entrypoints, copy the bundle because

compile/compile_test.go

@@ -538,6 +538,57 @@ func TestCompilerOptimizationL2(t *testing.T) {
 	}
 }
 
+func TestCompilerOptimizationWithConfiguredNamespace(t *testing.T) {
+
+	files := map[string]string{
+		"test.rego": `
+			package test
+
+			p { not q }
+			q { k[input.a]; k[input.b] }  # generate a product that is not inlined
+			k = {1,2,3}
+		`,
+	}
+
+	for _, useMemoryFS := range []bool{false, true} {
+		test.WithTestFS(files, useMemoryFS, func(root string, fsys fs.FS) {
+
+			compiler := New().
+				WithFS(fsys).
+				WithPaths(root).
+				WithOptimizationLevel(1).
+				WithEntrypoints("test/p").
+				WithPartialNamespace("custom")
+
+			err := compiler.Build(context.Background())
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			if len(compiler.bundle.Modules) != 2 {
+				t.Fatalf("expected two modules but got: %v", len(compiler.bundle.Modules))
+			}
+
+			optimizedExp := ast.MustParseModule(`package custom
+
+			__not1_0_2__ = true { data.test.q = _; _ }`)
+
+			if optimizedExp.String() != compiler.bundle.Modules[0].Parsed.String() {
+				t.Fatalf("expected optimized module to be:\n\n%v\n\ngot:\n\n%v", optimizedExp, compiler.bundle.Modules[0])
+			}
+
+			expected := ast.MustParseModule(`package test
+k = {1, 2, 3} { true }
+p = true { not data.custom.__not1_0_2__ }
+q = true { __local0__3 = input.a; data.test.k[__local0__3] = _; _; __local1__3 = input.b; data.test.k[__local1__3] = _; _ }`)
+
+			if expected.String() != compiler.bundle.Modules[1].Parsed.String() {
+				t.Fatalf("expected module to be:\n\n%v\n\ngot:\n\n%v", expected, compiler.bundle.Modules[1])
+			}
+		})
+	}
+}
+
 // NOTE(sr): we override this to not depend on build tags in tests
 func wasmABIVersions(vs ...ast.WasmABIVersion) *ast.Capabilities {
 	caps := ast.CapabilitiesForThisVersion()
@@ -1428,7 +1479,7 @@ func TestOptimizerNoops(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
-			o := getOptimizer(tc.modules, "", tc.entrypoints, nil)
+			o := getOptimizer(tc.modules, "", tc.entrypoints, nil, "")
 			cpy := o.bundle.Copy()
 			err := o.Do(context.Background())
 			if err != nil {
@@ -1479,7 +1530,7 @@ func TestOptimizerErrors(t *testing.T) {
 
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
-			o := getOptimizer(tc.modules, "", tc.entrypoints, nil)
+			o := getOptimizer(tc.modules, "", tc.entrypoints, nil, "")
 			cpy := o.bundle.Copy()
 			got := o.Do(context.Background())
 			if got == nil || got.Error() != tc.wantErr.Error() {
@@ -1499,6 +1550,7 @@ func TestOptimizerOutput(t *testing.T) {
 		modules     map[string]string
 		data        string
 		roots       []string
+		namespace   string
 		wantModules map[string]string
 	}{
 		{
@@ -1767,6 +1819,46 @@ func TestOptimizerOutput(t *testing.T) {
 				`,
 			},
 		},
+		{
+			note:        "configured package namespace",
+			entrypoints: []string{"data.test.p"},
+			namespace:   "custom",
+			modules: map[string]string{
+				"test.rego": `
+					package test
+
+					p { not q }
+					q { k[input.a]; k[input.b] }  # generate a product that is not inlined
+					k = {1,2,3}
+				`,
+			},
+			wantModules: map[string]string{
+				"optimized/custom.rego": `
+					package custom
+
+					__not1_0_2__ = true { 1 = input.a; 1 = input.b }
+					__not1_0_2__ = true { 1 = input.a; 2 = input.b }
+					__not1_0_2__ = true { 1 = input.a; 3 = input.b }
+					__not1_0_2__ = true { 2 = input.a; 1 = input.b }
+					__not1_0_2__ = true { 2 = input.a; 2 = input.b }
+					__not1_0_2__ = true { 2 = input.a; 3 = input.b }
+					__not1_0_2__ = true { 3 = input.a; 1 = input.b }
+					__not1_0_2__ = true { 3 = input.a; 2 = input.b }
+					__not1_0_2__ = true { 3 = input.a; 3 = input.b }
+				`,
+				"optimized/test.rego": `
+					package test
+
+					p = __result__ { not data.custom.__not1_0_2__; __result__ = true }
+				`,
+				"test.rego": `
+					package test
+
+					q = true { k[input.a]; k[input.b] }
+					k = {1, 2, 3} { true }
+				`,
+			},
+		},
 		{
 			note:        "infer unknowns from roots",
 			entrypoints: []string{"data.test.p"},
@@ -1935,7 +2027,7 @@ func TestOptimizerOutput(t *testing.T) {
 	for _, tc := range tests {
 		t.Run(tc.note, func(t *testing.T) {
 
-			o := getOptimizer(tc.modules, tc.data, tc.entrypoints, tc.roots)
+			o := getOptimizer(tc.modules, tc.data, tc.entrypoints, tc.roots, tc.namespace)
 			original := o.bundle.Copy()
 			err := o.Do(context.Background())
 			if err != nil {
@@ -1949,7 +2041,13 @@ func TestOptimizerOutput(t *testing.T) {
 
 			if len(tc.roots) > 0 {
 				exp.Manifest.Roots = &tc.roots
-				exp.Manifest.AddRoot("partial") // optimizer will add this automatically
+
+				// optimizer will add the manifest root in the optimized bundle automatically
+				if tc.namespace != "" {
+					exp.Manifest.AddRoot(tc.namespace)
+				} else {
+					exp.Manifest.AddRoot("partial")
+				}
 			}
 
 			exp.Manifest.Revision = "" // optimizations must reset the revision.
@@ -2017,7 +2115,7 @@ func TestRefSet(t *testing.T) {
 
 }
 
-func getOptimizer(modules map[string]string, data string, entries []string, roots []string) *optimizer {
+func getOptimizer(modules map[string]string, data string, entries []string, roots []string, ns string) *optimizer {
 
 	b := &bundle.Bundle{
 		Modules: getModuleFiles(modules, true),
@@ -2042,6 +2140,10 @@ func getOptimizer(modules map[string]string, data string, entries []string, root
 	o := newOptimizer(ast.CapabilitiesForThisVersion(), b).
 		WithEntrypoints(entrypoints)
 
+	if ns != "" {
+		o = o.WithPartialNamespace(ns)
+	}
+
 	o.resultsymprefix = ""
 
 	return o