Paginate library pins query

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

Hide file tree

Showing 3 changed files with 47 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,14 @@ 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,
+};
 
 #[ComplexObject]
 impl pin_library::Model {
@@ -23,10 +28,46 @@ 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<String, 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();
+        query(
+            after,
+            before,
+            first,
+            last,
+            |after, before, first, last| async move {
+                let mut pin_cursor = pin_query.cursor_by(pin_library::Column::Barcode);
+                if let Some(after) = after {
+                    pin_cursor.after(after);
+                }
+                if let Some(before) = before {
+                    pin_cursor.before(before);
+                }
+                if let Some(first) = first {
+                    pin_cursor.first(first as u64);
+                }
+                if let Some(last) = last {
+                    pin_cursor.last(last as u64);
+                }
+
+                let pins = pin_cursor.all(database).await?;
+
+                let mut connection = Connection::new(true, true);
+                connection.edges.extend(
+                    pins.into_iter()
+                        .map(|pin| Edge::new(pin.barcode.clone(), pin)),
+                );
+                Ok::<_, async_graphql::Error>(connection)
+            },
+        )
+        .await
     }
 }