From 52ff5516da374afa9af02e42c4a00fc04ca3baf6 Mon Sep 17 00:00:00 2001
From: Kyle Espinola <kyle.s.espinola@gmail.com>
Date: Thu, 19 Oct 2023 14:03:38 +0200
Subject: [PATCH] fix: supply dataloader

---
 api/src/dataloaders/collection.rs | 59 ++++++++++---------------------
 1 file changed, 19 insertions(+), 40 deletions(-)

diff --git a/api/src/dataloaders/collection.rs b/api/src/dataloaders/collection.rs
index 66f9e38..3eb5c58 100644
--- a/api/src/dataloaders/collection.rs
+++ b/api/src/dataloaders/collection.rs
@@ -1,7 +1,6 @@
 use std::collections::HashMap;
 
 use async_graphql::{dataloader::Loader as DataLoader, FieldError, Result};
-use hub_core::tracing::info;
 use poem::async_trait;
 use redis::{AsyncCommands, Client as Redis};
 use sea_orm::{prelude::*, FromQueryResult, QueryFilter, QuerySelect};
@@ -151,69 +150,53 @@ impl DataLoader<Uuid> for SupplyLoader {
     async fn load(&self, keys: &[Uuid]) -> Result<HashMap<Uuid, Self::Value>, Self::Error> {
         let mut results: HashMap<Uuid, Self::Value> = HashMap::new();
         let mut missing_keys: Vec<Uuid> = Vec::new();
+        let mut compute_keys: Vec<Uuid> = Vec::new();
 
         let mut redis_connection = self.redis.get_async_connection().await?;
+        let conn = self.db.get();
 
         for key in keys {
             let redis_key = format!("collection:{key}:supply");
-            match redis_connection.get::<_, Option<i64>>(&redis_key).await {
+
+            match redis_connection.get::<_, i64>(&redis_key).await {
                 Ok(value) => {
-                    info!("Got value from Redis for key: {}", key);
-                    results.insert(*key, value);
+                    results.insert(*key, Some(value));
                 },
                 Err(_) => {
-                    info!("Failed to get value from Redis for key: {}", key);
                     missing_keys.push(*key);
                 },
             }
         }
 
-        if missing_keys.is_empty() {
-            info!("No missing keys, returning results");
-            return Ok(results);
-        }
-
-        let conn = self.db.get();
-        let mut computed_supplies: Vec<Uuid> = Vec::new();
-        info!("Missing keys: {:?}", missing_keys);
-
         let collection_with_drops = collections::Entity::find()
             .filter(collections::Column::Id.is_in(missing_keys.iter().map(ToOwned::to_owned)))
             .inner_join(drops::Entity)
             .select_also(drops::Entity)
             .all(conn)
             .await?;
-        info!("Collection with drops: {:?}", collection_with_drops);
 
         for (collection, drop) in collection_with_drops {
-            if let Some(drop) = drop {
-                if drop.drop_type == DropType::Open {
-                    info!("Open drop for collection: {}", collection.id);
-                    computed_supplies.push(collection.id);
-                } else {
-                    let redis_key = format!("collection:{}:supply", collection.id);
-
-                    redis_connection.set(&redis_key, collection.supply).await?;
-                    info!(
-                        "supply: {:?} for collection {}",
-                        collection.supply, collection.id
-                    );
-
+            match drop {
+                Some(drop) if drop.drop_type == DropType::Edition => {
                     results.insert(collection.id, collection.supply);
-                }
-            } else {
-                info!("No drop for collection: {}", collection.id);
-                computed_supplies.push(collection.id);
+                },
+                Some(_) | None => {
+                    compute_keys.push(collection.id);
+                },
             }
         }
 
+        if compute_keys.is_empty() {
+            return Ok(results);
+        }
+
         let count_results = collection_mints::Entity::find()
             .select_only()
             .column_as(collection_mints::Column::Id.count(), "count")
             .column_as(collection_mints::Column::CollectionId, "id")
             .filter(
                 collection_mints::Column::CollectionId
-                    .is_in(computed_supplies.iter().map(ToOwned::to_owned)),
+                    .is_in(compute_keys.iter().map(ToOwned::to_owned)),
             )
             .group_by(collection_mints::Column::CollectionId)
             .into_model::<CollectionSupplyCount>()
@@ -223,17 +206,13 @@ impl DataLoader<Uuid> for SupplyLoader {
             .map(|result| (result.id, result.count))
             .collect::<HashMap<_, _>>();
 
-        for key in computed_supplies {
-            let count = count_results.get(&key).copied();
+        for key in compute_keys {
+            let count = count_results.get(&key).copied().unwrap_or_default();
             let redis_key = format!("collection:{key}:supply");
 
             redis_connection.set(&redis_key, count).await?;
-            info!(
-                "Set Redis key for computed supply: {} count: {:?}",
-                key, count
-            );
 
-            results.insert(key, count);
+            results.insert(key, Some(count));
         }
 
         Ok(results)