Grafbase Gateway SDK for Extensions

This crate provides building blocks for creating Grafbase Gateway extensions.

Usage

Extensions are still under development. Expect issues if you try them out before we complete development.

Initialize a new project with the Grafbase CLI:

grafbase extension init --type auth/resolver my-extension

This creates a new project with the necessary files and dependencies to get you started. Edit the src/lib.rs file to add your extension logic. The Grafbase Gateway initializes the struct TestProject once during the initial extension call. The Gateway maintains extensions in a connection pool and reuses the struct for multiple requests. Because an extension runs single-threaded, we maintain multiple instances in the gateway memory to handle multiple requests concurrently.

Resolver Example

You can initialize a new resolver extension with the Grafbase CLI:

grafbase extension init --type resolver my-extension

The initialization accepts a list of schema directives from the federated schema (defined in the schema file) and a Configuration object that remains empty for resolver extensions. The [ResolverExtension] derive macro generates the necessary code to initialize a resolver extension and guides you to implement two traits: [Extension] and [Resolver]. The [Extension] trait initializes the extension, and the [Resolver] trait implements the extension logic to resolve a field:

# use grafbase_sdk::{
#    types::{Configuration, Directive, FieldDefinition, FieldInputs, FieldOutput},
#    Error, Extension, Resolver, ResolverExtension, SharedContext,
# };
#[derive(ResolverExtension)]
struct TestProject;

impl Extension for TestProject {
    fn new(schema_directives: Vec<Directive>, config: Configuration) -> Result<Self, Box<dyn std::error::Error>> {
        Ok(Self)
    }
}

impl Resolver for TestProject {
    fn resolve_field(
        &mut self,
        context: SharedContext,
        directive: Directive,
        field_definition: FieldDefinition,
        inputs: FieldInputs,
    ) -> Result<FieldOutput, Error> {
        todo!()
    }
}

The schema_directives in the constructor provides serialized access to all the SCHEMA directives from the subgraph SDL defined in the definitions.graphql file. The directive in the resolve_field provides serialized access to the directive that triggered the resolver extension.

The FieldOutput contains the serialized output of the resolver which transfers back to the gateway. Remember to match the serialized response to the type of the field resolver.

You can find a full example of a REST resolver extension in the Grafbase repository.

Authentication Example

You can initialize a new authentication extension with the Grafbase CLI:

grafbase extension init --type auth my-extension

The initialization needs a list of schema directives from the federated schema (empty for authentication extensions) and a Configuration object that reflects the extension provider configuration in grafbase.toml. The [AuthenticationExtension] derive macro generates code to initialize a resolver extension and guides you to implement two traits: [Extension] and [Authenticator]. The [Extension] trait initializes the extension, and the [Authenticator] trait implements the extension logic to authenticate a request:

# use grafbase_sdk::{
#     types::{Configuration, Directive, ErrorResponse, Token},
#     AuthenticationExtension, Authenticator, Extension, Headers,
# };

#[derive(AuthenticationExtension)]
struct TestProject;

impl Extension for TestProject {
    fn new(schema_directives: Vec<Directive>, config: Configuration) -> Result<Self, Box<dyn std::error::Error>>
        where
            Self: Sized,
        {
            todo!()
        }
    }

impl Authenticator for TestProject {
    fn authenticate(&mut self, headers: Headers) -> Result<Token, ErrorResponse> {
        todo!()
    }
}

The system deserializes the configuration from the grafbase.toml configuration. As an example, here's the configuration data from the JWT extension:

[[authentication.providers]]

[authentication.providers.extension]
extension = "jwt"

[authentication.providers.extension.config]
url = "https://example.com/.well-known/jwks.json"
issuer = "example.com"
audience = "my-project"
poll_interval = 60
header_name = "Authorization"
header_value_prefix = "Bearer "

The config section becomes available through the Configuration struct, and structs implementing serde::Deserialize can deserialize it using with the correspondig deserialization method.

The authenticate method receives request headers as input. A returned token allows the request to continue. You can serialize any data with serde::Serialize and pass it to the token initializer. For certain directives like @requiredScopes, define the scope claim in the token.

Find a complete example of a JWT authentication extension in the Grafbase repository.

Building

You can build your extension with the Grafbase CLI. For this to work, you must have a working rustup installation:

grafbase extension build

This compiles your extension and creates two files:

build/
├── extension.wasm
└── manifest.json

You can use the path to the build directory in your gateway configuration to try out the extension.

Testing

You can enable the test-utils feature for this crate in your extension. The feature provides tooling for testing the extensions against the Grafbase Gateway. Keep in mind to add it as a dev dependency, the test utils do not compile to WebAssembly. Your tests run as native code, and only the extension compiles into a WebAssembly component and tests with the gateway binary.

cargo add --dev grafbase-gateway --features test-utils

Write your tests in the tests/integration_tests.rs file in your extension project.

See the integration tests of the REST extension for an example of how to use the test utils.

You can run the tests with cargo:

cargo test