`GraphQL::Schema::Interface#resolve_type` appears to resolve lazy results too early

[BytewaveMLP]

Describe the bug

When lazily-loading values within resolve_type, the promise appears to be immediately synced. This causes N+1 queries when using graphql-batch, where the result of a resolve_type depends on another record.

Versions

graphql version: 2.2.8
rails (or other framework): ActiveRecord 7.1.3
other applicable versions (graphql-batch, etc): graphql_batch 0.5.4

Example test cases

resolve_type

# frozen_string_literal: true

require 'bundler/inline'

gemfile(true) do
  source 'https://rubygems.org'

  gem 'activerecord'

  gem 'sqlite3'

  gem 'graphql'
  gem 'graphql-batch'
end

require 'active_record'
require 'minitest/autorun'
require 'logger'

$log = Logger.new($stdout)
ActiveRecord::Base.establish_connection(adapter: 'sqlite3', database: ':memory:')
ActiveRecord::Base.logger = $log

ActiveRecord::Schema.define do
  create_table :versionables, force: true do |t|
    t.string :type, null: false
  end

  create_table :versions, force: true do |t|
    t.references :versionable, null: false, foreign_key: true

    t.string :foo
    t.string :bar
  end

  create_table :version_references, force: true do |t|
    t.references :version, null: false, foreign_key: true
  end
end

class ApplicationRecord < ActiveRecord::Base
  self.abstract_class = true
end

class Versionable < ApplicationRecord
  has_many :versions
end

class FooVersionable < Versionable
end

class BarVersionable < Versionable
end

class Version < ApplicationRecord
  belongs_to :versionable

  has_many :version_references
end

class VersionReference < ApplicationRecord
  belongs_to :version
end

# Loader taken directly from shopify/graphql-batch
class AssociationLoader < GraphQL::Batch::Loader
  def self.validate(model, association_name)
    new(model, association_name)
    nil
  end

  def initialize(model, association_name)
    super()
    @model = model
    @association_name = association_name
    validate
  end

  def load(record)
    raise TypeError, "#{@model} loader can't load association for #{record.class}" unless record.is_a?(@model)
    return Promise.resolve(read_association(record)) if association_loaded?(record)
    super
  end

  # We want to load the associations on all records, even if they have the same id
  def cache_key(record)
    record.object_id
  end

  def perform(records)
    preload_association(records)
    records.each { |record| fulfill(record, read_association(record)) }
  end

  private

  def validate
    unless @model.reflect_on_association(@association_name)
      raise ArgumentError, "No association #{@association_name} on #{@model}"
    end
  end

  def preload_association(records)
    ::ActiveRecord::Associations::Preloader.new(records: records, associations: @association_name).call
  end

  def read_association(record)
    record.public_send(@association_name)
  end

  def association_loaded?(record)
    record.association(@association_name).loaded?
  end
end

module VersionType
  include GraphQL::Schema::Interface

  def self.resolve_type(object, _ctx)
    $log.debug "self.resolve_type(#{object.inspect}, _ctx)"
    AssociationLoader.for(Version, :versionable).load(object).then do |versionable|
      $log.debug "Promise resolved to #{versionable.class}"

      case versionable
      when FooVersionable
        FooVersionType
      when BarVersionable
        BarVersionType
      end
    end
  end

  field :id, ID, null: false
end

class FooVersionType < GraphQL::Schema::Object
  implements VersionType

  field :foo, String, null: true
end

class BarVersionType < GraphQL::Schema::Object
  implements VersionType

  field :bar, String, null: true
end

class VersionReferenceType < GraphQL::Schema::Object
  field :id, ID, null: false
  field :version, VersionType, null: false

  def version
    AssociationLoader.for(VersionReference, :version).load(object)
  end
end

class QueryType < GraphQL::Schema::Object
  field :version_references, [VersionReferenceType], null: false

  def version_references
    VersionReference.all
  end
end

class MySchema < GraphQL::Schema
  query(QueryType)

  use(GraphQL::Batch)

  orphan_types(FooVersionType, BarVersionType)

  def self.resolve_type(_type, object, _context)
    "::#{object.class.name}Type".constantize
  end
end

class BugTest < Minitest::Test
  def test_dataloading_bug
    foo_versionable = FooVersionable.create!
    bar_versionable = BarVersionable.create!

    foo_version = foo_versionable.versions.create!(foo: 'foo')
    bar_version = bar_versionable.versions.create!(bar: 'bar')

    foo_ref = VersionReference.create!(version: foo_version)
    bar_ref = VersionReference.create!(version: bar_version)

    result = MySchema.execute(<<~QUERY)
      query VersionReferences {
        versionReferences {
          id
          version {
            id
            __typename
            ... on FooVersion {
              foo
            }
            ... on BarVersion {
              bar
            }
          }
        }
      }
    QUERY

    assert_equal(
      {
        'data' => {
          'versionReferences' => [
            {
              'id'      => foo_ref.id.to_s,
              'version' => {
                'id'         => foo_version.id.to_s,
                '__typename' => 'FooVersion',
                'foo'        => 'foo'
              }
            },
            {
              'id'      => bar_ref.id.to_s,
              'version' => {
                'id'         => bar_version.id.to_s,
                '__typename' => 'BarVersion',
                'bar'        => 'bar'
              }
            }
          ]
        }
      },
      result.to_h
    )
  end
end

Standard field-based lazy loading, for comparison

# frozen_string_literal: true

require 'bundler/inline'

gemfile(true) do
  source 'https://rubygems.org'

  gem 'activerecord'

  gem 'sqlite3'

  gem 'graphql'
  gem 'graphql-batch'
end

require 'active_record'
require 'minitest/autorun'
require 'logger'

$log = Logger.new($stdout)
ActiveRecord::Base.establish_connection(adapter: 'sqlite3', database: ':memory:')
ActiveRecord::Base.logger = $log

ActiveRecord::Schema.define do
  create_table :versionables, force: true do |t|
  end

  create_table :versions, force: true do |t|
    t.references :versionable, null: false, foreign_key: true
  end

  create_table :version_references, force: true do |t|
    t.references :version, null: false, foreign_key: true
  end
end

class ApplicationRecord < ActiveRecord::Base
  self.abstract_class = true
end

class Versionable < ApplicationRecord
  has_many :versions
end

class Version < ApplicationRecord
  belongs_to :versionable

  has_many :version_references
end

class VersionReference < ApplicationRecord
  belongs_to :version
end

# Loader taken directly from shopify/graphql-batch
class AssociationLoader < GraphQL::Batch::Loader
  def self.validate(model, association_name)
    new(model, association_name)
    nil
  end

  def initialize(model, association_name)
    super()
    @model = model
    @association_name = association_name
    validate
  end

  def load(record)
    raise TypeError, "#{@model} loader can't load association for #{record.class}" unless record.is_a?(@model)
    return Promise.resolve(read_association(record)) if association_loaded?(record)
    super
  end

  # We want to load the associations on all records, even if they have the same id
  def cache_key(record)
    record.object_id
  end

  def perform(records)
    preload_association(records)
    records.each { |record| fulfill(record, read_association(record)) }
  end

  private

  def validate
    unless @model.reflect_on_association(@association_name)
      raise ArgumentError, "No association #{@association_name} on #{@model}"
    end
  end

  def preload_association(records)
    ::ActiveRecord::Associations::Preloader.new(records: records, associations: @association_name).call
  end

  def read_association(record)
    record.public_send(@association_name)
  end

  def association_loaded?(record)
    record.association(@association_name).loaded?
  end
end

class VersionableType < GraphQL::Schema::Object
  field :id, ID, null: false
end

class VersionType < GraphQL::Schema::Object
  field :id, ID, null: false
  field :versionable, VersionableType, null: false

  def versionable
    AssociationLoader.for(Version, :versionable).load(object)
  end
end

class VersionReferenceType < GraphQL::Schema::Object
  field :id, ID, null: false
  field :version, VersionType, null: false

  def version
    AssociationLoader.for(VersionReference, :version).load(object)
  end
end

class QueryType < GraphQL::Schema::Object
  field :version_references, [VersionReferenceType], null: false

  def version_references
    VersionReference.all
  end
end

class MySchema < GraphQL::Schema
  query(QueryType)

  use(GraphQL::Batch)

  def self.resolve_type(_type, object, _context)
    "::#{object.class.name}Type".constantize
  end
end

class BugTest < Minitest::Test
  def test_dataloading_bug
    foo_versionable = Versionable.create!
    bar_versionable = Versionable.create!

    foo_version = foo_versionable.versions.create!
    bar_version = bar_versionable.versions.create!

    foo_ref = VersionReference.create!(version: foo_version)
    bar_ref = VersionReference.create!(version: bar_version)

    result = MySchema.execute(<<~QUERY)
      query VersionReferences {
        versionReferences {
          id
          version {
            id
            versionable {
              id
            }
          }
        }
      }
    QUERY

    assert_equal(
      {
        'data' => {
          'versionReferences' => [
            {
              'id'      => foo_ref.id.to_s,
              'version' => {
                'id'          => foo_version.id.to_s,
                'versionable' => {
                  'id' => foo_versionable.id.to_s
                }
              }
            },
            {
              'id'      => bar_ref.id.to_s,
              'version' => {
                'id'          => bar_version.id.to_s,
                'versionable' => {
                  'id' => bar_versionable.id.to_s
                }
              }
            }
          ]
        }
      },
      result.to_h
    )
  end
end

Expected behavior

The resolve_type example above should only issue 3 queries, same as the "standard" test.

Log output from "standard" test

D, [2024-02-08T16:19:28.193666 #58311] DEBUG -- :   VersionReference Load (0.0ms)  SELECT "version_references".* FROM "version_references"
D, [2024-02-08T16:19:28.196335 #58311] DEBUG -- :   Version Load (0.0ms)  SELECT "versions".* FROM "versions" WHERE "versions"."id" IN (?, ?)  [["id", 1], ["id", 2]]
D, [2024-02-08T16:19:28.196778 #58311] DEBUG -- :   Versionable Load (0.0ms)  SELECT "versionables".* FROM "versionables" WHERE "versionables"."id" IN (?, ?)  [["id", 1], ["id", 2]]

Actual behavior

The resolve_type test issues one query per resolve_type call, and output implies the promise is being immediately synced upon return from resolve_type.

Log output from "resolve_type" example

D, [2024-02-08T16:07:04.550537 #58025] DEBUG -- :   VersionReference Load (0.0ms)  SELECT "version_references".* FROM "version_references"
D, [2024-02-08T16:07:04.553402 #58025] DEBUG -- :   Version Load (0.1ms)  SELECT "versions".* FROM "versions" WHERE "versions"."id" IN (?, ?)  [["id", 1], ["id", 2]]
D, [2024-02-08T16:07:04.553569 #58025] DEBUG -- : self.resolve_type(#<Version id: 1, versionable_id: 1, foo: "foo", bar: nil>, _ctx)
D, [2024-02-08T16:07:04.553905 #58025] DEBUG -- :   Versionable Load (0.0ms)  SELECT "versionables".* FROM "versionables" WHERE "versionables"."id" = ?  [["id", 1]]
D, [2024-02-08T16:07:04.553997 #58025] DEBUG -- : Promise resolved to FooVersionable
D, [2024-02-08T16:07:04.554151 #58025] DEBUG -- : self.resolve_type(#<Version id: 2, versionable_id: 2, foo: nil, bar: "bar">, _ctx)
D, [2024-02-08T16:07:04.554310 #58025] DEBUG -- :   Versionable Load (0.0ms)  SELECT "versionables".* FROM "versionables" WHERE "versionables"."id" = ?  [["id", 2]]
D, [2024-02-08T16:07:04.554364 #58025] DEBUG -- : Promise resolved to BarVersionable

Additional context

I'm not sure if resolve_type is intended to be used in this way, but as far as I can tell from reading over previous issues and PRs, lazy resolution is an intended feature. If this is expected behavior, or if it would be better reported to graphql-batch, feel free to let me know and close this issue.

Also, I recognize the schema presented is rather cursed. However, it mirrors a similar example to something we're encountering in our application; we have a versioning system which can have one of two STI types as its parent, and we want to transparently return this versioned data to our clients. In order to know which type to return however, we must know the type of the parent, which requires loading the associated parent within resolve_type.

[rmosolgo]

Hey, thanks for the detailed write-up. This is the right place for the issue -- GraphQL-Batch has done its job, now it's up to GraphQL-Ruby to resolve those promises as best it can!

I bet this can be improved in GraphQL-Ruby's runtime somehow. I'll take a closer look soon and work up a small replication based on your script above, then see about improving the promise resolving code to make it work better, and follow up here 👍

[rmosolgo]

Thanks again for the detailed report. I copied your example into the GraphQL-Ruby tests and spent some time trying to understand why it's not working "properly." I couldn't quite grok what needs to be better ... but I'll keep trying!