Skip to content

Commit

Permalink
wip
Browse files Browse the repository at this point in the history
  • Loading branch information
@andrewjstone
andrewjstone committed Nov 15, 2023
1 parent fdf4f79 commit 1721f7b
Showing 1 changed file with 47 additions and 31 deletions.
78 changes: 47 additions & 31 deletions nexus/db-queries/src/db/datastore/inventory.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -818,7 +818,11 @@ impl DataStore {

// TODO: Fix me to return a proper error
Ok(self
.inventory_collection_read_all_or_nothing(collection_id, limit)
.inventory_collection_read_all_or_nothing(
opctx,
collection_id,
limit,
)
.await
.unwrap())
}
Expand All @@ -827,11 +831,13 @@ impl DataStore {
/// records and returning nothing if `limit` is not large enough.
async fn inventory_collection_read_all_or_nothing(
&self,
opctx: &OpContext,
id: Uuid,
limit: NonZeroU32,
) -> Result<Collection, Error> {
let (collection, limit_reached) =
self.inventory_collection_read_best_effort(id, limit).await?;
let (collection, limit_reached) = self
.inventory_collection_read_best_effort(opctx, id, limit)
.await?;
bail_unless!(
!limit_reached,
"hit limit of {} records while loading collection",
Expand All @@ -844,10 +850,11 @@ impl DataStore {
/// records and returning up to `limit` records.
async fn inventory_collection_read_best_effort(
&self,
opctx: &OpContext,
id: Uuid,
limit: NonZeroU32,
) -> Result<(Collection, bool), Error> {
let conn = &self.pool_connection_for_tests().await?;
let conn = self.pool_connection_authorized(opctx).await?;
let sql_limit = i64::from(u32::from(limit));
let usize_limit = usize::try_from(u32::from(limit)).unwrap();
let mut limit_reached = false;
Expand All @@ -858,7 +865,7 @@ impl DataStore {
.filter(dsl::id.eq(id))
.limit(2)
.select(InvCollection::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| {
public_error_from_diesel(e, ErrorHandler::Server)
Expand All @@ -879,7 +886,7 @@ impl DataStore {
.order_by(dsl::idx)
.limit(sql_limit)
.select(InvCollectionError::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
Expand All @@ -894,7 +901,7 @@ impl DataStore {
.filter(dsl::inv_collection_id.eq(id))
.limit(sql_limit)
.select(InvServiceProcessor::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
Expand All @@ -915,7 +922,7 @@ impl DataStore {
.filter(dsl::inv_collection_id.eq(id))
.limit(sql_limit)
.select(InvRootOfTrust::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
Expand All @@ -940,7 +947,7 @@ impl DataStore {
.filter(dsl::id.eq_any(baseboard_id_ids))
.limit(sql_limit)
.select(HwBaseboardId::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
Expand Down Expand Up @@ -989,7 +996,7 @@ impl DataStore {
.filter(dsl::inv_collection_id.eq(id))
.limit(sql_limit)
.select(InvCaboose::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| {
public_error_from_diesel(e, ErrorHandler::Server)
Expand All @@ -1009,7 +1016,7 @@ impl DataStore {
.filter(dsl::id.eq_any(sw_caboose_ids))
.limit(sql_limit)
.select(SwCaboose::as_select())
.load_async(&**conn)
.load_async(&*conn)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
Expand Down Expand Up @@ -1117,6 +1124,7 @@ impl DataStoreInventoryTest for DataStore {

#[cfg(test)]
mod test {
use crate::context::OpContext;
use crate::db::datastore::datastore_test;
use crate::db::datastore::inventory::DataStoreInventoryTest;
use crate::db::datastore::DataStore;
Expand All @@ -1139,11 +1147,14 @@ mod test {
use uuid::Uuid;

async fn read_collection(
opctx: &OpContext,
datastore: &DataStore,
id: Uuid,
) -> anyhow::Result<Collection> {
let limit = NonZeroU32::new(1000).unwrap();
datastore.inventory_collection_read_all_or_nothing(id, limit).await
Ok(datastore
.inventory_collection_read_all_or_nothing(opctx, id, limit)
.await?)
}

async fn count_baseboards_cabooses(
Expand Down Expand Up @@ -1189,9 +1200,10 @@ mod test {

// Read it back.
let conn = datastore.pool_connection_for_tests().await.unwrap();
let collection_read = read_collection(&datastore, collection1.id)
.await
.expect("failed to read collection back");
let collection_read =
read_collection(&opctx, &datastore, collection1.id)
.await
.expect("failed to read collection back");
assert_eq!(collection1, collection_read);

// There ought to be no baseboards or cabooses in the databases from
Expand All @@ -1211,9 +1223,10 @@ mod test {
.inventory_insert_collection(&opctx, &collection2)
.await
.expect("failed to insert collection");
let collection_read = read_collection(&datastore, collection2.id)
.await
.expect("failed to read collection back");
let collection_read =
read_collection(&opctx, &datastore, collection2.id)
.await
.expect("failed to read collection back");
assert_eq!(collection2, collection_read);
// Verify that we have exactly the set of cabooses and baseboards in the
// databases that came from this first non-empty collection.
Expand All @@ -1234,9 +1247,10 @@ mod test {
.inventory_insert_collection(&opctx, &collection3)
.await
.expect("failed to insert collection");
let collection_read = read_collection(&datastore, collection3.id)
.await
.expect("failed to read collection back");
let collection_read =
read_collection(&opctx, &datastore, collection3.id)
.await
.expect("failed to read collection back");
assert_eq!(collection3, collection_read);
// Verify that we have the same number of cabooses and baseboards, since
// those didn't change.
Expand Down Expand Up @@ -1278,9 +1292,10 @@ mod test {
.inventory_insert_collection(&opctx, &collection4)
.await
.expect("failed to insert collection");
let collection_read = read_collection(&datastore, collection4.id)
.await
.expect("failed to read collection back");
let collection_read =
read_collection(&opctx, &datastore, collection4.id)
.await
.expect("failed to read collection back");
assert_eq!(collection4, collection_read);
// Verify the number of baseboards and collections again.
assert_eq!(
Expand All @@ -1305,9 +1320,10 @@ mod test {
.inventory_insert_collection(&opctx, &collection5)
.await
.expect("failed to insert collection");
let collection_read = read_collection(&datastore, collection5.id)
.await
.expect("failed to read collection back");
let collection_read =
read_collection(&opctx, &datastore, collection5.id)
.await
.expect("failed to read collection back");
assert_eq!(collection5, collection_read);
assert_eq!(collection5.baseboards.len(), collection3.baseboards.len());
assert_eq!(collection5.cabooses.len(), collection3.cabooses.len());
Expand Down Expand Up @@ -1436,19 +1452,19 @@ mod test {
);

// If we try to fetch a pruned collection, we should get nothing.
let _ = read_collection(&datastore, collection4.id)
let _ = read_collection(&opctx, &datastore, collection4.id)
.await
.expect_err("unexpectedly read pruned collection");

// But we should still be able to fetch the collections that do exist.
let collection_read =
read_collection(&datastore, collection5.id).await.unwrap();
read_collection(&opctx, &datastore, collection5.id).await.unwrap();
assert_eq!(collection5, collection_read);
let collection_read =
read_collection(&datastore, collection6.id).await.unwrap();
read_collection(&opctx, &datastore, collection6.id).await.unwrap();
assert_eq!(collection6, collection_read);
let collection_read =
read_collection(&datastore, collection7.id).await.unwrap();
read_collection(&opctx, &datastore, collection7.id).await.unwrap();
assert_eq!(collection7, collection_read);

// We should prune more than one collection, if needed. We'll wind up
Expand Down

0 comments on commit 1721f7b

Please sign in to comment.