Unify queries for all db_metadata_nexus states

Author: smklein

nexus/db-queries/src/db/datastore/db_metadata.rs

@@ -300,33 +300,18 @@ impl DatastoreSetupAction {
 }
 
 impl DataStore {
-    // Returns the active Nexuses
-    pub async fn get_active_db_metadata_nexus(
+    /// Returns [`DbMetadataNexus`] records in any of the supplied states.
+    pub async fn get_db_metadata_nexus_in_state(
         &self,
         opctx: &OpContext,
+        states: &[DbMetadataNexusState],
     ) -> Result<Vec<DbMetadataNexus>, Error> {
         use nexus_db_schema::schema::db_metadata_nexus::dsl;
 
         opctx.authorize(authz::Action::Read, &authz::FLEET).await?;
 
         dsl::db_metadata_nexus
-            .filter(dsl::state.eq(DbMetadataNexusState::Active))
-            .load_async(&*self.pool_connection_authorized(&opctx).await?)
-            .await
-            .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))
-    }
-
-    // Returns the 'not yet' Nexuses
-    pub async fn get_not_yet_db_metadata_nexus(
-        &self,
-        opctx: &OpContext,
-    ) -> Result<Vec<DbMetadataNexus>, Error> {
-        use nexus_db_schema::schema::db_metadata_nexus::dsl;
-
-        opctx.authorize(authz::Action::Read, &authz::FLEET).await?;
-
-        dsl::db_metadata_nexus
-            .filter(dsl::state.eq(DbMetadataNexusState::NotYet))
+            .filter(dsl::state.eq_any(states.to_vec()))
             .load_async(&*self.pool_connection_authorized(&opctx).await?)
             .await
             .map_err(|e| public_error_from_diesel(e, ErrorHandler::Server))

nexus/reconfigurator/execution/src/database.rs

@@ -5,6 +5,7 @@
 //! Manages deployment of records into the database.
 
 use anyhow::anyhow;
+use nexus_db_model::DbMetadataNexusState;
 use nexus_db_queries::context::OpContext;
 use nexus_db_queries::db::DataStore;
 use nexus_types::deployment::Blueprint;
@@ -23,7 +24,7 @@ pub(crate) async fn deploy_db_metadata_nexus_records(
     // can lag behind it if we are one of those Nexuses running after quiescing
     // has started.
     let active_nexus_zones = datastore
-        .get_active_db_metadata_nexus(opctx)
+        .get_db_metadata_nexus_in_state(opctx, &[DbMetadataNexusState::Active])
         .await?
         .into_iter()
         .map(|z| z.nexus_id())

nexus/reconfigurator/execution/src/dns.rs

@@ -313,6 +313,7 @@ mod test {
     use internal_dns_types::names::BOUNDARY_NTP_DNS_NAME;
     use internal_dns_types::names::DNS_ZONE;
     use internal_dns_types::names::ServiceName;
+    use nexus_db_model::DbMetadataNexusState;
     use nexus_db_model::DnsGroup;
     use nexus_db_model::Silo;
     use nexus_db_queries::authn;
@@ -1503,7 +1504,10 @@ mod test {
         let ip_pool_range_rows =
             fetch_all_service_ip_pool_ranges(&datastore, &opctx).await;
         let active_nexus_zones = datastore
-            .get_active_db_metadata_nexus(&opctx)
+            .get_db_metadata_nexus_in_state(
+                &opctx,
+                &[DbMetadataNexusState::Active],
+            )
             .await
             .internal_context("fetching active nexuses")
             .unwrap()

nexus/reconfigurator/preparation/src/lib.rs

@@ -6,6 +6,7 @@
 
 use anyhow::Context;
 use futures::StreamExt;
+use nexus_db_model::DbMetadataNexusState;
 use nexus_db_model::DnsGroup;
 use nexus_db_model::Generation;
 use nexus_db_queries::context::OpContext;
@@ -204,20 +205,26 @@ impl PlanningInputFromDb<'_> {
             .await
             .internal_context("fetching oximeter read policy")?;
 
-        let active_nexus_zones = datastore
-            .get_active_db_metadata_nexus(opctx)
-            .await
-            .internal_context("fetching active nexuses")?
-            .into_iter()
-            .map(|z| z.nexus_id())
-            .collect::<Vec<_>>();
-        let not_yet_nexus_zones = datastore
-            .get_not_yet_db_metadata_nexus(opctx)
-            .await
-            .internal_context("fetching 'not yet' nexuses")?
-            .into_iter()
-            .map(|z| z.nexus_id())
-            .collect::<Vec<_>>();
+        let (active_nexus_zones, not_yet_nexus_zones): (Vec<_>, Vec<_>) =
+            datastore
+                .get_db_metadata_nexus_in_state(
+                    opctx,
+                    &[
+                        DbMetadataNexusState::Active,
+                        DbMetadataNexusState::NotYet,
+                    ],
+                )
+                .await
+                .internal_context("fetching db_metdata_nexus records")?
+                .into_iter()
+                .partition(|nexus| {
+                    nexus.state() == DbMetadataNexusState::Active
+                });
+
+        let active_nexus_zones =
+            active_nexus_zones.into_iter().map(|n| n.nexus_id()).collect();
+        let not_yet_nexus_zones =
+            not_yet_nexus_zones.into_iter().map(|n| n.nexus_id()).collect();
 
         let planning_input = PlanningInputFromDb {
             sled_rows: &sled_rows,