Enhance Your GraphQL Federation API with NATS Subscriptions using Grafbase

Earlier this month, we launched Grafbase Extensions, which provides a fast, reliable, and backwards-compatible way to customize and enhance federated GraphQL APIs. Think of them as modules you can install and develop outside of the Grafbase Gateway. We maintain a stable ABI, allowing our customers to update the Grafbase Gateway freely without worrying about breaking existing extensions.

Today, we will demonstrate the power of our NATS extension. The NATS extension enables you to publish and subscribe to NATS subjects directly from your GraphQL API, and lets you use the key-value store to store and retrieve data.

A Grafbase extension consists of code written in Rust using the Grafbase Extension SDK. The SDK offers necessary interfaces to implement the extension's logic. After writing the extension, you compile it into a WebAssembly component with the Grafbase CLI.

Grafbase offers a selection of extensions in the Extensions Marketplace. We automatically build and publish these from our extensions repository. If you are running the Grafbase Enterprise Platform, you can also build and publish extensions from your own repository.

A resolver extension typically provides at least one field directive. The NATS extension implements the resolver interface, which enables it to resolve a field or subscription. For example, a subscription field with the @natsSubscription directive triggers a subscription to a NATS subject.

For most cases, just install and configure the extension.

The NATS extension acts as a virtual subgraph that integrates into the federated schema. First, extend your schema with the directives from the extension:

extend schema
  @link(
    url: "https://grafbase.com/extensions/nats/0.4.1"
    import: [
      "@natsPublish",
      "@natsSubscription",
      "@natsRequest",
      "@natsKeyValue"
    ]
  )

Then use these directives to define a virtual subgraph schema:

type Query {
  # https://docs.nats.io/nats-concepts/core-nats/reqreply
  requestReply(input: RequestReplyInput!): RequestReplyResult! @natsRequest(
    subject: "help.please",
    timeoutMs: 500,
  )
}

type Mutation {
  # https://docs.nats.io/nats-concepts/core-nats/pubsub/pubsub_walkthrough
  publishUserEvent(id: Int!, input: UserEventInput!): Boolean! @natsPublish(
    subject: "publish.user.{{ args.id }}.events"
  )
}

# https://docs.nats.io/nats-concepts/core-nats/pubsub/pubsub_walkthrough
type Subscription {
  # Subscription with a filter
  highPriorityBankEvents(limit: Int!): BankEvent! @natsSubscription(
    subject: "subscription.bank",
    selection: "select(.money > {{ args.limit }}) | { id, account, money }",
  )
}

You can find other supported directives in the extension documentation.

A WASM extension can serve only one request at a time. We launch instances of an extension internally in a connection pool, releasing them to other requests when no longer in use. This creates problems with subscriptions where requests live longer.

A subscription extension can implement an optional method to generate a subscription key. When implemented, this enables the engine to deduplicate streams. So, if multiple clients request the exact same subscription, we can serve all of them from a single instance.

This saves considerable CPU and memory, and our NATS extension demonstrates how deduplication works.

Our NATS extension supports only simple values or JSON objects in payloads. Companies often enforce static schemas for their event bus, whether Protobuf, Avro, or other static formats. For example, Protobuf requires you to compile your code with type definitions. This has traditionally presented a difficult problem to solve generically.

We licensed the NATS extension under the Apache 2.0 License, and made the code available on GitHub. The well-documented code makes it easy to modify for your needs.

If you face any issues or have questions, feel free to reach out to us and we can help you customize the extension to fit your needs.

Now that we have the infrastructure to write extensions for subscriptions, we can easily add extensions for other streaming platforms like Kafka, RabbitMQ, and others.

Extensions powerfully enhance Grafbase Gateway, and we've only begun to explore what they make possible. We look forward to seeing what the community builds with them.