From 4af6221dd26d8463c8edbc810e72c32db93095c2 Mon Sep 17 00:00:00 2001
From: Sean Klein <seanmarionklein@gmail.com>
Date: Thu, 6 Mar 2025 10:49:56 -0800
Subject: [PATCH] Convert paginate_multicolumn
---
nexus/db-model/src/inventory.rs | 2 +-
.../db-queries/src/db/datastore/inventory.rs | 526 +++++++++++-------
nexus/db-queries/src/db/datastore/role.rs | 27 +-
nexus/db-queries/src/db/datastore/saga.rs | 28 +-
nexus/db-queries/src/db/pagination.rs | 32 ++
5 files changed, 383 insertions(+), 232 deletions(-)
diff --git a/nexus/db-model/src/inventory.rs b/nexus/db-model/src/inventory.rs
index 371f0c2456a..fcca8bc7618 100644
--- a/nexus/db-model/src/inventory.rs
+++ b/nexus/db-model/src/inventory.rs
@@ -921,9 +921,9 @@ pub struct InvNvmeDiskFirmware {
inv_collection_id: DbTypedUuid<CollectionKind>,
sled_id: DbTypedUuid<SledKind>,
slot: i64,
+ number_of_slots: SqlU8,
active_slot: SqlU8,
next_active_slot: Option<SqlU8>,
- number_of_slots: SqlU8,
slot1_is_read_only: bool,
slot_firmware_versions: Vec<Option<String>>,
}
diff --git a/nexus/db-queries/src/db/datastore/inventory.rs b/nexus/db-queries/src/db/datastore/inventory.rs
index b6cb5c5609f..0cf5e4019dc 100644
--- a/nexus/db-queries/src/db/datastore/inventory.rs
+++ b/nexus/db-queries/src/db/datastore/inventory.rs
@@ -6,10 +6,12 @@ use super::DataStore;
use crate::authz;
use crate::context::OpContext;
use crate::db;
+use crate::db::column_walker::AllColumnsOf;
use crate::db::error::ErrorHandler;
use crate::db::error::public_error_from_diesel;
use crate::db::error::public_error_from_diesel_lookup;
-use crate::db::pagination::{Paginator, paginated, paginated_multicolumn};
+use crate::db::pagination::Paginator;
+use crate::db::pagination::RawPaginator;
use crate::db::queries::ALLOW_FULL_TABLE_SCAN_SQL;
use anyhow::Context;
use async_bb8_diesel::AsyncConnection;
@@ -25,7 +27,10 @@ use diesel::OptionalExtension;
use diesel::QueryDsl;
use diesel::Table;
use diesel::expression::SelectableHelper;
+use diesel::helper_types::AsSelect;
use diesel::sql_types::Nullable;
+use diesel::pg::Pg;
+use diesel::pg::sql_types::Array;
use futures::FutureExt;
use futures::future::BoxFuture;
use nexus_db_model::CabooseWhichEnum;
@@ -1571,19 +1576,23 @@ impl DataStore {
let mut errors = Vec::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_collection_error,
- dsl::idx,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .order_by(dsl::idx)
- .select(InvCollectionError::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_collection_error>::as_str())
+ .sql(" FROM inv_collection_error WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::idx>()
+ .query::<AsSelect<InvCollectionError, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_collection_error")
+ })?;
paginator =
p.found_batch(&batch, &|row: &InvCollectionError| row.idx);
errors.extend(batch.into_iter().map(|e| e.message));
@@ -1598,18 +1607,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_service_processor,
- dsl::hw_baseboard_id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvServiceProcessor::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_service_processor>::as_str())
+ .sql(" FROM inv_service_processor WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::hw_baseboard_id>()
+ .query::<AsSelect<InvServiceProcessor, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_service_processor")
+ })?;
paginator = p.found_batch(&batch, &|row| row.hw_baseboard_id);
sps.extend(batch.into_iter().map(|row| {
let baseboard_id = row.hw_baseboard_id;
@@ -1629,18 +1643,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_root_of_trust,
- dsl::hw_baseboard_id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvRootOfTrust::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_root_of_trust>::as_str())
+ .sql(" FROM inv_root_of_trust WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::hw_baseboard_id>()
+ .query::<AsSelect<InvRootOfTrust, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_root_of_trust")
+ })?;
paginator = p.found_batch(&batch, &|row| row.hw_baseboard_id);
rots.extend(batch.into_iter().map(|rot_row| {
let baseboard_id = rot_row.hw_baseboard_id;
@@ -1660,18 +1679,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let mut batch = paginated(
- dsl::inv_sled_agent,
- dsl::sled_id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvSledAgent::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_sled_agent>::as_str())
+ .sql(" FROM inv_sled_agent WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let mut batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::sled_id>()
+ .query::<AsSelect<InvSledAgent, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_sled_agent")
+ })?;
paginator = p.found_batch(&batch, &|row| row.sled_id);
rows.append(&mut batch);
}
@@ -1692,18 +1716,26 @@ impl DataStore {
>::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated_multicolumn(
- dsl::inv_nvme_disk_firmware,
- (dsl::sled_id, dsl::slot),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvNvmeDiskFirmware::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_nvme_disk_firmware>::as_str())
+ .sql(" FROM inv_nvme_disk_firmware WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::sled_id, dsl::slot>()
+ .query::<AsSelect<InvNvmeDiskFirmware, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_nvme_disk_firmware")
+ })?;
+
+ println!("Batch: {batch:?}");
+
paginator =
p.found_batch(&batch, &|row| (row.sled_id(), row.slot()));
for firmware in batch {
@@ -1734,18 +1766,24 @@ impl DataStore {
>::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated_multicolumn(
- dsl::inv_physical_disk,
- (dsl::sled_id, dsl::slot),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvPhysicalDisk::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_physical_disk>::as_str())
+ .sql(" FROM inv_physical_disk WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::sled_id, dsl::slot>()
+ .query::<AsSelect<InvPhysicalDisk, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_physical_disk")
+ })?;
+
paginator =
p.found_batch(&batch, &|row| (row.sled_id, row.slot));
for disk in batch {
@@ -1780,18 +1818,23 @@ impl DataStore {
BTreeMap::<Uuid, Vec<nexus_types::inventory::Zpool>>::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated_multicolumn(
- dsl::inv_zpool,
- (dsl::sled_id, dsl::id),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvZpool::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_zpool>::as_str())
+ .sql(" FROM inv_zpool WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::sled_id, dsl::id>()
+ .query::<AsSelect<InvZpool, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_zpool")
+ })?;
paginator = p.found_batch(&batch, &|row| (row.sled_id, row.id));
for zpool in batch {
zpools
@@ -1811,18 +1854,24 @@ impl DataStore {
BTreeMap::<Uuid, Vec<nexus_types::inventory::Dataset>>::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated_multicolumn(
- dsl::inv_dataset,
- (dsl::sled_id, dsl::name),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvDataset::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_dataset>::as_str())
+ .sql(" FROM inv_dataset WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::sled_id, dsl::name>()
+ .query::<AsSelect<InvDataset, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_dataset")
+ })?;
+
paginator = p.found_batch(&batch, &|row| {
(row.sled_id, row.name.clone())
});
@@ -1838,12 +1887,15 @@ impl DataStore {
// Collect the unique baseboard ids referenced by SPs, RoTs, and Sled
// Agents.
- let baseboard_id_ids: BTreeSet<_> = sps
+ let baseboard_id_ids: Vec<_> = sps
.keys()
.chain(rots.keys())
.cloned()
.chain(sled_agent_rows.iter().filter_map(|s| s.hw_baseboard_id))
+ .collect::<BTreeSet<_>>()
+ .into_iter()
.collect();
+
// Fetch the corresponding baseboard records.
let baseboards_by_id: BTreeMap<_, _> = {
use db::schema::hw_baseboard_id::dsl;
@@ -1852,18 +1904,24 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::hw_baseboard_id,
- dsl::id,
- &p.current_pagparams(),
- )
- .filter(dsl::id.eq_any(baseboard_id_ids.clone()))
- .select(HwBaseboardId::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::hw_baseboard_id>::as_str())
+ .sql(" FROM hw_baseboard_id WHERE id = ANY(")
+ .param()
+ .bind::<Array<diesel::sql_types::Uuid>, _>(baseboard_id_ids.clone())
+ .sql(")");
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::id>()
+ .query::<AsSelect<HwBaseboardId, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing hw_baseboard_id")
+ })?;
paginator = p.found_batch(&batch, &|row| row.id);
bbs.extend(batch.into_iter().map(|bb| {
(
@@ -1912,18 +1970,24 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let mut batch = paginated_multicolumn(
- dsl::inv_caboose,
- (dsl::hw_baseboard_id, dsl::which),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvCaboose::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_caboose>::as_str())
+ .sql(" FROM inv_caboose WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let mut batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::hw_baseboard_id, dsl::which>()
+ .query::<AsSelect<InvCaboose, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_caboose")
+ })?;
+
paginator = p.found_batch(&batch, &|row| {
(row.hw_baseboard_id, row.which)
});
@@ -1934,9 +1998,11 @@ impl DataStore {
};
// Collect the unique sw_caboose_ids for those cabooses.
- let sw_caboose_ids: BTreeSet<_> = inv_caboose_rows
+ let sw_caboose_ids: Vec<_> = inv_caboose_rows
.iter()
.map(|inv_caboose| inv_caboose.sw_caboose_id)
+ .collect::<BTreeSet<_>>()
+ .into_iter()
.collect();
// Fetch the corresponing records.
let cabooses_by_id: BTreeMap<_, _> = {
@@ -1946,15 +2012,24 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch =
- paginated(dsl::sw_caboose, dsl::id, &p.current_pagparams())
- .filter(dsl::id.eq_any(sw_caboose_ids.clone()))
- .select(SwCaboose::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::sw_caboose>::as_str())
+ .sql(" FROM sw_caboose WHERE id = ANY(")
+ .param()
+ .bind::<Array<diesel::sql_types::Uuid>, _>(sw_caboose_ids.clone())
+ .sql(")");
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::id>()
+ .query::<AsSelect<SwCaboose, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing sw_caboose")
+ })?;
paginator = p.found_batch(&batch, &|row| row.id);
cabooses.extend(batch.into_iter().map(|sw_caboose_row| {
(
@@ -2014,18 +2089,24 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let mut batch = paginated_multicolumn(
- dsl::inv_root_of_trust_page,
- (dsl::hw_baseboard_id, dsl::which),
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvRotPage::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_root_of_trust_page>::as_str())
+ .sql(" FROM inv_root_of_trust_page WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let mut batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::hw_baseboard_id, dsl::which>()
+ .query::<AsSelect<InvRotPage, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_root_of_trust_page")
+ })?;
+
paginator = p.found_batch(&batch, &|row| {
(row.hw_baseboard_id, row.which)
});
@@ -2036,9 +2117,11 @@ impl DataStore {
};
// Collect the unique sw_rot_page_ids for those pages.
- let sw_rot_page_ids: BTreeSet<_> = inv_rot_page_rows
+ let sw_rot_page_ids: Vec<_> = inv_rot_page_rows
.iter()
.map(|inv_rot_page| inv_rot_page.sw_root_of_trust_page_id)
+ .collect::<BTreeSet<_>>()
+ .into_iter()
.collect();
// Fetch the corresponding records.
let rot_pages_by_id: BTreeMap<_, _> = {
@@ -2048,18 +2131,25 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::sw_root_of_trust_page,
- dsl::id,
- &p.current_pagparams(),
- )
- .filter(dsl::id.eq_any(sw_rot_page_ids.clone()))
- .select(SwRotPage::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::sw_root_of_trust_page>::as_str())
+ .sql(" FROM sw_root_of_trust_page WHERE id = ANY(")
+ .param()
+ .bind::<Array<diesel::sql_types::Uuid>, _>(sw_rot_page_ids.clone())
+ .sql(")");
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::id>()
+ .query::<AsSelect<SwRotPage, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing sw_root_of_trust_page")
+ })?;
+
paginator = p.found_batch(&batch, &|row| row.id);
rot_pages.extend(batch.into_iter().map(|sw_rot_page_row| {
(
@@ -2128,18 +2218,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_sled_omicron_zones,
- dsl::sled_id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvSledOmicronZones::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_sled_omicron_zones>::as_str())
+ .sql(" FROM inv_sled_omicron_zones WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::sled_id>()
+ .query::<AsSelect<InvSledOmicronZones, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_sled_omicron_zones")
+ })?;
paginator = p.found_batch(&batch, &|row| row.sled_id);
zones.extend(batch.into_iter().map(|sled_zones_config| {
(
@@ -2167,18 +2262,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_omicron_zone_nic,
- dsl::id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvOmicronZoneNic::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_omicron_zone_nic>::as_str())
+ .sql(" FROM inv_omicron_zone_nic WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::id>()
+ .query::<AsSelect<InvOmicronZoneNic, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_omicron_zone_nic")
+ })?;
paginator = p.found_batch(&batch, &|row| row.id);
nics.extend(batch.into_iter().map(|found_zone_nic| {
(found_zone_nic.id, found_zone_nic)
@@ -2196,22 +2296,23 @@ impl DataStore {
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let mut batch = paginated(
- dsl::inv_omicron_zone,
- dsl::id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- // It's not strictly necessary to order these by id. Doing so
- // ensures a consistent representation for `Collection`, which
- // makes testing easier. It's already indexed to do this, too.
- .order_by(dsl::id)
- .select(InvOmicronZone::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_omicron_zone>::as_str())
+ .sql(" FROM inv_omicron_zone WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let mut batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::id>()
+ .query::<AsSelect<InvOmicronZone, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_omicron_zone")
+ })?;
paginator = p.found_batch(&batch, &|row| row.id);
zones.append(&mut batch);
}
@@ -2264,18 +2365,23 @@ impl DataStore {
let mut memberships = BTreeSet::new();
let mut paginator = Paginator::new(batch_size);
while let Some(p) = paginator.next() {
- let batch = paginated(
- dsl::inv_clickhouse_keeper_membership,
- dsl::queried_keeper_id,
- &p.current_pagparams(),
- )
- .filter(dsl::inv_collection_id.eq(db_id))
- .select(InvClickhouseKeeperMembership::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| {
- public_error_from_diesel(e, ErrorHandler::Server)
- })?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::inv_clickhouse_keeper_membership>::as_str())
+ .sql(" FROM inv_clickhouse_keeper_membership WHERE inv_collection_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(db_id);
+ let batch = pagination_query
+ .with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_column::<dsl::queried_keeper_id>()
+ .query::<AsSelect<InvClickhouseKeeperMembership, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| {
+ public_error_from_diesel(e, ErrorHandler::Server)
+ .internal_context("SELECT-ing inv_clickhouse_keeper_membership")
+ })?;
paginator = p.found_batch(&batch, &|row| row.queried_keeper_id);
for membership in batch.into_iter() {
memberships.insert(
diff --git a/nexus/db-queries/src/db/datastore/role.rs b/nexus/db-queries/src/db/datastore/role.rs
index c4e106705d5..1797345fa7d 100644
--- a/nexus/db-queries/src/db/datastore/role.rs
+++ b/nexus/db-queries/src/db/datastore/role.rs
@@ -9,6 +9,7 @@ use crate::authz;
use crate::authz::AuthorizedResource;
use crate::context::OpContext;
use crate::db;
+use crate::db::column_walker::AllColumnsOf;
use crate::db::datastore::RunnableQuery;
use crate::db::datastore::RunnableQueryNoReturn;
use crate::db::error::ErrorHandler;
@@ -18,10 +19,12 @@ use crate::db::model::DatabaseString;
use crate::db::model::IdentityType;
use crate::db::model::RoleAssignment;
use crate::db::model::RoleBuiltin;
-use crate::db::pagination::paginated_multicolumn;
+use crate::db::pagination::RawPaginator;
use crate::db::pool::DbConnection;
use async_bb8_diesel::AsyncRunQueryDsl;
use diesel::prelude::*;
+use diesel::helper_types::AsSelect;
+use diesel::pg::Pg;
use nexus_db_fixed_data::role_assignment::BUILTIN_ROLE_ASSIGNMENTS;
use nexus_db_fixed_data::role_builtin::BUILTIN_ROLES;
use nexus_types::external_api::shared;
@@ -38,18 +41,22 @@ impl DataStore {
pagparams: &DataPageParams<'_, (String, String)>,
) -> ListResultVec<RoleBuiltin> {
use db::schema::role_builtin::dsl;
+
opctx.authorize(authz::Action::ListChildren, &authz::FLEET).await?;
let conn = self.pool_connection_authorized(opctx).await?;
- paginated_multicolumn(
- dsl::role_builtin,
- (dsl::resource_type, dsl::role_name),
- pagparams,
- )
- .select(RoleBuiltin::as_select())
- .load_async::<RoleBuiltin>(&*conn)
- .await
- .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
+
+ let mut paginator = RawPaginator::new();
+ paginator.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::role_builtin>::as_str());
+ paginator
+ .with_parameters(pagparams)
+ .paginate_by_diesel_columns::<dsl::resource_type, dsl::role_name>()
+ .query::<AsSelect<RoleBuiltin, Pg>>()
+ .load_async::<RoleBuiltin>(&*conn)
+ .await
+ .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
}
/// Load built-in roles into the database
diff --git a/nexus/db-queries/src/db/datastore/saga.rs b/nexus/db-queries/src/db/datastore/saga.rs
index 1e9e24b68ad..7e9bd14a9b9 100644
--- a/nexus/db-queries/src/db/datastore/saga.rs
+++ b/nexus/db-queries/src/db/datastore/saga.rs
@@ -7,15 +7,18 @@
use super::DataStore;
use super::SQL_BATCH_SIZE;
use crate::db;
+use crate::db::column_walker::AllColumnsOf;
use crate::db::error::ErrorHandler;
use crate::db::error::public_error_from_diesel;
use crate::db::pagination::Paginator;
+use crate::db::pagination::RawPaginator;
use crate::db::pagination::paginated;
-use crate::db::pagination::paginated_multicolumn;
use crate::db::update_and_check::UpdateAndCheck;
use crate::db::update_and_check::UpdateStatus;
use async_bb8_diesel::AsyncRunQueryDsl;
use diesel::prelude::*;
+use diesel::helper_types::AsSelect;
+use diesel::pg::Pg;
use nexus_auth::authz;
use nexus_auth::context::OpContext;
use omicron_common::api::external::Error;
@@ -175,16 +178,19 @@ impl DataStore {
let conn = self.pool_connection_authorized(opctx).await?;
while let Some(p) = paginator.next() {
use db::schema::saga_node_event::dsl;
- let batch = paginated_multicolumn(
- dsl::saga_node_event,
- (dsl::node_id, dsl::event_type),
- &p.current_pagparams(),
- )
- .filter(dsl::saga_id.eq(saga_id))
- .select(db::saga_types::SagaNodeEvent::as_select())
- .load_async(&*conn)
- .await
- .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?;
+ let mut pagination_query = RawPaginator::new();
+ pagination_query.source()
+ .sql("SELECT ")
+ .sql(AllColumnsOf::<dsl::saga_node_event>::as_str())
+ .sql(" FROM saga_node_event WHERE saga_id = ")
+ .param()
+ .bind::<diesel::sql_types::Uuid, _>(saga_id);
+ let batch = pagination_query.with_parameters(&p.current_pagparams())
+ .paginate_by_diesel_columns::<dsl::node_id, dsl::event_type>()
+ .query::<AsSelect<db::saga_types::SagaNodeEvent, Pg>>()
+ .load_async(&*conn)
+ .await
+ .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))?;
paginator = p.found_batch(&batch, &|row| {
(row.node_id, row.event_type.clone())
diff --git a/nexus/db-queries/src/db/pagination.rs b/nexus/db-queries/src/db/pagination.rs
index b10f0d9701c..b5e041d2d6a 100644
--- a/nexus/db-queries/src/db/pagination.rs
+++ b/nexus/db-queries/src/db/pagination.rs
@@ -274,6 +274,20 @@ pub struct RawPaginatorWithParams<'a, M> {
}
impl<M> RawPaginatorWithParams<'_, M> {
+ /// Short-hand for [Self::paginate_by_column], when paginating on a
+ /// Diesel column.
+ pub fn paginate_by_diesel_column<C>(
+ self,
+ ) -> QueryBuilder
+ where
+ C: diesel::Column,
+ C::SqlType: Send + 'static,
+ Pg: diesel::sql_types::HasSqlType<C::SqlType>,
+ M: diesel::serialize::ToSql<C::SqlType, Pg> + Send + Clone + 'static,
+ {
+ self.paginate_by_column(C::NAME)
+ }
+
pub fn paginate_by_column<BindSt>(
mut self,
key_column: &'static str,
@@ -314,6 +328,24 @@ impl<M> RawPaginatorWithParams<'_, M> {
}
impl<M1, M2> RawPaginatorWithParams<'_, (M1, M2)> {
+ /// Short-hand for [Self::paginate_by_columns], when paginating on a
+ /// Diesel column.
+ pub fn paginate_by_diesel_columns<C1, C2>(
+ self,
+ ) -> QueryBuilder
+ where
+ C1: diesel::Column,
+ C1::SqlType: Send + 'static,
+ C2: diesel::Column,
+ C2::SqlType: Send + 'static,
+ Pg: diesel::sql_types::HasSqlType<C1::SqlType>,
+ Pg: diesel::sql_types::HasSqlType<C2::SqlType>,
+ M1: diesel::serialize::ToSql<C1::SqlType, Pg> + Send + Clone + 'static,
+ M2: diesel::serialize::ToSql<C2::SqlType, Pg> + Send + Clone + 'static,
+ {
+ self.paginate_by_columns(C1::NAME, C2::NAME)
+ }
+
pub fn paginate_by_columns<BindSt1, BindSt2>(
mut self,
key_column1: &'static str,