Fix the multi get objects

MystenLabs · Dec 19, 2024 · 0dfcc3a · 0dfcc3a
1 parent aec9627
commit 0dfcc3a
Show file tree

Hide file tree

Showing 6 changed files with 66 additions and 6 deletions.
diff --git a/crates/sui-graphql-rpc/schema.graphql b/crates/sui-graphql-rpc/schema.graphql
@@ -3329,7 +3329,7 @@ type Query {
 	"""
 	Fetch a list of objects by their IDs and versions.
 	"""
-	multiGetObjects(objectKeys: [ObjectKey!]!): [Object!]!
+	multiGetObjects(keys: [ObjectKey!]!): [Object!]!
 	"""
 	The coin objects that exist in the network.
 	

diff --git a/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs b/crates/sui-graphql-rpc/src/extensions/query_limits_checker.rs
@@ -381,7 +381,7 @@ impl<'a> LimitsTraversal<'a> {
     ///
     /// This check must be done after the input limit check, because it relies on the query depth
     /// being bounded to protect it from recursing too deeply.
-    fn check_output_limits(&mut self, op: &'a Positioned<OperationDefinition>) -> ServerResult<()> {
+    fn check_output_limits(&mut self, op: &Positioned<OperationDefinition>) -> ServerResult<()> {
         for selection in &op.node.selection_set.node.items {
             self.traverse_selection_for_output(selection, 1, None)?;
         }
@@ -397,7 +397,7 @@ impl<'a> LimitsTraversal<'a> {
     /// size of the connection's page.
     fn traverse_selection_for_output(
         &mut self,
-        selection: &'a Positioned<Selection>,
+        selection: &Positioned<Selection>,
         multiplicity: u32,
         page_size: Option<u32>,
     ) -> ServerResult<()> {

diff --git a/crates/sui-graphql-rpc/src/server/builder.rs b/crates/sui-graphql-rpc/src/server/builder.rs
@@ -2031,4 +2031,64 @@ pub mod tests {
              bytes or fewer."
         );
     }
+
+    #[tokio::test]
+    async fn test_multi_get_objects_query_limits() {
+        let cluster = prep_executor_cluster().await;
+        let db_url = cluster.graphql_connection_config.db_url.clone();
+        assert_eq!(
+            execute_for_error(
+                &db_url,
+                Limits {
+                    max_output_nodes: 5,
+                    ..Default::default()
+                }, // the query will have 6 output nodes: 2 keys * 3 fields, thus exceeding the
+                   // limit
+                r#"
+                    query {
+                          multiGetObjects(
+                            keys: [
+                              {objectId: "0x01dcb4674affb04e68d8088895e951f4ea335ef1695e9e50c166618f6789d808", version: 2},
+                              {objectId: "0x23e340e97fb41249278c85b1f067dc88576f750670c6dc56572e90971f857c8c", version: 2},
+                            ]
+                          ) {
+                                address
+                                status
+                                version
+                            }
+                    }
+                "#
+                .into(),
+            )
+            .await,
+            "Estimated output nodes exceeds 5"
+        );
+        assert_eq!(
+            execute_for_error(
+                &db_url,
+                Limits {
+                    max_output_nodes: 4,
+                    ..Default::default()
+                }, // the query will have 5 output nodes, thus exceeding the limit
+                r#"
+                    query {
+                          multiGetObjects(
+                            keys: [
+                              {objectId: "0x01dcb4674affb04e68d8088895e951f4ea335ef1695e9e50c166618f6789d808", version: 2},
+                              {objectId: "0x23e340e97fb41249278c85b1f067dc88576f750670c6dc56572e90971f857c8c", version: 2},
+                              {objectId: "0x23e340e97fb41249278c85b1f067dc88576f750670c6dc56572e90971f857c8c", version: 2}, 
+                              {objectId: "0x33032e0706337632361f2607b79df8c9d1079e8069259b27b1fa5c0394e79893", version: 2},
+                              {objectId: "0x388295e3ecad53986ebf9a7a1e5854b7df94c3f1f0bba934c5396a2a9eb4550b", version: 2},
+                            ]
+                          ) {
+                                address
+                            }
+                    }
+                "#
+                .into(),
+            )
+            .await,
+            "Estimated output nodes exceeds 4"
+        );
+    }
 }
diff --git a/crates/sui-graphql-rpc/staging.graphql b/crates/sui-graphql-rpc/staging.graphql
@@ -3329,7 +3329,7 @@ type Query {
 	"""
 	Fetch a list of objects by their IDs and versions.
 	"""
-	multiGetObjects(objectKeys: [ObjectKey!]!): [Object!]!
+	multiGetObjects(keys: [ObjectKey!]!): [Object!]!
 	"""
 	The coin objects that exist in the network.
 	

diff --git a/crates/sui-graphql-rpc/tests/snapshots/snapshot_tests__schema.graphql.snap b/crates/sui-graphql-rpc/tests/snapshots/snapshot_tests__schema.graphql.snap
@@ -3333,7 +3333,7 @@ type Query {
 	"""
 	Fetch a list of objects by their IDs and versions.
 	"""
-	multiGetObjects(objectKeys: [ObjectKey!]!): [Object!]!
+	multiGetObjects(keys: [ObjectKey!]!): [Object!]!
 	"""
 	The coin objects that exist in the network.
 	

diff --git a/crates/sui-graphql-rpc/tests/snapshots/snapshot_tests__staging.graphql.snap b/crates/sui-graphql-rpc/tests/snapshots/snapshot_tests__staging.graphql.snap
@@ -3333,7 +3333,7 @@ type Query {
 	"""
 	Fetch a list of objects by their IDs and versions.
 	"""
-	multiGetObjects(objectKeys: [ObjectKey!]!): [Object!]!
+	multiGetObjects(keys: [ObjectKey!]!): [Object!]!
 	"""
 	The coin objects that exist in the network.