diff --git a/v1/ast/compile.go b/v1/ast/compile.go index f3ca101735..2ffc0f59a2 100644 --- a/v1/ast/compile.go +++ b/v1/ast/compile.go @@ -1681,6 +1681,31 @@ func (c *Compiler) init() { return } + if defaultModuleLoader != nil { + if c.moduleLoader == nil { + c.moduleLoader = defaultModuleLoader + } else { + first := c.moduleLoader + c.moduleLoader = func(res map[string]*Module) (map[string]*Module, error) { + res0, err := first(res) + if err != nil { + return nil, err + } + res1, err := defaultModuleLoader(res) + if err != nil { + return nil, err + } + // merge res1 into res0, based on module "file" names, to avoid clashes + for k, v := range res1 { + if _, ok := res0[k]; !ok { + res0[k] = v + } + } + return res0, nil + } + } + } + if c.capabilities == nil { c.capabilities = CapabilitiesForThisVersion() } diff --git a/v1/ast/compile_test.go b/v1/ast/compile_test.go index 5c7e01b348..04514957eb 100644 --- a/v1/ast/compile_test.go +++ b/v1/ast/compile_test.go @@ -11500,3 +11500,63 @@ func TestCompile_DefaultRegoVersion(t *testing.T) { }) } } + +func TestCompilerInitWithDefaultModuleLoader(t *testing.T) { + // Reset the global variable after the test + defer func() { defaultModuleLoader = nil }() + + // a dummy loader that adds "foo" + loader1 := func(res map[string]*Module) (map[string]*Module, error) { + mod := MustParseModule(`package foo`) + resCopy := map[string]*Module{} + maps.Copy(resCopy, res) + resCopy["foo.rego"] = mod + return resCopy, nil + } + + // a dummy loader that adds "bar" + loader2 := func(res map[string]*Module) (map[string]*Module, error) { + mod := MustParseModule(`package bar`) + resCopy := map[string]*Module{} + maps.Copy(resCopy, res) + resCopy["bar.rego"] = mod + return resCopy, nil + } + + DefaultModuleLoader(loader2) + + c := NewCompiler().WithModuleLoader(loader1) + c.init() + + got, err := c.moduleLoader(make(map[string]*Module)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + + expected := map[string]*Module{ + "foo.rego": MustParseModule(`package foo`), + "bar.rego": MustParseModule(`package bar`), + } + // check both modules are present + for k, v := range expected { + gotMod, ok := got[k] + if !ok { + t.Errorf("expected key %q in result", k) + continue + } + if !reflect.DeepEqual(gotMod, v) { + t.Errorf("unexpected module for %q: got %v want %v", k, gotMod, v) + } + } + + // Now, test defaultModuleLoader only + c2 := NewCompiler() + c2.init() + got2, err := c2.moduleLoader(make(map[string]*Module)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if _, ok := got2["bar.rego"]; !ok { + t.Error("expected bar.rego from defaultModuleLoader in result") + } +} diff --git a/v1/ast/default_module_loader.go b/v1/ast/default_module_loader.go new file mode 100644 index 0000000000..528c253e16 --- /dev/null +++ b/v1/ast/default_module_loader.go @@ -0,0 +1,14 @@ +// Copyright 2025 The OPA Authors. All rights reserved. +// Use of this source code is governed by an Apache2 +// license that can be found in the LICENSE file. + +package ast + +var defaultModuleLoader ModuleLoader + +// DefaultModuleLoader lets you inject an `ast.ModuleLoader` that will +// always be used. If another one is provided with the ast package, +// they will both be consulted to enrich the set of modules dynamically. +func DefaultModuleLoader(ml ModuleLoader) { + defaultModuleLoader = ml +}