diff --git a/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs b/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs index 6159cd8405c8b3..715241653542ba 100644 --- a/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs +++ b/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs @@ -23,9 +23,16 @@ use tokio::sync::Mutex; use tracing::info; use uuid::Uuid; +pub(crate) const CONNECTION_FIELDS: [&str; 2] = ["edges", "nodes"]; + /// Extension factory for adding checks that the query is within configurable limits. pub(crate) struct QueryLimitsChecker; +#[derive(Debug, Default)] +struct QueryLimitsCheckerExt { + validation_result: Mutex>, +} + /// Only display usage information if this header was in the request. pub(crate) struct ShowUsage; @@ -39,13 +46,13 @@ struct ValidationRes { query_payload: u32, } -#[derive(Debug, Default)] -struct QueryLimitsCheckerExt { - validation_result: Mutex>, +#[derive(Debug)] +struct ComponentCost { + pub input_nodes: u32, + pub output_nodes: u32, + pub depth: u32, } -pub(crate) const CONNECTION_FIELDS: [&str; 2] = ["edges", "nodes"]; - impl ShowUsage { pub(crate) fn name() -> &'static HeaderName { &LIMITS_HEADER @@ -60,25 +67,6 @@ impl ExtensionFactory for QueryLimitsChecker { } } -#[derive(Debug)] -struct ComponentCost { - pub input_nodes: u32, - pub output_nodes: u32, - pub depth: u32, -} - -impl std::ops::Add for ComponentCost { - type Output = Self; - - fn add(self, rhs: Self) -> Self::Output { - Self { - input_nodes: self.input_nodes + rhs.input_nodes, - output_nodes: self.output_nodes + rhs.output_nodes, - depth: self.depth + rhs.depth, - } - } -} - #[async_trait::async_trait] impl Extension for QueryLimitsCheckerExt { async fn request(&self, ctx: &ExtensionContext<'_>, next: NextRequest<'_>) -> Response { @@ -210,6 +198,18 @@ impl Extension for QueryLimitsCheckerExt { } } +impl std::ops::Add for ComponentCost { + type Output = Self; + + fn add(self, rhs: Self) -> Self::Output { + Self { + input_nodes: self.input_nodes + rhs.input_nodes, + output_nodes: self.output_nodes + rhs.output_nodes, + depth: self.depth + rhs.depth, + } + } +} + /// Parse the selected fields in one operation and check if it conforms to configured limits. fn analyze_selection_set( limits: &Limits,