grafbase_sdk/extension/
hooks.rs

1use crate::{
2    component::AnyExtension,
3    host_io::{event_queue::EventQueue, http::StatusCode},
4    types::{Configuration, Error, ErrorResponse, GatewayHeaders, Headers, HttpRequestParts, OnRequestOutput},
5};
6
7/// The Hooks extension allows you to hook into an incoming request or an outgoing response.
8///
9/// You have mutable access to the headers, and information about the request or response
10/// to decide whether to continue processing or not.
11///
12/// Keep in mind this is not meant for authentication purposes.
13///
14/// # Example
15///
16/// ```rust
17/// use grafbase_sdk::{
18///     HooksExtension,
19///     types::{GatewayHeaders, Headers, Configuration, Error, ErrorResponse},
20///     host_io::event_queue::EventQueue,
21/// };
22///
23/// #[derive(HooksExtension)]
24/// struct MyHooks {
25///     config: Config,
26/// }
27///
28/// #[derive(serde::Deserialize)]
29/// struct Config {
30///     // Define your configuration fields here. They are parsed from
31///     // the grafbase.toml configuration.
32///     something: String,
33/// }
34///
35/// impl HooksExtension for MyHooks {
36///     fn new(config: Configuration) -> Result<Self, Error> {
37///         let config = config.deserialize()?;
38///         Ok(Self { config })
39///     }
40///
41///     #[allow(refining_impl_trait)]
42///     fn on_request(&mut self, url: &str, method: http::Method, headers: &mut GatewayHeaders) -> Result<(), ErrorResponse> {
43///         // Implement your request hook logic here.
44///         Ok(())
45///     }
46///
47///     fn on_response(
48///         &mut self,
49///         status: http::StatusCode,
50///         headers: &mut Headers,
51///         event_queue: EventQueue,
52///     ) -> Result<(), Error> {
53///         // Implement your response hook logic here.
54///         Ok(())
55///     }
56/// }
57/// ```
58#[allow(unused_variables)]
59pub trait HooksExtension: Sized + 'static {
60    /// Creates a new instance of the extension. The [`Configuration`] will contain all the
61    /// configuration defined in the `grafbase.toml` by the extension user in a serialized format.
62    ///
63    /// # Example
64    ///
65    /// The following TOML configuration:
66    /// ```toml
67    /// [extensions.my-hooks.config]
68    /// my_custom_key = "value"
69    /// ```
70    ///
71    /// can be easily deserialized with:
72    ///
73    /// ```rust
74    /// # use grafbase_sdk::types::{Configuration, Error};
75    /// # fn dummy(config: Configuration) -> Result<(), Error> {
76    /// #[derive(serde::Deserialize)]
77    /// struct Config {
78    ///     my_custom_key: String
79    /// }
80    ///
81    /// let config: Config = config.deserialize()?;
82    /// # Ok(())
83    /// # }
84    /// ```
85    fn new(config: Configuration) -> Result<Self, Error>;
86
87    /// Called immediately when a request is received, before entering the GraphQL engine.
88    ///
89    /// This hook can be used to modify the request headers before they are processed by the GraphQL engine, and provides a way to audit the headers, URL, and method before processing the operation.
90    fn on_request(
91        &mut self,
92        url: &str,
93        method: http::Method,
94        headers: &mut GatewayHeaders,
95    ) -> Result<impl IntoOnRequestOutput, ErrorResponse> {
96        Ok(())
97    }
98
99    /// Called right before the response is sent back to the client.
100    ///
101    /// This hook can be used to modify the response headers before the response is sent back to the client.
102    fn on_response(
103        &mut self,
104        status: http::StatusCode,
105        headers: &mut Headers,
106        event_queue: EventQueue,
107    ) -> Result<(), Error> {
108        Ok(())
109    }
110
111    /// Called when a (GraphQL) subgraph request is made, allowing you to modify the request parts before they are sent to the subgraph.
112    fn on_subgraph_request(&mut self, parts: &mut HttpRequestParts) -> Result<(), Error> {
113        Ok(())
114    }
115}
116
117pub trait IntoOnRequestOutput {
118    fn into_on_request_output(self) -> OnRequestOutput;
119}
120
121impl IntoOnRequestOutput for OnRequestOutput {
122    fn into_on_request_output(self) -> OnRequestOutput {
123        self
124    }
125}
126
127impl IntoOnRequestOutput for () {
128    fn into_on_request_output(self) -> OnRequestOutput {
129        OnRequestOutput::default()
130    }
131}
132
133#[doc(hidden)]
134pub fn register<T: HooksExtension>() {
135    pub(super) struct Proxy<T: HooksExtension>(T);
136
137    impl<T: HooksExtension> AnyExtension for Proxy<T> {
138        fn on_request(
139            &mut self,
140            url: &str,
141            method: http::Method,
142            headers: &mut Headers,
143        ) -> Result<OnRequestOutput, ErrorResponse> {
144            HooksExtension::on_request(&mut self.0, url, method, headers)
145                .map(IntoOnRequestOutput::into_on_request_output)
146        }
147
148        fn on_response(
149            &mut self,
150            status: StatusCode,
151            headers: &mut Headers,
152            event_queue: EventQueue,
153        ) -> Result<(), Error> {
154            HooksExtension::on_response(&mut self.0, status, headers, event_queue)
155        }
156
157        fn on_subgraph_request(&mut self, parts: &mut HttpRequestParts) -> Result<(), Error> {
158            HooksExtension::on_subgraph_request(&mut self.0, parts)
159        }
160    }
161
162    crate::component::register_extension(Box::new(|_, config| {
163        <T as HooksExtension>::new(config).map(|extension| Box::new(Proxy(extension)) as Box<dyn AnyExtension>)
164    }))
165}