From 1721f7b24f9fc7abb14a918438fbf3af2fe172fa Mon Sep 17 00:00:00 2001 From: "Andrew J. Stone" Date: Wed, 15 Nov 2023 22:44:43 +0000 Subject: [PATCH] wip --- .../db-queries/src/db/datastore/inventory.rs | 78 +++++++++++-------- 1 file changed, 47 insertions(+), 31 deletions(-) diff --git a/nexus/db-queries/src/db/datastore/inventory.rs b/nexus/db-queries/src/db/datastore/inventory.rs index d94a2aec29e..7e8ac0b7dea 100644 --- a/nexus/db-queries/src/db/datastore/inventory.rs +++ b/nexus/db-queries/src/db/datastore/inventory.rs @@ -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()) } @@ -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 { - 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", @@ -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; @@ -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) @@ -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() @@ -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() @@ -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() @@ -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() @@ -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) @@ -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() @@ -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; @@ -1139,11 +1147,14 @@ mod test { use uuid::Uuid; async fn read_collection( + opctx: &OpContext, datastore: &DataStore, id: Uuid, ) -> anyhow::Result { 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( @@ -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 @@ -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. @@ -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. @@ -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!( @@ -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()); @@ -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