ensure query can be evaluated during preparation

Continuing my trend of shoving more validation into
`OPA.Engine.prepareForEvaluation`, this method now checks that the
provided query can be evaluated. This should help catch
misconfigurations earlier in application bootstrapping, without
necessarily having to evaluate something first.

I think it's a bit of a deviation from how query preparation works with
the Go implementation. And I could see this check becoming much more
difficult or impossible to implement as query capabilities expand, so I
won't be upset if this is rejected.

There are already some tests checking for this behavior, but they were
expecting it to happen on `evaluate`. Now the error pops up during
`prepareForEvaluation`, but the tests pass either way.

Signed-off-by: Chris Rice <[email protected]>

Author: dev-rice

Sources/Rego/Engine.swift

@@ -104,20 +104,21 @@ extension OPA.Engine {
             try await store.write(to: StoreKeyPath(["data"]), value: bundle.data)
         }
 
+        let evaluator: Evaluator
         if self.policies.count > 0 {
             guard loadedBundles.isEmpty else {
                 throw RegoError.init(code: .invalidArgumentError, message: "Cannot mix direct IR policies with bundles")
             }
-            return PreparedQuery(
-                query: query,
-                evaluator: IREvaluator(policies: self.policies),
-                store: self.store
-            )
+            evaluator = IREvaluator(policies: self.policies)
+        } else {
+            evaluator = try IREvaluator(bundles: loadedBundles)
         }
 
+        try evaluator.ensureQueryIsSupported(query)
+
         return PreparedQuery(
             query: query,
-            evaluator: try IREvaluator(bundles: loadedBundles),
+            evaluator: evaluator,
             store: self.store
         )
     }

Sources/Rego/Evaluator.swift

@@ -3,6 +3,7 @@ import AST
 import IR
 
 protocol Evaluator {
+    func ensureQueryIsSupported(_ query: String) throws
     func evaluate(withContext ctx: EvaluationContext) async throws -> ResultSet
 }
 

Sources/Rego/IREvaluator.swift

@@ -32,24 +32,42 @@ internal struct IREvaluator {
     init(policies: [IR.Policy]) {
         self.policies = policies.map { IndexedIRPolicy(policy: $0) }
     }
+
+    func findPlanAndPolicy(entrypoint: String) -> (IndexedIRPolicy, Plan)? {
+        // TODO: We're assuming that queries are only ever defined in a single policy... that _should_ hold true.. but who's checkin?
+        guard let policy = policies.first(where: { $0.plans[entrypoint] != nil }) else {
+            return nil
+        }
+
+        // The above guard has already ensured that the policy has a plan for the given entrypoint, so force unwrap is safe.
+        return (policy, policy.plans[entrypoint]!)
+    }
+    
+    func unknownQueryError(query: String) -> RegoError {
+        RegoError(code: .unknownQuery, message: "query not found in plan: \(query)")
+    }
 }
 
 extension IREvaluator: Evaluator {
-    func evaluate(withContext ctx: EvaluationContext) async throws -> ResultSet {
-        // TODO: We're assuming that queries are only ever defined in a single policy... that _should_ hold true.. but who's checkin?
+    func ensureQueryIsSupported(_ query: String) throws {
+        let entrypoint = try queryToEntryPoint(query)
 
+        guard findPlanAndPolicy(entrypoint: entrypoint) != nil else {
+            throw unknownQueryError(query: query)
+        }
+    }
+    
+    func evaluate(withContext ctx: EvaluationContext) async throws -> ResultSet {
         let entrypoint = try queryToEntryPoint(ctx.query)
 
-        for policy in policies {
-            if let plan = policy.plans[entrypoint] {
-                let ctx = IREvaluationContext(ctx: ctx, policy: policy)
-                return try await evalPlan(
-                    withContext: ctx,
-                    plan: plan
-                )
-            }
+        guard let (policy, plan) = findPlanAndPolicy(entrypoint: entrypoint) else {
+            throw unknownQueryError(query: ctx.query)
         }
-        throw RegoError(code: .unknownQuery, message: "query not found in plan: \(ctx.query)")
+
+        return try await evalPlan(
+            withContext: IREvaluationContext(ctx: ctx, policy: policy),
+            plan: plan
+        )
     }
 }