maglev-rb 0.1.1 → 0.2.1

data/lib/maglev/planner.rb

@@ -0,0 +1,172 @@
+# frozen_string_literal: true
+
+require "securerandom"
+
+require_relative "planner_adapter"
+require_relative "query_validator"
+require_relative "trace"
+
+module Maglev
+  class Planner
+    QUERY_IR_SCHEMA = {
+      type: "object", additionalProperties: false,
+      required: %w[version root operation scopes filters joins sort distinct limit],
+      properties: {
+        version: {const: 1}, root: {type: "string"}, operation: {enum: %w[records aggregate]},
+        scopes: {type: "array", items: {type: "object", additionalProperties: false,
+                                        required: %w[name parameters], properties: {name: {type: "string"}, parameters: {type: "object"}}}},
+        filters: {type: "array", items: {type: "object", additionalProperties: false,
+                                         required: %w[field operator value], properties: {field: {type: "string"}, operator: {type: "string"}, value: {}}}},
+        joins: {type: "array", items: {type: "string"}},
+        sort: {type: "array", items: {type: "object", additionalProperties: false,
+                                      required: %w[field direction], properties: {field: {type: "string"}, direction: {enum: %w[asc desc]}}}},
+        distinct: {type: "boolean"}, limit: {type: "integer", minimum: 1},
+        aggregate: {type: ["object", "null"]}
+      }
+    }.freeze
+    OUTCOMES = %w[ready clarification_required unsupported].freeze
+
+    Plan = Struct.new(:status, :route, :ir, :resource, :explanation, :warnings, :errors,
+      :clarification, :constraints, :trace_id, :validation, :base_relation,
+      :evidence_requirements, :policy_limits) do
+      def initialize(**attributes)
+        attributes[:status] = attributes.fetch(:status).to_sym
+        attributes[:route] = (attributes[:route] || :structured).to_sym
+        attributes[:warnings] = Array(attributes[:warnings]).freeze
+        attributes[:errors] = Array(attributes[:errors]).freeze
+        attributes[:constraints] = (attributes[:constraints] || {}).freeze
+        attributes[:evidence_requirements] = (attributes[:evidence_requirements] || {}).freeze
+        attributes[:policy_limits] = (attributes[:policy_limits] || {}).freeze
+        attributes[:clarification] = attributes[:clarification]&.freeze
+        super
+        freeze
+      end
+
+      def ready? = status == :ready
+    end
+
+    def initialize(adapter:)
+      unless adapter.respond_to?(:plan)
+        raise ConfigurationError, "planner adapter must implement #plan"
+      end
+
+      @adapter = adapter
+    end
+
+    def plan(question:, snapshot:, resource:, constraints: {}, base_relation: nil)
+      root = resource.to_s
+      trace_id = SecureRandom.uuid
+      request_constraints = normalize_constraints(constraints)
+      output = Trace.instrument(:planning, trace_id: trace_id, resource: root) do |payload|
+        provider_plan(question, snapshot, request_constraints, nil).tap do |result|
+          payload[:status] = result.is_a?(Hash) ? result["status"]&.to_sym || :invalid : :invalid
+        end
+      end
+      return outcome_plan(output, root, request_constraints, trace_id) unless output.is_a?(Hash) && output["status"] == "ready"
+
+      validation = Trace.instrument(:validation, trace_id: trace_id, resource: root,
+        operation: output.dig("ir", "operation")) do |payload|
+        validate(output["ir"], snapshot, root, request_constraints).tap do |result|
+          payload[:status] = result.valid? ? :ready : :invalid
+          payload[:error_codes] = result.errors.map(&:code).uniq.sort if result.errors.any?
+        end
+      end
+      unless validation.valid?
+        repair = {errors: safe_errors(validation.errors)}.freeze
+        output = Trace.instrument(:planning, trace_id: trace_id, resource: root) do |payload|
+          provider_plan(question, snapshot, request_constraints, repair).tap do |result|
+            payload[:status] = result.is_a?(Hash) ? result["status"]&.to_sym || :invalid : :invalid
+          end
+        end
+        return outcome_plan(output, root, request_constraints, trace_id) unless output.is_a?(Hash) && output["status"] == "ready"
+        validation = Trace.instrument(:validation, trace_id: trace_id, resource: root,
+          operation: output.dig("ir", "operation")) do |payload|
+          validate(output["ir"], snapshot, root, request_constraints).tap do |result|
+            payload[:status] = result.valid? ? :ready : :invalid
+            payload[:error_codes] = result.errors.map(&:code).uniq.sort if result.errors.any?
+          end
+        end
+      end
+
+      return invalid_plan(root, request_constraints, validation.errors, trace_id) unless validation.valid?
+
+      evidence_requirements = {kind: validation.ir.operation, max_rows: validation.ir.limit}.freeze
+      Plan.new(status: :ready, ir: validation.ir, resource: root, explanation: validation.explanation,
+        validation: validation, constraints: request_constraints, trace_id: trace_id,
+        base_relation: base_relation, evidence_requirements: evidence_requirements,
+        policy_limits: validator_limits(snapshot, root, request_constraints))
+    end
+
+    private
+
+    def provider_plan(question, snapshot, constraints, repair)
+      @adapter.plan(question: question.to_s, schema_snapshot: snapshot, constraints: constraints,
+        query_ir_schema: QUERY_IR_SCHEMA, repair: repair)
+    end
+
+    def validate(input, snapshot, root, constraints)
+      QueryValidator.new(snapshot: snapshot, root: root, limits: constraints).call(input)
+    end
+
+    def validator_limits(snapshot, root, constraints)
+      QueryValidator.new(snapshot: snapshot, root: root, limits: constraints).policy_limits
+    end
+
+    def normalize_constraints(constraints)
+      values = constraints.transform_keys(&:to_sym)
+      allowed = QueryValidator::DEFAULT_LIMITS.keys
+      unless (values.keys - allowed).empty? && values.values.all? { |value| value.is_a?(Integer) && value.positive? }
+        raise ArgumentError, "invalid planner constraints"
+      end
+      values.freeze
+    end
+
+    def safe_errors(errors)
+      errors.map { |error| {code: error.code, message: error.message, path: error.path}.freeze }.freeze
+    end
+
+    def outcome_plan(output, root, constraints, trace_id)
+      return invalid_plan(root, constraints, [], trace_id) unless output.is_a?(Hash) && OUTCOMES.include?(output["status"])
+
+      status = output["status"].to_sym
+      case status
+      when :clarification_required
+        message = output["message"]
+        choices = output["choices"]
+        return invalid_plan(root, constraints, [], trace_id) unless message.is_a?(String) && choices.is_a?(Array) &&
+          choices.length.between?(1, 10) && choices.all? { |choice| choice.is_a?(String) && choice.bytesize <= 200 }
+        Plan.new(status: status, resource: root, clarification: {message: message, choices: choices.freeze},
+          constraints: constraints, trace_id: trace_id)
+      when :unsupported
+        return invalid_plan(root, constraints, [], trace_id) unless output["message"].is_a?(String)
+        Plan.new(status: status, resource: root, warnings: [output["message"]], constraints: constraints,
+          trace_id: trace_id)
+      else
+        invalid_plan(root, constraints, [], trace_id)
+      end
+    end
+
+    def invalid_plan(root, constraints, errors, trace_id)
+      Plan.new(status: :invalid, resource: root, errors: errors, constraints: constraints,
+        trace_id: trace_id)
+    end
+  end
+
+  def self.plan(question, resource:, base_relation:, resources: nil, constraints: {}, user: nil,
+    authorizer: nil, adapter: Maglev.configuration.planner_adapter)
+    raise ConfigurationError, "planner adapter is not configured" unless adapter
+    unless defined?(ActiveRecord::Relation) && base_relation.is_a?(ActiveRecord::Relation)
+      raise ConfigurationError, "an ActiveRecord base relation is required for planning"
+    end
+
+    identifiers = [resource, *Array(resources)].compact.map(&:to_s).uniq
+    snapshot = Registry.snapshot(resources: identifiers, user: user, authorizer: authorizer)
+    root_model = snapshot.model_class_for(resource)
+    unless root_model && base_relation.klass == root_model
+      raise ConfigurationError, "base relation does not match the authorized resource"
+    end
+
+    Planner.new(adapter: adapter).plan(question: question, snapshot: snapshot,
+      resource: resource, constraints: constraints, base_relation: base_relation)
+  end
+end

data/lib/maglev/planner_adapter.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module Maglev
+  class PlannerAdapter
+    def plan(question:, schema_snapshot:, constraints:, query_ir_schema:, repair: nil)
+      raise NotImplementedError, "#{self.class.name} must implement #plan"
+    end
+  end
+
+  class FakePlannerAdapter < PlannerAdapter
+    attr_reader :requests
+
+    def initialize(outputs)
+      @outputs = Array(outputs).dup
+      @requests = []
+    end
+
+    def plan(**request)
+      @requests << request.freeze
+      raise PermanentProviderError, "Fake planner has no remaining output" if @outputs.empty?
+
+      @outputs.shift
+    end
+  end
+end

data/lib/maglev/planner_evaluation.rb

@@ -0,0 +1,49 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Maglev
+  module PlannerEvaluation
+    module_function
+
+    def score(cases)
+      results = cases.map { |test_case| score_case(test_case) }.freeze
+      passed = results.count { |result| result[:passed] }
+      total = results.length
+      {total: total, passed: passed, failed: total - passed,
+       score: total.zero? ? 0.0 : passed.fdiv(total), cases: results}.freeze
+    end
+
+    def load(path)
+      corpus = JSON.parse(File.read(path))
+      raise ArgumentError, "unsupported planner evaluation corpus" unless corpus["version"] == 1 && corpus["cases"].is_a?(Array)
+
+      corpus
+    end
+
+    def score_case(test_case)
+      expected = test_case.fetch("expected")
+      actual = test_case.fetch("actual")
+      failure_class = if expected["status"] != actual["status"]
+        :status_mismatch
+      elsif expected["status"] == "ready" && canonical(expected["ir"]) != canonical(actual["ir"])
+        :ir_mismatch
+      end
+      {id: test_case.fetch("id"), passed: failure_class.nil?, failure_class: failure_class}.freeze
+    end
+    private_class_method :score_case
+
+    def canonical(value, parent = nil)
+      case value
+      when Hash
+        value.keys.sort.to_h { |key| [key, canonical(value[key], key)] }
+      when Array
+        items = value.map { |item| canonical(item) }
+        %w[filters joins].include?(parent) ? items.sort_by { |item| JSON.generate(item) } : items
+      else
+        value
+      end
+    end
+    private_class_method :canonical
+  end
+end

data/lib/maglev/query_compiler.rb

@@ -0,0 +1,197 @@
+# frozen_string_literal: true
+
+require_relative "errors"
+require_relative "query_validator"
+
+module Maglev
+  class QueryCompilationError < Error
+  end
+
+  class StructuredPlan
+    attr_reader :relation, :operations, :explanation, :warnings, :aggregate, :aggregate_column
+
+    def initialize(relation:, operations:, explanation:, aggregate: nil, aggregate_column: nil, warnings: [])
+      @relation = relation
+      @operations = operations.freeze
+      @explanation = explanation.to_s.freeze
+      @warnings = warnings.freeze
+      @aggregate = aggregate
+      @aggregate_column = aggregate_column
+      freeze
+    end
+
+    def records? = aggregate.nil?
+    def aggregate? = !records?
+    def executed? = false
+    def to_sql = relation.to_sql
+  end
+
+  module ReadOnlyRelation
+    MUTATION_METHODS = %i[
+      build create create! create_or_find_by create_or_find_by! create_with create_with!
+      delete delete_all delete_by destroy destroy_all destroy_by find_or_create_by
+      find_or_create_by! first_or_create first_or_create! insert insert! insert_all
+      insert_all! new touch_all update update! update_all update_counters upsert upsert_all
+    ].freeze
+
+    MUTATION_METHODS.each do |method_name|
+      define_method(method_name) do |*|
+        raise QueryCompilationError, "Structured query relations are read-only"
+      end
+    end
+  end
+
+  class QueryCompiler
+    SAFE_SCOPE_KEYS = %i[where order limit distinct].freeze
+
+    def initialize(snapshot:)
+      @snapshot = snapshot
+      @resources = snapshot.resources.to_h { |resource| [resource.identifier, resource] }
+    end
+
+    def compile(validation:, base_relation:)
+      unless validation.is_a?(QueryValidator::Result) && validation.valid? && validation.ir && validation.snapshot.equal?(@snapshot)
+        raise QueryCompilationError, "A valid query validation result is required"
+      end
+      unless defined?(ActiveRecord::Relation) && base_relation.is_a?(ActiveRecord::Relation)
+        raise QueryCompilationError, "An ActiveRecord base relation is required"
+      end
+
+      ir = validation.ir
+      resource = @resources[ir.root]
+      root_model = @snapshot.model_class_for(ir.root)
+      unless resource && root_model && base_relation.klass == root_model
+        raise QueryCompilationError, "The base relation does not match the validated root resource"
+      end
+
+      relation = base_relation
+      ir.scopes.each { |scope| relation = apply_scope(relation, resource, scope) }
+      ir.joins.each { |path| relation = relation.joins(join_argument(path.segments)) }
+      ir.filters.each { |predicate| relation = apply_predicate(relation, root_model, predicate) }
+      ir.sort.each { |sort| relation = apply_sort(relation, root_model, sort) }
+      relation = relation.distinct if ir.distinct
+      relation = relation.limit([relation.limit_value, ir.limit].compact.min)
+      relation = relation.readonly.extending(ReadOnlyRelation)
+
+      StructuredPlan.new(
+        relation: relation,
+        aggregate: ir.aggregate,
+        aggregate_column: aggregate_column(root_model, ir.aggregate),
+        operations: operation_descriptions(ir),
+        explanation: validation.explanation
+      )
+    rescue QueryCompilationError
+      raise
+    rescue NoMethodError, ArgumentError, ActiveRecord::ActiveRecordError => error
+      raise QueryCompilationError, "Validated query could not be compiled: #{error.class.name}"
+    end
+
+    private
+
+    def apply_scope(relation, resource, scope)
+      declaration = resource.scopes.find { |candidate| candidate.fetch(:name) == scope.name }
+      raise QueryCompilationError, "The validated scope is unavailable" unless declaration
+
+      values = declaration.fetch(:parameters).keys.map { |name| literal_value(scope.parameters.fetch(name)) }
+      scoped = ActiveRecord::Base.while_preventing_writes do
+        relation.public_send(scope.name, *values)
+      end
+      if scoped.respond_to?(:unscope_values) && scoped.unscope_values.any?
+        raise QueryCompilationError, "Registered scope cannot remove relation constraints"
+      end
+      unless scoped.is_a?(ActiveRecord::Relation) && scoped.klass == relation.klass
+        raise QueryCompilationError, "Registered scope returned an incompatible relation"
+      end
+      raise QueryCompilationError, "Registered scope cannot widen the base relation" unless preserves_relation?(relation, scoped)
+
+      scoped
+    rescue ActiveRecord::ReadOnlyError
+      raise QueryCompilationError, "Registered scope cannot widen the base relation or introduce unsafe clauses"
+    end
+
+    def apply_predicate(relation, root_model, predicate)
+      model = model_for_path(root_model, predicate.field.segments)
+      column = model.arel_table[predicate.field.segments.last]
+      value = literal_value(predicate.value)
+      node = case predicate.operator
+      when :eq then column.eq(value)
+      when :not_eq then column.not_eq(value)
+      when :gt then column.gt(value)
+      when :gte then column.gteq(value)
+      when :lt then column.lt(value)
+      when :lte then column.lteq(value)
+      when :in then column.in(value)
+      when :not_in then column.not_in(value)
+      when :is_null then column.eq(nil)
+      when :is_not_null then column.not_eq(nil)
+      when :between then column.between(value.first..value.last)
+      else raise QueryCompilationError, "The validated predicate is unavailable"
+      end
+      relation.where(node)
+    end
+
+    def apply_sort(relation, root_model, sort)
+      model = model_for_path(root_model, sort.field.segments)
+      column = model.arel_table[sort.field.segments.last]
+      relation.order((sort.direction == :desc) ? column.desc : column.asc)
+    end
+
+    def model_for_path(root_model, segments)
+      segments[0...-1].reduce(root_model) { |model, name| model.reflect_on_association(name).klass }
+    end
+
+    def aggregate_column(root_model, aggregate)
+      return unless aggregate&.field
+
+      model_for_path(root_model, aggregate.field.segments).arel_table[aggregate.field.segments.last]
+    end
+
+    def preserves_relation?(base, scoped)
+      base_values = base.values
+      scoped_values = scoped.values
+      return false unless (base_values.keys - scoped_values.keys).empty?
+      return false unless (scoped_values.keys - base_values.keys - SAFE_SCOPE_KEYS).empty?
+      return false unless base.where_clause.send(:predicates).all? { |predicate| scoped.where_clause.send(:predicates).include?(predicate) }
+      return false unless Array(base.joins_values).all? { |join| scoped.joins_values.include?(join) }
+      return false unless Array(base.left_outer_joins_values).all? { |join| scoped.left_outer_joins_values.include?(join) }
+      return false unless Array(base.order_values).all? { |order| scoped.order_values.include?(order) }
+      return false if base.limit_value && (!scoped.limit_value || scoped.limit_value > base.limit_value)
+      return false if base.distinct_value && !scoped.distinct_value
+
+      protected_keys = base_values.keys - SAFE_SCOPE_KEYS
+      protected_keys.all? { |key| base_values[key] == scoped_values[key] }
+    end
+
+    def join_argument(segments)
+      return segments.first.to_sym if segments.one?
+
+      segments.reverse.reduce { |nested, segment| {segment.to_sym => nested} }
+    end
+
+    def literal_value(value)
+      case value
+      when QueryIR::Literal
+        value.value
+      when QueryIR::RelativeTime
+        amount = value.amount.public_send(value.unit)
+        (value.direction == "ago") ? Time.current - amount : Time.current + amount
+      when Array
+        value.map { |item| literal_value(item) }
+      else
+        value
+      end
+    end
+
+    def operation_descriptions(ir)
+      [
+        *ir.scopes.map { |scope| "scope #{scope.name}" },
+        *ir.joins.map { |join| "join #{join}" },
+        *ir.filters.map { |filter| "filter #{filter.field} #{filter.operator}" },
+        *ir.sort.map { |sort| "sort #{sort.field} #{sort.direction}" },
+        ("distinct" if ir.distinct),
+        "limit #{ir.limit}",
+        ("aggregate #{ir.aggregate.function}" if ir.aggregate)
+      ].compact.freeze
+    end
+  end
+end

data/lib/maglev/query_ir.rb

@@ -0,0 +1,143 @@
+# frozen_string_literal: true
+
+require "json"
+
+module Maglev
+  module QueryIR
+    VERSION = 1
+
+    module Value
+      def to_json(*arguments)
+        JSON.generate(to_h, *arguments)
+      end
+    end
+
+    class Path
+      include Value
+
+      attr_reader :segments
+
+      def initialize(value)
+        @segments = value.to_s.split(".").map { |segment| segment.dup.freeze }.freeze
+        freeze
+      end
+
+      def to_s = segments.join(".")
+      def to_h = to_s
+    end
+
+    class Literal
+      include Value
+
+      attr_reader :value, :type
+
+      def initialize(value:, type:)
+        @value = value.freeze
+        @type = type.to_sym
+        freeze
+      end
+
+      def to_h = value
+    end
+
+    class RelativeTime
+      include Value
+
+      UNITS = %w[seconds minutes hours days weeks months years].freeze
+      DIRECTIONS = %w[ago from_now].freeze
+      attr_reader :amount, :unit, :direction
+
+      def initialize(amount:, unit:, direction:)
+        @amount = amount
+        @unit = unit.to_s.freeze
+        @direction = direction.to_s.freeze
+        freeze
+      end
+
+      def to_h = {"relative" => {"amount" => amount, "unit" => unit, "direction" => direction}.freeze}.freeze
+    end
+
+    Predicate = Struct.new(:field, :operator, :value) do
+      include Value
+
+      def initialize(**attributes) = super.tap { freeze }
+
+      def to_h
+        result = {"field" => field.to_s, "operator" => operator.to_s}
+        unless value.nil?
+          result["value"] = if value.is_a?(Array)
+            value.map { |item| item.respond_to?(:to_h) ? item.to_h : item }
+          elsif value.respond_to?(:to_h)
+            value.to_h
+          else
+            value
+          end
+        end
+        result.freeze
+      end
+    end
+    Scope = Struct.new(:name, :parameters) do
+      include Value
+
+      def initialize(**attributes)
+        attributes[:parameters] = attributes.fetch(:parameters).freeze
+        super
+        freeze
+      end
+
+      def to_h = {"name" => name.to_s, "parameters" => parameters.transform_values { |value| value.respond_to?(:to_h) ? value.to_h : value }.freeze}.freeze
+    end
+    Sort = Struct.new(:field, :direction) do
+      include Value
+
+      def initialize(**attributes) = super.tap { freeze }
+      def to_h = {"field" => field.to_s, "direction" => direction.to_s}.freeze
+    end
+    Aggregate = Struct.new(:function, :field) do
+      include Value
+
+      def initialize(**attributes) = super.tap { freeze }
+      def to_h = {"function" => function.to_s}.tap { |hash| hash["field"] = field.to_s if field }.freeze
+    end
+
+    class Query
+      include Value
+
+      attr_reader :version, :root, :operation, :scopes, :filters, :joins, :sort, :distinct, :limit, :aggregate
+
+      def initialize(version:, root:, operation:, scopes:, filters:, joins:, sort:, distinct:, limit:, aggregate: nil)
+        @version = version
+        @root = root.to_s.freeze
+        @operation = operation.to_sym
+        @scopes = scopes.freeze
+        @filters = filters.freeze
+        @joins = joins.freeze
+        @sort = sort.freeze
+        @distinct = distinct
+        @limit = limit
+        @aggregate = aggregate
+        freeze
+      end
+
+      def to_h
+        hash = {"version" => version, "root" => root, "operation" => operation.to_s,
+                "scopes" => scopes.map(&:to_h), "filters" => filters.map(&:to_h),
+                "joins" => joins.map(&:to_s), "sort" => sort.map(&:to_h),
+                "distinct" => distinct, "limit" => limit}
+        hash["aggregate"] = aggregate.to_h if aggregate
+        hash.freeze
+      end
+    end
+
+    def self.from_json(json)
+      data = JSON.parse(json)
+      Query.new(version: data.fetch("version"), root: data.fetch("root"), operation: data.fetch("operation"),
+        scopes: data.fetch("scopes").map { |scope| Scope.new(name: scope.fetch("name"), parameters: scope.fetch("parameters")) },
+        filters: data.fetch("filters").map { |filter| Predicate.new(field: Path.new(filter.fetch("field")), operator: filter.fetch("operator").to_sym, value: filter["value"]) },
+        joins: data.fetch("joins").map { |path| Path.new(path) },
+        sort: data.fetch("sort").map { |sort| Sort.new(field: Path.new(sort.fetch("field")), direction: sort.fetch("direction").to_sym) },
+        distinct: data.fetch("distinct"), limit: data.fetch("limit"),
+        aggregate: data["aggregate"] && Aggregate.new(function: data.dig("aggregate", "function").to_sym, field: data.dig("aggregate", "field") && Path.new(data.dig("aggregate", "field"))))
+    end
+  end
+end