Paginate library pins query

DiamondLightSource · Oct 9, 2023 · fe21ba3 · fe21ba3
1 parent ddf96e7
commit fe21ba3
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 4 deletions.
diff --git a/backend/Cargo.lock b/backend/Cargo.lock
diff --git a/backend/pin_packing/Cargo.toml b/backend/pin_packing/Cargo.toml
@@ -4,6 +4,7 @@ version = "0.1.0"
 edition = "2021"
 
 [dependencies]
+anyhow = { workspace = true }
 async-graphql = { workspace = true }
 axum = { workspace = true }
 clap = { workspace = true }

diff --git a/backend/pin_packing/src/resolvers/pin_library.rs b/backend/pin_packing/src/resolvers/pin_library.rs
@@ -2,9 +2,15 @@ use crate::tables::{
     pin_library::{self, PinStatus},
     pin_mount,
 };
-use async_graphql::{ComplexObject, Context, Object};
+use async_graphql::{
+    connection::{query, Connection, Edge, EmptyFields},
+    ComplexObject, Context, Object,
+};
 use opa_client::subject_authorization;
-use sea_orm::{ActiveValue, DatabaseConnection, EntityTrait, IntoActiveModel, ModelTrait};
+use sea_orm::{
+    ActiveValue, CursorTrait, DatabaseConnection, EntityTrait, IntoActiveModel, ModelTrait,
+    PaginatorTrait,
+};
 
 #[ComplexObject]
 impl pin_library::Model {
@@ -23,10 +29,48 @@ impl PinLibraryQuery {
     async fn library_pins(
         &self,
         ctx: &Context<'_>,
-    ) -> async_graphql::Result<Vec<pin_library::Model>> {
+        after: Option<String>,
+        before: Option<String>,
+        first: Option<i32>,
+        last: Option<i32>,
+    ) -> async_graphql::Result<Connection<usize, pin_library::Model, EmptyFields, EmptyFields>>
+    {
         subject_authorization!("xchemlab.pin_packing.read_pin_library", ctx).await?;
         let database = ctx.data::<DatabaseConnection>()?;
-        Ok(pin_library::Entity::find().all(database).await?)
+        let pin_query = pin_library::Entity::find();
+        let pin_count = pin_query.clone().count(database).await? as usize;
+        let mut pin_cursor = pin_query.cursor_by(pin_library::Column::Barcode);
+        query(
+            after,
+            before,
+            first,
+            last,
+            |after, before, first, last| async move {
+                let (after, before) = match (after, before, first, last) {
+                    (Some(after), None, Some(first), None) => Ok((after, after + first)),
+                    (None, Some(before), None, Some(last)) => Ok((before - last, before)),
+                    (None, None, None, None) => Ok((0, pin_count)),
+                    _ => Err(anyhow::anyhow!(
+                        "Could not determine desired pagination behaviour."
+                    )),
+                }?;
+
+                let pins = pin_cursor
+                    .after(after as u64)
+                    .before(before as u64)
+                    .all(database)
+                    .await?;
+
+                let mut connection = Connection::new(after > 0, before < pin_count);
+                connection.edges.extend(
+                    pins.into_iter()
+                        .enumerate()
+                        .map(|(idx, pin)| Edge::new(after + idx, pin)),
+                );
+                Ok::<_, async_graphql::Error>(connection)
+            },
+        )
+        .await
     }
 }