Set schema depth and complexity

SeaQL · Jan 8, 2025 · e472f3d · e472f3d
1 parent 66053f7
commit e472f3d
Show file tree

Hide file tree

Showing 11 changed files with 91 additions and 125 deletions.
diff --git a/examples/mysql/src/query_root.rs b/examples/mysql/src/query_root.rs
@@ -33,16 +33,10 @@ pub fn schema(
         ]
     );
     builder.register_enumeration::<crate::entities::sea_orm_active_enums::Rating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
diff --git a/examples/mysql/tests/guard_mutation_tests.rs b/examples/mysql/tests/guard_mutation_tests.rs
@@ -54,18 +54,12 @@ pub fn schema(
         ]
     );
     builder.register_enumeration::<sea_orm_active_enums::Rating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/examples/mysql/tests/guard_tests.rs b/examples/mysql/tests/guard_tests.rs
@@ -259,18 +259,12 @@ pub fn schema(
     ]);
     builder
         .register_enumeration::<seaography_mysql_example::entities::sea_orm_active_enums::Rating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/examples/postgres/src/query_root.rs b/examples/postgres/src/query_root.rs
@@ -32,16 +32,10 @@ pub fn schema(
         ]
     );
     builder.register_enumeration::<crate::entities::sea_orm_active_enums::MpaaRating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
diff --git a/examples/postgres/tests/guard_mutation_tests.rs b/examples/postgres/tests/guard_mutation_tests.rs
@@ -53,18 +53,12 @@ pub fn schema(
         ]
     );
     builder.register_enumeration::<sea_orm_active_enums::MpaaRating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/examples/postgres/tests/guard_tests.rs b/examples/postgres/tests/guard_tests.rs
@@ -247,18 +247,12 @@ pub fn schema(
             ),
     ]);
     builder.register_enumeration::<seaography_postgres_example::entities::sea_orm_active_enums::MpaaRating>();
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/examples/sqlite/src/query_root.rs b/examples/sqlite/src/query_root.rs
@@ -32,16 +32,10 @@ pub fn schema(
             store,
         ]
     );
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
diff --git a/examples/sqlite/tests/guard_mutation_tests.rs b/examples/sqlite/tests/guard_mutation_tests.rs
@@ -53,18 +53,12 @@ pub fn schema(
             store,
         ]
     );
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/examples/sqlite/tests/guard_tests.rs b/examples/sqlite/tests/guard_tests.rs
@@ -257,18 +257,12 @@ pub fn schema(
                 seaography_sqlite_example::entities::address::Relation::Store.def(),
             ),
     ]);
-    let schema = builder.schema_builder();
-    let schema = if let Some(depth) = depth {
-        schema.limit_depth(depth)
-    } else {
-        schema
-    };
-    let schema = if let Some(complexity) = complexity {
-        schema.limit_complexity(complexity)
-    } else {
-        schema
-    };
-    schema.data(database).finish()
+    builder
+        .set_depth_limit(depth)
+        .set_complexity_limit(complexity)
+        .schema_builder()
+        .data(database)
+        .finish()
 }
 
 pub async fn get_schema() -> Schema {

diff --git a/generator/src/writer.rs b/generator/src/writer.rs
@@ -107,21 +107,12 @@ pub fn generate_query_root<P: AsRef<Path>>(entities_path: &P) -> TokenStream {
 
             #(#enumerations)*
 
-            let schema = builder.schema_builder();
-
-            let schema = if let Some(depth) = depth {
-                schema.limit_depth(depth)
-            } else {
-                schema
-            };
-
-            let schema = if let Some(complexity) = complexity {
-                schema.limit_complexity(complexity)
-            } else {
-                schema
-            };
-
-            schema.data(database).finish()
+            builder
+                .set_depth_limit(depth)
+                .set_complexity_limit(complexity)
+                .schema_builder()
+                .data(database)
+                .finish()
         }
     }
 }

diff --git a/src/builder.rs b/src/builder.rs
@@ -42,6 +42,12 @@ pub struct Builder {
 
     /// configuration for builder
     pub context: &'static BuilderContext,
+
+    /// Set the maximum depth a query can have
+    pub depth: Option<usize>,
+
+    /// Set the maximum complexity a query can have
+    pub complexity: Option<usize>,
 }
 
 impl Builder {
@@ -66,6 +72,8 @@ impl Builder {
             mutations: Vec::new(),
             connection,
             context,
+            depth: None,
+            complexity: None,
         }
     }
 
@@ -219,6 +227,16 @@ impl Builder {
             .push(filter_types_map_helper.generate_filter_input(&filter_info));
     }
 
+    pub fn set_depth_limit(mut self, depth: Option<usize>) -> Self {
+        self.depth = depth;
+        self
+    }
+
+    pub fn set_complexity_limit(mut self, complexity: Option<usize>) -> Self {
+        self.complexity = complexity;
+        self
+    }
+
     /// used to consume the builder context and generate a ready to be completed GraphQL schema
     pub fn schema_builder(self) -> SchemaBuilder {
         let query = self.query;
@@ -264,7 +282,7 @@ impl Builder {
             .into_iter()
             .fold(schema, |schema, cur| schema.register(cur));
 
-        schema
+        let schema = schema
             .register(
                 OrderByEnumBuilder {
                     context: self.context,
@@ -314,7 +332,18 @@ impl Builder {
                 .to_object(),
             )
             .register(query)
-            .register(mutation)
+            .register(mutation);
+
+        let schema = if let Some(depth) = self.depth {
+            schema.limit_depth(depth)
+        } else {
+            schema
+        };
+        if let Some(complexity) = self.complexity {
+            schema.limit_complexity(complexity)
+        } else {
+            schema
+        }
     }
 }