From 8f2e1f0796731aa3c5c90518c3962ead8c200bbb Mon Sep 17 00:00:00 2001
From: iamvigneshwars <iamvigneshwars@gmail.com>
Date: Thu, 11 Apr 2024 11:30:40 +0000
Subject: [PATCH] Move dataloader to router handler for tracing

---
 .../charts/processed_data/Chart.yaml          |  2 +-
 processed_data/src/graphql/mod.rs             | 21 ++++++++++++-------
 processed_data/src/main.rs                    | 11 +++++-----
 processed_data/src/route_handlers.rs          | 21 +++++++------------
 4 files changed, 28 insertions(+), 27 deletions(-)

diff --git a/charts/processed_data/charts/processed_data/Chart.yaml b/charts/processed_data/charts/processed_data/Chart.yaml
index e0db19f..f37dccd 100644
--- a/charts/processed_data/charts/processed_data/Chart.yaml
+++ b/charts/processed_data/charts/processed_data/Chart.yaml
@@ -5,4 +5,4 @@ type: application
 
 version: 0.1.0
 
-appVersion: 0.1.0-rc8
+appVersion: 0.1.0-rc9
diff --git a/processed_data/src/graphql/mod.rs b/processed_data/src/graphql/mod.rs
index 52841da..c28ea93 100644
--- a/processed_data/src/graphql/mod.rs
+++ b/processed_data/src/graphql/mod.rs
@@ -26,12 +26,14 @@ use self::entities::AutoProcProgram;
 /// The GraphQL schema exposed by the service
 pub type RootSchema = Schema<Query, EmptyMutation, EmptySubscription>;
 
-/// A schema builder for the service
-pub fn root_schema_builder(
-    database: DatabaseConnection,
-) -> SchemaBuilder<Query, EmptyMutation, EmptySubscription> {
-    Schema::build(Query, EmptyMutation, EmptySubscription)
-        .data(DataLoader::new(
+pub trait AddDataLoadersExt {
+    fn add_data_loaders(self, database: DatabaseConnection) -> Self;
+}
+
+impl AddDataLoadersExt for async_graphql::Request {
+    #[instrument(name = "add_data_loaders", skip(self))]
+    fn add_data_loaders(self, database: DatabaseConnection) -> Self {
+        self.data(DataLoader::new(
             ProcessedDataLoader::new(database.clone()),
             tokio::spawn,
         ))
@@ -72,7 +74,12 @@ pub fn root_schema_builder(
             tokio::spawn,
         ))
         .data(database)
-        .enable_federation()
+    }
+}
+
+/// A schema builder for the service
+pub fn root_schema_builder() -> SchemaBuilder<Query, EmptyMutation, EmptySubscription> {
+    Schema::build(Query, EmptyMutation, EmptySubscription).enable_federation()
 }
 
 /// The root query of the service
diff --git a/processed_data/src/main.rs b/processed_data/src/main.rs
index 6d9a4fa..aad8308 100644
--- a/processed_data/src/main.rs
+++ b/processed_data/src/main.rs
@@ -143,7 +143,7 @@ async fn setup_database(database_url: Url) -> Result<DatabaseConnection, Transac
 }
 
 /// Creates an [`axum::Router`] serving GraphiQL, synchronous GraphQL and GraphQL subscriptions
-fn setup_router(schema: RootSchema) -> Router {
+fn setup_router(schema: RootSchema, database: DatabaseConnection) -> Router {
     #[allow(clippy::missing_docs_in_private_items)]
     const GRAPHQL_ENDPOINT: &str = "/";
 
@@ -153,7 +153,7 @@ fn setup_router(schema: RootSchema) -> Router {
             get(Html(
                 GraphiQLSource::build().endpoint(GRAPHQL_ENDPOINT).finish(),
             ))
-            .post(GraphQLHandler::new(schema)),
+            .post(GraphQLHandler::new(schema, database)),
         )
         .layer(OtelInResponseLayer)
         .layer(OtelAxumLayer::default())
@@ -241,13 +241,12 @@ 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(database).finish();
-            let router = setup_router(schema);
+            let schema = root_schema_builder().finish();
+            let router = setup_router(schema, database);
             serve(router, args.port).await.unwrap();
         }
         Cli::Schema(args) => {
-            let database = setup_database(args.database_url).await.unwrap();
-            let schema = root_schema_builder(database).finish();
+            let schema = root_schema_builder().finish();
             let schema_string = schema.sdl_with_options(SDLExportOptions::new().federation());
             if let Some(path) = args.path {
                 let mut file = File::create(path).unwrap();
diff --git a/processed_data/src/route_handlers.rs b/processed_data/src/route_handlers.rs
index 135be2d..bb5ff74 100644
--- a/processed_data/src/route_handlers.rs
+++ b/processed_data/src/route_handlers.rs
@@ -7,23 +7,23 @@ use axum::{
     response::{IntoResponse, Response},
     RequestExt,
 };
-use axum_extra::{
-    headers::{authorization::Bearer, Authorization},
-    TypedHeader,
-};
+use sea_orm::DatabaseConnection;
 use std::{future::Future, pin::Pin};
 
+use crate::graphql::AddDataLoadersExt;
+
 /// An [`Handler`] which executes an [`Executor`] including the [`Authorization<Bearer>`] in the [`async_graphql::Context`]
 #[derive(Debug, Clone)]
 pub struct GraphQLHandler<E: Executor> {
     /// The GraphQL executor used to process the request
     executor: E,
+    database: DatabaseConnection,
 }
 
 impl<E: Executor> GraphQLHandler<E> {
     /// Constructs an instance of the handler with the provided schema.
-    pub fn new(executor: E) -> Self {
-        Self { executor }
+    pub fn new(executor: E, database: DatabaseConnection) -> Self {
+        Self { executor, database }
     }
 }
 
@@ -33,18 +33,13 @@ where
 {
     type Future = Pin<Box<dyn Future<Output = Response> + Send + 'static>>;
 
-    fn call(self, mut req: Request, _state: S) -> Self::Future {
+    fn call(self, req: Request, _state: S) -> Self::Future {
         Box::pin(async move {
-            let token = req
-                .extract_parts::<TypedHeader<Authorization<Bearer>>>()
-                .await
-                .ok()
-                .map(|token| token.0);
             let request = req.extract::<GraphQLRequest, _>().await;
             match request {
                 Ok(request) => GraphQLResponse::from(
                     self.executor
-                        .execute(request.into_inner().data(token))
+                        .execute(request.into_inner().add_data_loaders(self.database))
                         .await,
                 )
                 .into_response(),