Relationship loading with federation

I’m trying to use federation where the second instance extends on the data provided by the first. Let’s imagine a simple setup where you have two instances, each with a data source (via JPA): A for companies and B for employees. A provides only the data about the companies and is unaware of the employees. B on the other hand has methods that extend the Company model to return all the Employees underneath that company, while Employee also has a relation back to Company to return their current employer.

Several classes on instance B:

// Company

@KeyDirective(fields = FieldSet("id"))
@ExtendsDirective
data class Company(
    @ExternalDirective val id: Int
) {
    @OneToMany(mappedBy = "company")
    lateinit var employees: List<Employee>
}

// Employee

@Entity
@Table(name = "employees")
@KeyDirective(fields = FieldSet("id"))
data class Employee(
    @Id @GeneratedValue(strategy = GenerationType.IDENTITY)
    val id: Int = 0,

    @Column(name = "company_id")
    val companyId: Int = 0
) {
    @Transient
    lateinit var company: Company
}

// Data fetcher

class CompanyToEmployeesDataFetcher : DataFetcher<CompletableFuture<List<Employee>>>, BeanFactoryAware {
    private lateinit var beanFactory: BeanFactory

    override fun setBeanFactory(beanFactory: BeanFactory) {
        this.beanFactory = beanFactory
    }

    override fun get(environment: DataFetchingEnvironment): CompletableFuture<List<Employee>> {
        val id = environment.getSource<Company>().id
        return environment
            .getDataLoader<Int, List<Employee>>("employeesForCompany")
            .load(id)
    }
}

// Data loader configuration

@Configuration
open class DataLoaderConfiguration(private val employeeRepository: EmployeeRepository) {

    @Bean
    open fun dataLoaderRegistryFactory(): DataLoaderRegistryFactory {
        return object : DataLoaderRegistryFactory {
            override fun generate(): DataLoaderRegistry {
                val registry = DataLoaderRegistry()

                registry.register("employeesForCompany", DataLoader<Int, List<Employee>> { ids ->
                    CompletableFuture.supplyAsync {
                        val employees = employeeRepository.findEmployeesForCompany(ids)
                        ids.map { id -> employees.filter { it.companyId == id } }
                    }
                })

                return registry
            }
        }
    }

}

Querying the company -> employees relationship works. Please let me know if the approach I used above is not the correct one.

The other way around (employee -> company) I can’t even seem to figure out. I figured it would use the resolver registered in the federated type registry but it seems to be asking for a data fetcher. So I’ve tried creating a data fetcher, connecting it to the data loader in which I create a new Company with the ID. Result: Unable to resolve federated type, representation={__typename=Company, id=1} on instance A.

{
  # Company -> Employees works
  company(id: 1) {
    name
    employees {
      id
    }
  }

  # Employee -> Company doesn't work
  employee(id: 100) {
    id
    company {
      id
      name
    }
  }
}

Can you point me in the right direction for implementing relationships across instances? I’ve looked at several examples but without luck. Any help would be greatly appreciated, thanks!

[idrm]

@dariuszkuc After a bit of digging into the library source code, here’s what I came up with (no modification to graphql-kotlin):

import com.expediagroup.graphql.execution.FunctionDataFetcher
import com.expediagroup.graphql.execution.KotlinDataFetcherFactoryProvider
import com.expediagroup.graphql.execution.SimpleKotlinDataFetcherFactoryProvider
import com.fasterxml.jackson.databind.ObjectMapper
import graphql.schema.DataFetcherFactory
import graphql.schema.DataFetchingEnvironment
import org.springframework.beans.factory.annotation.Autowired
import org.springframework.beans.factory.annotation.Qualifier
import org.springframework.context.ApplicationContext
import org.springframework.context.annotation.Bean
import org.springframework.context.annotation.Configuration
import kotlin.reflect.*
import kotlin.reflect.full.findAnnotation
import kotlin.reflect.jvm.javaType

@Configuration
class DataFetcherConfiguration {

    @Bean
    fun dataFetcherFactoryProvider(objectMapper: ObjectMapper, 
                                   appContext: ApplicationContext): KotlinDataFetcherFactoryProvider =
        CustomDataFetcherFactoryProvider(objectMapper, appContext)
}

class CustomFunctionDataFetcher(target: Any?, fn: KFunction<*>, objectMapper: ObjectMapper, 
                                private val appContext: ApplicationContext) :
    FunctionDataFetcher(target, fn, objectMapper) {

    override fun mapParameterToValue(param: KParameter, environment: DataFetchingEnvironment): Any? =
        when {
            param.findAnnotation<Autowired>() != null -> {
                val qualifierAnnotation = param.findAnnotation<Qualifier>()
                if (qualifierAnnotation != null) {
                    appContext.getBean(qualifierAnnotation.value)
                } else {
                    appContext.getBean(param.type.javaType as Class<*>)
                }
            }
            else -> super.mapParameterToValue(param, environment)
        }

}

class CustomDataFetcherFactoryProvider(
    private val objectMapper: ObjectMapper,
    private val appContext: ApplicationContext
) : SimpleKotlinDataFetcherFactoryProvider(objectMapper) {

    override fun functionDataFetcherFactory(target: Any?, kFunction: KFunction<*>): DataFetcherFactory<Any?> {
        return DataFetcherFactory {
            CustomFunctionDataFetcher(
                target = target,
                fn = kFunction,
                objectMapper = objectMapper,
                appContext = appContext)
        }
    }
}

This allows one to inject beans by decorating a type method’s arguments with the @Autowired and, optionally, @Qualifier annotations. The @GraphQLIgnore annotation is required as well, in order to prevent the parameter from being considered for the schema (further tweaks to the beans and underlying processors may allow skipping this annotation, the same way a DataFetchingEnvironment param doesn’t require it). Here’s a usage example:

data class Product(val name: String) {
  fun getOrders(@GraphQLIgnore @Autowired orderRepository: OrderRepository) : List<Order> =
    orderRepository.findByProduct(name)
}

P.S. This being a closed issue, I’m not sure how useful posting the above code snippets here is going to be. Feel free to publish and/or integrate them as you wish.

[dariuszkuc]

@idrm yep that integrations looks good. I’m thinking we can add it as something you have to opt in in 3.x.x branch but we probably we could make it a default one in 4.x.x.