Dataloader testing

Author: iamvigneshwars

Cargo.lock

@@ -5,4 +5,4 @@ type: application
 
 version: 0.1.0
 
-appVersion: 0.1.0-rc4
+appVersion: 0.1.0-rc5

charts/processed_data/charts/processed_data/Chart.yaml

@@ -11,6 +11,7 @@ async-graphql = { version = "7.0.2", default-features = false, features = [
     "chrono",
     "graphiql",
     "tracing",
+    "dataloader",
 ] }
 async-graphql-axum = { version = "7.0.2" }
 aws-credential-types = { version = "0.56.0" }

processed_data/Cargo.toml

@@ -2,6 +2,7 @@
 mod entities;
 use crate::S3Bucket;
 use async_graphql::{
+    dataloader::{DataLoader, Loader},
     ComplexObject, Context, EmptyMutation, EmptySubscription, Object, Schema, SchemaBuilder,
 };
 use aws_sdk_s3::presigning::PresigningConfig;
@@ -15,7 +16,9 @@ use models::{
     processing_job_parameter,
 };
 use sea_orm::{ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter};
+use std::collections::HashMap;
 use std::time::Duration;
+use tracing::instrument;
 use url::Url;
 
 use self::entities::AutoProcProgram;
@@ -24,29 +27,62 @@ use self::entities::AutoProcProgram;
 pub type RootSchema = Schema<Query, EmptyMutation, EmptySubscription>;
 
 /// A schema builder for the service
-pub fn root_schema_builder() -> SchemaBuilder<Query, EmptyMutation, EmptySubscription> {
-    Schema::build(Query, EmptyMutation, EmptySubscription).enable_federation()
+pub fn root_schema_builder(
+    database: DatabaseConnection,
+) -> SchemaBuilder<Query, EmptyMutation, EmptySubscription> {
+    Schema::build(Query, EmptyMutation, EmptySubscription)
+        .data(DataLoader::new(
+            DataCollectionLoader::new(database.clone()),
+            tokio::spawn,
+        ))
+        .data(database)
+        .enable_federation()
 }
 
 /// The root query of the service
 #[derive(Debug, Clone, Default)]
 pub struct Query;
 
+pub struct DataCollectionLoader(DatabaseConnection);
+
+impl DataCollectionLoader {
+    fn new(database: DatabaseConnection) -> Self {
+        Self(database)
+    }
+}
+
+impl Loader<u32> for DataCollectionLoader {
+    type Value = DataProcessing;
+    type Error = async_graphql::Error;
+
+    async fn load(&self, keys: &[u32]) -> Result<HashMap<u32, Self::Value>, Self::Error> {
+        let mut results = HashMap::new();
+        let keys_vec: Vec<u32> = keys.iter().cloned().collect();
+        let records = data_collection_file_attachment::Entity::find()
+            .filter(data_collection_file_attachment::Column::DataCollectionId.is_in(keys_vec))
+            .all(&self.0)
+            .await?;
+
+        for record in records {
+            let data_collection_id = record.data_collection_id;
+            let data = DataProcessing::from(record);
+
+            results.insert(data_collection_id, data);
+        }
+
+        Ok(results)
+    }
+}
+
 #[ComplexObject]
 impl DataCollection {
     /// Fetched all the processed data from data collection during a session
     async fn processed_data(
         &self,
         ctx: &Context<'_>,
-    ) -> Result<Vec<DataProcessing>, async_graphql::Error> {
-        let database = ctx.data::<DatabaseConnection>()?;
-        Ok(data_collection_file_attachment::Entity::find()
-            .filter(data_collection_file_attachment::Column::DataCollectionId.eq(self.id))
-            .all(database)
-            .await?
-            .into_iter()
-            .map(DataProcessing::from)
-            .collect())
+    ) -> Result<Option<DataProcessing>, async_graphql::Error> {
+        let loader = ctx.data_unchecked::<DataLoader<DataCollectionLoader>>();
+        Ok(loader.load_one(self.id).await?)
     }
 
     /// Fetched all the processing jobs

processed_data/src/graphql/mod.rs

@@ -126,6 +126,8 @@ struct SchemaArgs {
     /// The path to write the schema to, if not set the schema will be printed to stdout
     #[arg(short, long)]
     path: Option<PathBuf>,
+    #[arg(long, env = "DATABASE_URL")]
+    database_url: Url,
 }
 
 /// Creates a connection pool to access the database
@@ -239,12 +241,13 @@ async fn main() {
         Cli::Serve(args) => {
             setup_telemetry(args.log_level, args.otel_collector_url).unwrap();
             let database = setup_database(args.database_url).await.unwrap();
-            let schema = root_schema_builder().data(database).finish();
+            let schema = root_schema_builder(database).finish();
             let router = setup_router(schema);
             serve(router, args.port).await.unwrap();
         }
         Cli::Schema(args) => {
-            let schema = root_schema_builder().finish();
+            let database = setup_database(args.database_url).await.unwrap();
+            let schema = root_schema_builder(database).finish();
             let schema_string = schema.sdl_with_options(SDLExportOptions::new().federation());
             if let Some(path) = args.path {
                 let mut file = File::create(path).unwrap();