grafbase_sdk/extension/
resolver.rs

1use crate::{
2    component::AnyExtension,
3    types::{
4        Configuration, Error, IndexedSchema, ResolvedField, Response, SubgraphHeaders, SubgraphSchema,
5        SubscriptionItem, Variables,
6    },
7};
8
9/// A resolver extension is called by the gateway to resolve a specific field.
10///
11/// # Example
12///
13/// You can initialize a new resolver extension with the Grafbase CLI:
14///
15/// ```bash
16/// grafbase extension init --type resolver my-resolver
17/// ```
18///
19/// This will generate the following:
20///
21/// ```rust
22/// use grafbase_sdk::{
23///     ResolverExtension,
24///     types::{Configuration, Error, ResolvedField, Response, SubgraphHeaders, SubgraphSchema, Variables},
25/// };
26///
27/// #[derive(ResolverExtension)]
28/// struct MyResolver {
29///     config: Config
30/// }
31///
32/// // Configuration in the TOML for this extension
33/// #[derive(serde::Deserialize)]
34/// struct Config {
35///     #[serde(default)]
36///     key: Option<String>
37/// }
38///
39/// impl ResolverExtension for MyResolver {
40///     fn new(subgraph_schemas: Vec<SubgraphSchema>, config: Configuration) -> Result<Self, Error> {
41///         let config: Config = config.deserialize()?;
42///         Ok(Self { config })
43///     }
44///
45///     fn resolve(
46///         &mut self,
47///         prepared: &[u8],
48///         headers: SubgraphHeaders,
49///         variables: Variables,
50///     ) -> Result<Response, Error> {
51///         // field which must be resolved. The prepared bytes can be customized to store anything you need in the operation cache.
52///         let field = ResolvedField::try_from(prepared)?;
53///         Ok(Response::null())
54///     }
55/// }
56/// ```
57/// ## Configuration
58///
59/// The configuration provided in the `new` method is the one defined in the `grafbase.toml`
60/// file by the extension user:
61///
62/// ```toml
63/// [extensions.my-resolver.config]
64/// key = "value"
65/// ```
66///
67/// Once your business logic is written down you can compile your extension with:
68///
69/// ```bash
70/// grafbase extension build
71/// ```
72///
73/// It will generate all the necessary files in a `build` directory which you can specify in the
74/// `grafbase.toml` configuration with:
75///
76/// ```toml
77/// [extensions.my-resolver]
78/// path = "<project path>/build"
79/// ```
80///
81/// ## Directives
82///
83/// In addition to the Rust extension, a `definitions.graphql` file will be also generated. It
84/// should define directives for subgraph owners and any necessary input types, scalars or enum
85/// necessary for those. Directives have two purposes for resolvers: define which fields can be
86/// resolved, providing the necessary metadata for it, and provide global metadata with schema
87/// directive.
88///
89/// A HTTP resolver extension could have the following directives for example:
90///
91/// ```graphql
92/// scalar URL
93///
94/// directive @httpEndpoint(name: String!, url: URL!) on SCHEMA
95///
96/// directive @http(endpoint: String!, path: String!) on FIELD_DEFINITION
97/// ```
98///
99/// The `@httpEndpoint` directive would be provided during the [new()](ResolverExtension::new())
100/// method as a schema [crate::types::Directive]. The whole subgraph schema is also provided for
101/// each subgraph where this extension is used. While the latter can be accessed with
102/// [ResolvedField::directive()] in the [resolve()](ResolverExtension::resolve()) method.
103///
104pub trait ResolverExtension: Sized + 'static {
105    /// Creates a new instance of the extension. The [Configuration] will contain all the
106    /// configuration defined in the `grafbase.toml` by the extension user in a serialized format.
107    /// Furthermore the complete subgraph schema is provided whenever this extension is used.
108    ///
109    /// # Configuration example
110    ///
111    /// The following TOML configuration:
112    /// ```toml
113    /// [extensions.my-auth.config]
114    /// my_custom_key = "value"
115    /// ```
116    ///
117    /// can be easily deserialized with:
118    ///
119    /// ```rust
120    /// # use grafbase_sdk::types::{Configuration, Error};
121    /// # fn dummy(config: Configuration) -> Result<(), Error> {
122    /// #[derive(serde::Deserialize)]
123    /// struct Config {
124    ///     my_custom_key: String
125    /// }
126    ///
127    /// let config: Config = config.deserialize()?;
128    /// # Ok(())
129    /// # }
130    /// ```
131    ///
132    /// # Directive example
133    ///
134    /// ```graphql
135    /// extend schema @httpEdnpoint(name: "example", url: "https://example.com")
136    /// ```
137    ///
138    /// can be easily deserialized with:
139    ///
140    /// ```rust
141    /// # use grafbase_sdk::types::{Error, SubgraphSchema};
142    /// # fn dummy(subgraph_schemas: Vec<SubgraphSchema>) -> Result<(), Error> {
143    /// #[derive(serde::Deserialize)]
144    /// struct HttpEndpoint {
145    ///     name: String,
146    ///     url: String
147    /// }
148    /// for schema in subgraph_schemas {
149    ///     for directive in schema.directives() {
150    ///         match directive.name() {
151    ///             "httpEndpoint" => {
152    ///                  let http_endpoint: HttpEndpoint = directive.arguments()?;
153    ///             }
154    ///             _ => unreachable!()
155    ///         }
156    ///     }
157    /// }
158    /// # Ok(())
159    /// # }
160    /// ```
161    fn new(subgraph_schemas: Vec<SubgraphSchema>, config: Configuration) -> Result<Self, Error>;
162
163    /// Prepares the field for resolution. The resulting byte array will be part of the operation
164    /// cache. Backwards compatibility is not a concern as the cache is only re-used for the same
165    /// schema and extension versions.
166    /// By default the [ResolvedField] is cached for a simpler implementation.
167    fn prepare(&mut self, field: ResolvedField<'_>) -> Result<Vec<u8>, Error> {
168        Ok(field.into())
169    }
170
171    /// Resolves the field with the provided prepared bytes, headers and variables. With the
172    /// default [prepare()](ResolverExtension::prepare()) you can retrieve all the relevant
173    /// information with:
174    /// ```rust
175    /// # use grafbase_sdk::types::{SubgraphHeaders, Variables, Response, Error, ResolvedField};
176    /// # fn resolve(prepared: &[u8], headers: SubgraphHeaders, variables: Variables) -> Result<Response, Error> {
177    /// let field = ResolvedField::try_from(prepared)?;
178    /// # Ok(Response::null())
179    /// # }
180    /// ```
181    ///
182    /// If you're not doing any data transformation it's best to forward the JSON or CBOR bytes,
183    /// with [Response::json] and [Response::cbor] respectively, directly to the gateway. The
184    /// gateway will always validate the subgraph data and deal with error propagation. Otherwise
185    /// use [Response::data] to use the fastest supported serialization format.
186    fn resolve(&mut self, prepared: &[u8], headers: SubgraphHeaders, variables: Variables) -> Result<Response, Error>;
187
188    /// Resolves a subscription for the given prepared bytes, headers, and variables.
189    /// Subscriptions must implement the [Subscription] trait. It's also possible to provide a
190    /// de-duplication key. If provided the gateway will first check if there is an existing
191    /// subscription with the same key and if there is, re-use it for the new client. This greatly
192    /// limits the impact on the upstream service. So you have two choices for the result type:
193    /// - return a `Ok(subscription)` directly, without any de-duplication
194    /// - return a `Ok((key, callback))` with an optional de-duplication key and a callback
195    ///   function. The latter is only called if no existing subscription exists for the given key.
196    #[allow(unused_variables)]
197    fn resolve_subscription<'s>(
198        &'s mut self,
199        prepared: &'s [u8],
200        headers: SubgraphHeaders,
201        variables: Variables,
202    ) -> Result<impl IntoSubscription<'s>, Error> {
203        unimplemented!("Subscription resolution not implemented for this resolver extension");
204
205        // So that Rust doesn't complain about the unknown type
206        #[allow(unused)]
207        Ok(PhantomSubscription)
208    }
209}
210
211/// A trait for consuming field outputs from streams.
212///
213/// This trait provides an abstraction over different implementations
214/// of subscriptions to field output streams. Implementors should handle
215/// the details of their specific transport mechanism while providing a
216/// consistent interface for consumers.
217pub trait Subscription {
218    /// Retrieves the next field output from the subscription.
219    ///
220    /// Returns:
221    /// - `Ok(Some(Data))` if a field output was available
222    /// - `Ok(None)` if the subscription has ended normally
223    /// - `Err(Error)` if an error occurred while retrieving the next field output
224    fn next(&mut self) -> Result<Option<SubscriptionItem>, Error>;
225}
226
227pub type SubscriptionCallback<'s> = Box<dyn FnOnce() -> Result<Box<dyn Subscription + 's>, Error> + 's>;
228
229pub trait IntoSubscription<'s> {
230    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>);
231}
232
233impl<'s, S> IntoSubscription<'s> for S
234where
235    S: Subscription + 's,
236{
237    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
238        (None, Box::new(move || Ok(Box::new(self))))
239    }
240}
241
242impl<'s, Callback, S> IntoSubscription<'s> for (Vec<u8>, Callback)
243where
244    Callback: FnOnce() -> Result<S, Error> + 's,
245    S: Subscription + 's,
246{
247    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
248        (
249            Some(self.0),
250            Box::new(move || {
251                let s = (self.1)()?;
252                Ok(Box::new(s) as Box<dyn Subscription + 's>)
253            }),
254        )
255    }
256}
257
258impl<'s, Callback, S> IntoSubscription<'s> for (Option<Vec<u8>>, Callback)
259where
260    Callback: FnOnce() -> Result<S, Error> + 's,
261    S: Subscription + 's,
262{
263    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
264        (
265            self.0,
266            Box::new(move || {
267                let s = (self.1)()?;
268                Ok(Box::new(s) as Box<dyn Subscription + 's>)
269            }),
270        )
271    }
272}
273
274impl<'s, Callback, S> IntoSubscription<'s> for (String, Callback)
275where
276    Callback: FnOnce() -> Result<S, Error> + 's,
277    S: Subscription + 's,
278{
279    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
280        (
281            Some(self.0.into()),
282            Box::new(move || {
283                let s = (self.1)()?;
284                Ok(Box::new(s) as Box<dyn Subscription + 's>)
285            }),
286        )
287    }
288}
289
290impl<'s, Callback, S> IntoSubscription<'s> for (Option<String>, Callback)
291where
292    Callback: FnOnce() -> Result<S, Error> + 's,
293    S: Subscription + 's,
294{
295    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
296        (
297            self.0.map(Into::into),
298            Box::new(move || {
299                let s = (self.1)()?;
300                Ok(Box::new(s) as Box<dyn Subscription + 's>)
301            }),
302        )
303    }
304}
305
306impl<'s, Callback, S> IntoSubscription<'s> for (Callback, Vec<u8>)
307where
308    Callback: FnOnce() -> Result<S, Error> + 's,
309    S: Subscription + 's,
310{
311    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
312        (self.1, self.0).into_deduplication_key_and_subscription_callback()
313    }
314}
315
316impl<'s, Callback, S> IntoSubscription<'s> for (Callback, Option<Vec<u8>>)
317where
318    Callback: FnOnce() -> Result<S, Error> + 's,
319    S: Subscription + 's,
320{
321    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
322        (self.1, self.0).into_deduplication_key_and_subscription_callback()
323    }
324}
325
326impl<'s, Callback, S> IntoSubscription<'s> for (Callback, String)
327where
328    Callback: FnOnce() -> Result<S, Error> + 's,
329    S: Subscription + 's,
330{
331    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
332        (self.1, self.0).into_deduplication_key_and_subscription_callback()
333    }
334}
335
336impl<'s, Callback, S> IntoSubscription<'s> for (Callback, Option<String>)
337where
338    Callback: FnOnce() -> Result<S, Error> + 's,
339    S: Subscription + 's,
340{
341    fn into_deduplication_key_and_subscription_callback(self) -> (Option<Vec<u8>>, SubscriptionCallback<'s>) {
342        (self.1, self.0).into_deduplication_key_and_subscription_callback()
343    }
344}
345
346#[doc(hidden)]
347pub fn register<T: ResolverExtension>() {
348    pub(super) struct Proxy<T: ResolverExtension>(T);
349
350    impl<T: ResolverExtension> AnyExtension for Proxy<T> {
351        fn prepare(&mut self, field: ResolvedField<'_>) -> Result<Vec<u8>, Error> {
352            self.0.prepare(field)
353        }
354
355        fn resolve(&mut self, prepared: &[u8], headers: SubgraphHeaders, variables: Variables) -> Response {
356            self.0.resolve(prepared, headers, variables).into()
357        }
358
359        fn resolve_subscription<'a>(
360            &'a mut self,
361            prepared: &'a [u8],
362            headers: SubgraphHeaders,
363            variables: Variables,
364        ) -> Result<(Option<Vec<u8>>, SubscriptionCallback<'a>), Error> {
365            let (key, callback) = self
366                .0
367                .resolve_subscription(prepared, headers, variables)?
368                .into_deduplication_key_and_subscription_callback();
369            Ok((key, callback))
370        }
371    }
372
373    crate::component::register_extension(Box::new(|subgraph_schemas, config| {
374        let schemas = subgraph_schemas
375            .into_iter()
376            .map(IndexedSchema::from)
377            .collect::<Vec<_>>();
378        <T as ResolverExtension>::new(schemas.into_iter().map(SubgraphSchema).collect(), config)
379            .map(|extension| Box::new(Proxy(extension)) as Box<dyn AnyExtension>)
380    }))
381}
382
383struct PhantomSubscription;
384
385impl Subscription for PhantomSubscription {
386    fn next(&mut self) -> Result<Option<SubscriptionItem>, Error> {
387        Ok(None) // No-op implementation for the phantom subscription
388    }
389}