From a8b7a88e8f32b1eb575640d7638cd63ed5b01a90 Mon Sep 17 00:00:00 2001
From: iamvigneshwars <iamvigneshwars@gmail.com>
Date: Wed, 1 May 2024 13:07:19 +0000
Subject: [PATCH] Add pagination

---
 datasets/src/graphql/mod.rs | 63 ++++++++++++++++++++++++++++++++-----
 1 file changed, 55 insertions(+), 8 deletions(-)

diff --git a/datasets/src/graphql/mod.rs b/datasets/src/graphql/mod.rs
index 9a83afb..3945d8d 100644
--- a/datasets/src/graphql/mod.rs
+++ b/datasets/src/graphql/mod.rs
@@ -2,11 +2,12 @@
 mod entity;
 
 use async_graphql::{
-    ComplexObject, Context, EmptyMutation, EmptySubscription, Object, Schema, SchemaBuilder,
+    types::connection::*, ComplexObject, Context, EmptyMutation, EmptySubscription, Object, Schema,
+    SchemaBuilder,
 };
 use entity::{DataCollection, Session};
 use models::data_collection;
-use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter};
+use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter, QuerySelect};
 
 /// The GraphQL schema exposed by the service
 pub type RootSchema = Schema<Query, EmptyMutation, EmptySubscription>;
@@ -22,25 +23,71 @@ pub struct Query;
 
 #[ComplexObject]
 impl Session {
-    /// Fetches all the data collected during a session
+    /// Fetches all the data collected during a session with pagination
     async fn data_collection(
         &self,
         ctx: &Context<'_>,
-    ) -> async_graphql::Result<Vec<DataCollection>> {
+        after: Option<String>,
+        before: Option<String>,
+        first: Option<i32>,
+        last: Option<i32>,
+    ) -> async_graphql::Result<Connection<usize, DataCollection, EmptyFields, EmptyFields>> {
         let database = ctx.data::<DatabaseConnection>()?;
-        Ok(data_collection::Entity::find()
+        let mut start: usize = after
+            .map(|after| after.parse::<usize>().ok().unwrap_or(0))
+            .unwrap_or(0);
+        let mut end: usize = before
+            .map(|before| before.parse::<usize>().ok().unwrap_or(usize::MAX))
+            .unwrap_or(usize::MAX);
+
+        if let Some(first) = first {
+            end = (start + first as usize).min(end);
+        }
+        if let Some(last) = last {
+            start = if last as usize > end - start {
+                end
+            } else {
+                end - last as usize
+            };
+        }
+
+        let limit: u64 = match end.checked_sub(start) {
+            Some(diff) => diff as u64,
+            None => return Err(async_graphql::Error::new("Pagination limit overflow")),
+        };
+
+        let data_collection_entities = data_collection::Entity::find()
             .filter(data_collection::Column::Sessionid.eq(self.id))
+            .limit(limit)
+            .offset(start as u64)
             .all(database)
-            .await?
+            .await?;
+
+        let data_collections: Vec<DataCollection> = data_collection_entities
             .into_iter()
             .map(DataCollection::from)
-            .collect())
+            .collect::<Vec<_>>();
+
+        let mut connection: Connection<usize, DataCollection> =
+            Connection::new(start > 0, end < usize::MAX);
+        connection
+            .edges
+            .extend(
+                data_collections
+                    .into_iter()
+                    .enumerate()
+                    .map(|(index, data_collection)| {
+                        Edge::with_additional_fields(index + start, data_collection, EmptyFields)
+                    }),
+            );
+
+        Ok(connection)
     }
 }
 
 #[Object]
 impl Query {
-    // /// Reference sessions resolver for the router
+    /// Reference sessions resolver for the router
     #[graphql(entity)]
     async fn router_sessions(&self, id: i32) -> Session {
         Session { id }