feat(cognitarium): handle limit & skip at plan build

axone-protocol · Jun 5, 2023 · a96258f · a96258f
1 parent 23536ae
commit a96258f
Showing 1 changed file with 37 additions and 5 deletions.
diff --git a/contracts/okp4-cognitarium/src/querier/plan_builder.rs b/contracts/okp4-cognitarium/src/querier/plan_builder.rs
@@ -12,6 +12,8 @@ pub struct PlanBuilder<'a> {
     storage: &'a dyn Storage,
     prefixes: HashMap<String, String>,
     variables: Vec<String>,
+    limit: Option<usize>,
+    skip: Option<usize>,
 }
 
 impl<'a> PlanBuilder<'a> {
@@ -20,9 +22,21 @@ impl<'a> PlanBuilder<'a> {
             storage,
             prefixes: Self::make_prefixes(prefixes),
             variables: Vec::new(),
+            skip: None,
+            limit: None,
         }
     }
 
+    pub fn with_limit(&mut self, limit: usize) -> &Self {
+        self.limit = Some(limit);
+        self
+    }
+
+    pub fn with_skip(&mut self, skip: usize) -> &Self {
+        self.skip = Some(skip);
+        self
+    }
+
     pub fn build_plan(&mut self, where_clause: WhereClause) -> StdResult<QueryPlan> {
         let bgp: Vec<QueryNode> = where_clause
             .iter()
@@ -33,9 +47,31 @@ impl<'a> PlanBuilder<'a> {
             .collect::<StdResult<Vec<QueryNode>>>()?;
 
         self.build_from_bgp(bgp)
+            .map(|mut node| {
+                if let Some(skip) = self.skip {
+                    node = QueryNode::Skip {
+                        child: Box::new(node),
+                        first: skip,
+                    }
+                }
+                node
+            })
+            .map(|mut node| {
+                if let Some(limit) = self.skip {
+                    node = QueryNode::Limit {
+                        child: Box::new(node),
+                        first: limit,
+                    }
+                }
+                node
+            })
+            .map(|node| QueryPlan {
+                entrypoint: Box::new(node),
+                variables: self.variables.clone(),
+            })
     }
 
-    fn build_from_bgp(&self, bgp: Vec<QueryNode>) -> StdResult<QueryPlan> {
+    fn build_from_bgp(&self, bgp: Vec<QueryNode>) -> StdResult<QueryNode> {
         bgp.into_iter()
             .reduce(|left: QueryNode, right: QueryNode| -> QueryNode {
                 if left
@@ -56,10 +92,6 @@ impl<'a> PlanBuilder<'a> {
             })
             .map(Ok)
             .unwrap_or(Err(StdError::generic_err("Empty basic graph pattern")))
-            .map(|query_node| QueryPlan {
-                entrypoint: Box::new(query_node),
-                variables: self.variables.clone(),
-            })
     }
 
     fn build_triple_pattern(&mut self, pattern: &TriplePattern) -> StdResult<QueryNode> {