How to implement DataLoader to batch load GraphQL class rows in a single DB query?

[paussu]

Hi,

Any example on how to implement a DataLoader for a GraphQL class field?

For example we have an InvoiceHeader class that returns list of InvoiceRow fields and we would want to batch the keys and execute a single DB query to get the rows.

This probably can be implemented with a future or a coroutine?

Asked AI and it gave a DataLoader class like this:

#pragma once

#include <unordered_map>
#include <vector>
#include <future>
#include <mutex>
#include <functional>
#include <queue>
#include <memory>

template<typename Key, typename Value>
class DataLoader {
public:
    using BatchFn = std::function<std::unordered_map<Key, Value>(const std::vector<Key>&)>;

    explicit DataLoader(BatchFn batchFn)
        : _batchFunction(std::move(batchFn))
    {
    }

    std::future<Value> load(const Key& key) {
        std::lock_guard lock(_mutex);

        // Check cache
        auto it = _cache.find(key);
        if (it != _cache.end()) {
            std::promise<Value> promise;
            promise.set_value(it->second);
            return promise.get_future();
        }
        // Check if already pending
        auto pit = _pending.find(key);
        if (pit != _pending.end())
            return pit->second->get_future();

        // New key: add to queue and pending
        auto promisePtr = std::make_shared<std::promise<Value>>();
        _pending[key] = promisePtr;
        _queue.push_back(key);
        return promisePtr->get_future();
    }

    // Call this after all load()s to resolve the batch
    void dispatch()
    {
        std::vector<Key> keys;
        {
            std::lock_guard lock(_mutex);
            if (_queue.empty()) return;
            keys.swap(_queue);
        }
        auto results = _batchFunction(keys);
        std::lock_guard lock(_mutex);
        for (const auto& key : keys)
        {
            auto pit = _pending.find(key);
            if (pit != _pending.end())
            {
                auto rit = results.find(key);
                if (rit != results.end())
                {
                    pit->second->set_value(rit->second);
                    _cache[key] = rit->second;
                }
                else
                {
                    pit->second->set_exception(std::make_exception_ptr(std::runtime_error("Key not found in batch result")));
                }
                _pending.erase(pit);
            }
        }
    }

    // Optional: clear cache for new request
    void clearCache()
    {
        std::lock_guard lock(_mutex);
        _cache.clear();
    }

private:
    BatchFn _batchFunction;
    std::unordered_map<Key, std::shared_ptr<std::promise<Value>>> _pending;
    std::unordered_map<Key, Value> _cache;
    std::vector<Key> _queue;
    std::mutex _mutex;
};

Invoice row

service::AwaitableObject<std::vector<std::shared_ptr<object::InvoiceRow>>> InvoiceHeader::getRows(service::FieldParams&& params) const
   {
      const auto &key = invoiceSeqNumber.value();
      const auto &loader = _invoiceQueries->getInvoiceRowsLoader();

      return loader->load(key);
   }

Does this look correct and where should the dispatch happen? In query endSelectionSet or somewhere else?