Skip to content

Commit

Permalink
fix: supply dataloader
Browse files Browse the repository at this point in the history
  • Loading branch information
kespinola committed Oct 19, 2023
1 parent 5f3f6c9 commit 52ff551
Showing 1 changed file with 19 additions and 40 deletions.
59 changes: 19 additions & 40 deletions api/src/dataloaders/collection.rs
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -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>()
Expand All @@ -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)
Expand Down

0 comments on commit 52ff551

Please sign in to comment.