ensure query can be evaluated during preparation

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,43 @@ 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?
+        for policy in policies {
+            if let plan = policy.plans[entrypoint] {
+                return (policy, plan)
+            }
+        }
+
+        return nil
+    }
+
+    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
+        )
     }
 }