Skip to content

Commit

Permalink
Manual pagination of both tables at once
Browse files Browse the repository at this point in the history
  • Loading branch information
@smklein
smklein committed Feb 24, 2025
1 parent 4e6362d commit f50749d
Show file tree
Hide file tree
Showing 5 changed files with 95 additions and 111 deletions.
8 changes: 4 additions & 4 deletions common/src/api/external/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1390,18 +1390,18 @@ impl SimpleIdentity for AffinityGroupMember {
#[derive(Clone, Debug, Deserialize, Serialize, JsonSchema, PartialEq)]
#[serde(tag = "type", content = "value", rename_all = "snake_case")]
pub enum AntiAffinityGroupMember {
/// An instance belonging to this group, identified by UUID.
Instance(InstanceUuid),

/// An affinity group belonging to this group, identified by UUID.
AffinityGroup(AffinityGroupUuid),

/// An instance belonging to this group, identified by UUID.
Instance(InstanceUuid),
}

impl SimpleIdentity for AntiAffinityGroupMember {
fn id(&self) -> Uuid {
match self {
AntiAffinityGroupMember::Instance(id) => *id.as_untyped_uuid(),
AntiAffinityGroupMember::AffinityGroup(id) => *id.as_untyped_uuid(),
AntiAffinityGroupMember::Instance(id) => *id.as_untyped_uuid(),
}
}
}
Expand Down
1 change: 1 addition & 0 deletions nexus/db-model/src/schema.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2079,6 +2079,7 @@ allow_tables_to_appear_in_same_query!(hw_baseboard_id, inv_sled_agent,);

allow_tables_to_appear_in_same_query!(
anti_affinity_group,
anti_affinity_group_affinity_membership,
anti_affinity_group_instance_membership,
affinity_group,
affinity_group_instance_membership,
Expand Down
181 changes: 84 additions & 97 deletions nexus/db-queries/src/db/datastore/affinity.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,13 +24,15 @@ use crate::db::model::AntiAffinityGroupUpdate;
use crate::db::model::Name;
use crate::db::model::Project;
use crate::db::pagination::paginated;
use crate::db::raw_query_builder::QueryBuilder;
use crate::transaction_retry::OptionalError;
use async_bb8_diesel::AsyncRunQueryDsl;
use chrono::Utc;
use diesel::prelude::*;
use omicron_common::api::external;
use omicron_common::api::external::http_pagination::PaginatedBy;
use omicron_common::api::external::CreateResult;
use omicron_common::api::external::DataPageParams;
use omicron_common::api::external::DeleteResult;
use omicron_common::api::external::Error;
use omicron_common::api::external::ListResultVec;
Expand All @@ -42,6 +44,7 @@ use omicron_uuid_kinds::AntiAffinityGroupUuid;
use omicron_uuid_kinds::GenericUuid;
use omicron_uuid_kinds::InstanceUuid;
use ref_cast::RefCast;
use uuid::Uuid;

impl DataStore {
pub async fn affinity_group_list(
Expand Down Expand Up @@ -365,39 +368,76 @@ impl DataStore {
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
}

pub async fn anti_affinity_group_member_instance_list(
pub async fn anti_affinity_group_member_list(
&self,
opctx: &OpContext,
authz_anti_affinity_group: &authz::AntiAffinityGroup,
pagparams: &PaginatedBy<'_>,
) -> ListResultVec<AntiAffinityGroupInstanceMembership> {
pagparams: &DataPageParams<'_, Uuid>,
) -> ListResultVec<external::AntiAffinityGroupMember> {
opctx.authorize(authz::Action::Read, authz_anti_affinity_group).await?;

// TODO: Need to also look up "group_membership" here
// TODO: definitely test listing both?
// TODO: atlernately - make this API "instance member" specific, make a
// different one for "affinity group members", and paginate over both.
//
// That might be preferable.

use db::schema::anti_affinity_group_instance_membership::dsl;
match pagparams {
PaginatedBy::Id(pagparams) => paginated(
dsl::anti_affinity_group_instance_membership,
dsl::instance_id,
&pagparams,
),
PaginatedBy::Name(_) => {
return Err(Error::invalid_request(
"Cannot paginate group members by name",
));
let mut query = QueryBuilder::new()
.sql(
"
SELECT instance_id as id, 'instance' as label
FROM anti_affinity_group_instance_membership
WHERE group_id = ",
)
.param()
.bind::<diesel::sql_types::Uuid, _>(authz_anti_affinity_group.id())
.sql(
"
UNION
SELECT affinity_group_id as id, 'affinity_group' as label
FROM anti_affinity_group_affinity_membership
WHERE anti_affinity_group_id = ",
)
.param()
.bind::<diesel::sql_types::Uuid, _>(authz_anti_affinity_group.id())
.sql(" ");

let (sort, cmp) = match pagparams.direction {
dropshot::PaginationOrder::Ascending => (" ORDER BY id ASC ", ">"),
dropshot::PaginationOrder::Descending => {
(" ORDER BY id DESC ", "<")
}
}
.filter(dsl::group_id.eq(authz_anti_affinity_group.id()))
.select(AntiAffinityGroupInstanceMembership::as_select())
.load_async(&*self.pool_connection_authorized(opctx).await?)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
};
if let Some(id) = pagparams.marker {
query = query
.sql("WHERE id ")
.sql(cmp)
.sql(" ")
.param()
.bind::<diesel::sql_types::Uuid, _>(*id);
};

query = query.sql(sort);
query =
query.sql(" LIMIT ").param().bind::<diesel::sql_types::BigInt, _>(
i64::from(pagparams.limit.get()),
);

Ok(query
.query::<(diesel::sql_types::Uuid, diesel::sql_types::Text)>()
.load_async::<(Uuid, String)>(
&*self.pool_connection_authorized(opctx).await?,
)
.await
.map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?
.into_iter()
.map(|(id, label)| {
use external::AntiAffinityGroupMember as Member;
match label.as_str() {
"affinity_group" => Member::AffinityGroup(
AffinityGroupUuid::from_untyped_uuid(id),
),
"instance" => {
Member::Instance(InstanceUuid::from_untyped_uuid(id))
}
other => panic!("Unexpected label from query: {other}"),
}
})
.collect())
}

pub async fn affinity_group_member_view(
Expand Down Expand Up @@ -1804,18 +1844,13 @@ mod tests {
.unwrap();

// A new group should have no members
let pagparams_id = DataPageParams {
let pagparams = DataPageParams {
marker: None,
limit: NonZeroU32::new(100).unwrap(),
direction: dropshot::PaginationOrder::Ascending,
};
let pagbyid = PaginatedBy::Id(pagparams_id);
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -1843,11 +1878,7 @@ mod tests {

// We should now be able to list the new member
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert_eq!(members.len(), 1);
Expand All @@ -1870,11 +1901,7 @@ mod tests {
.await
.unwrap();
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2048,18 +2075,13 @@ mod tests {
.unwrap();

// A new group should have no members
let pagparams_id = DataPageParams {
let pagparams = DataPageParams {
marker: None,
limit: NonZeroU32::new(100).unwrap(),
direction: dropshot::PaginationOrder::Ascending,
};
let pagbyid = PaginatedBy::Id(pagparams_id);
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2123,11 +2145,7 @@ mod tests {

// We should now be able to list the new member
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert_eq!(members.len(), 1);
Expand All @@ -2151,11 +2169,7 @@ mod tests {
.await
.unwrap();
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2266,18 +2280,13 @@ mod tests {
.unwrap();

// A new group should have no members
let pagparams_id = DataPageParams {
let pagparams = DataPageParams {
marker: None,
limit: NonZeroU32::new(100).unwrap(),
direction: dropshot::PaginationOrder::Ascending,
};
let pagbyid = PaginatedBy::Id(pagparams_id);
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2309,11 +2318,7 @@ mod tests {

// Confirm that no instance members exist
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2433,18 +2438,13 @@ mod tests {
.unwrap();

// A new group should have no members
let pagparams_id = DataPageParams {
let pagparams = DataPageParams {
marker: None,
limit: NonZeroU32::new(100).unwrap(),
direction: dropshot::PaginationOrder::Ascending,
};
let pagbyid = PaginatedBy::Id(pagparams_id);
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -2481,11 +2481,7 @@ mod tests {

// Confirm that no instance members exist
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down Expand Up @@ -3014,18 +3010,13 @@ mod tests {
//
// Two calls to "anti_affinity_group_member_add" should be the same
// as a single call.
let pagparams_id = DataPageParams {
let pagparams = DataPageParams {
marker: None,
limit: NonZeroU32::new(100).unwrap(),
direction: dropshot::PaginationOrder::Ascending,
};
let pagbyid = PaginatedBy::Id(pagparams_id);
let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert_eq!(members.len(), 1);
Expand Down Expand Up @@ -3063,11 +3054,7 @@ mod tests {
);

let members = datastore
.anti_affinity_group_member_instance_list(
&opctx,
&authz_group,
&pagbyid,
)
.anti_affinity_group_member_list(&opctx, &authz_group, &pagparams)
.await
.unwrap();
assert!(members.is_empty());
Expand Down
13 changes: 5 additions & 8 deletions nexus/src/app/affinity.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ use nexus_types::external_api::views;
use omicron_common::api::external;
use omicron_common::api::external::http_pagination::PaginatedBy;
use omicron_common::api::external::CreateResult;
use omicron_common::api::external::DataPageParams;
use omicron_common::api::external::DeleteResult;
use omicron_common::api::external::Error;
use omicron_common::api::external::ListResultVec;
Expand Down Expand Up @@ -254,22 +255,18 @@ impl super::Nexus {
&self,
opctx: &OpContext,
anti_affinity_group_lookup: &lookup::AntiAffinityGroup<'_>,
pagparams: &PaginatedBy<'_>,
pagparams: &DataPageParams<'_, uuid::Uuid>,
) -> ListResultVec<external::AntiAffinityGroupMember> {
let (.., authz_anti_affinity_group) = anti_affinity_group_lookup
.lookup_for(authz::Action::ListChildren)
.await?;
Ok(self
.db_datastore
.anti_affinity_group_member_instance_list(
self.db_datastore
.anti_affinity_group_member_list(
opctx,
&authz_anti_affinity_group,
pagparams,
)
.await?
.into_iter()
.map(Into::into)
.collect())
.await
}

pub(crate) async fn affinity_group_member_view(
Expand Down
Loading

0 comments on commit f50749d

Please sign in to comment.