diff --git a/Sources/Rego/Engine.swift b/Sources/Rego/Engine.swift index b870e57..307dfb9 100644 --- a/Sources/Rego/Engine.swift +++ b/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 ) } diff --git a/Sources/Rego/Evaluator.swift b/Sources/Rego/Evaluator.swift index c76a8ec..7c954f9 100644 --- a/Sources/Rego/Evaluator.swift +++ b/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 } diff --git a/Sources/Rego/IREvaluator.swift b/Sources/Rego/IREvaluator.swift index ffca6a4..b4bd38a 100644 --- a/Sources/Rego/IREvaluator.swift +++ b/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 + ) } }