[wip] Convert heavy generic CTEs to using the raw QueryBuilder

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

@@ -9,8 +9,9 @@ use crate::authz;
 use crate::authz::ApiResource;
 use crate::context::OpContext;
 use crate::db;
+use crate::db::collection_attach;
 use crate::db::collection_attach::AttachError;
-use crate::db::collection_attach::DatastoreAttachTarget;
+use crate::db::collection_attach::AttachQuery;
 use crate::db::collection_detach::DatastoreDetachTarget;
 use crate::db::collection_detach::DetachError;
 use crate::db::collection_insert::AsyncInsertError;
@@ -29,7 +30,7 @@ use crate::db::model::VirtualProvisioningResource;
 use crate::db::model::Volume;
 use crate::db::model::to_db_typed_uuid;
 use crate::db::pagination::paginated;
-use crate::db::queries::disk::DiskSetClauseForAttach;
+use crate::db::raw_query_builder::QueryBuilder;
 use crate::db::update_and_check::UpdateAndCheck;
 use crate::db::update_and_check::UpdateStatus;
 use async_bb8_diesel::AsyncRunQueryDsl;
@@ -58,6 +59,75 @@ use std::collections::HashSet;
 use std::net::SocketAddrV6;
 use uuid::Uuid;
 
+/// Creates a database query template for attaching disks to instances.
+///
+/// This defines the complete SQL structure for the attach operation,
+/// including the SET clause and filters.
+fn disk_attach_query_template() -> collection_attach::AttachQueryTemplate {
+    use crate::db::queries::disk::build_next_disk_slot_subquery;
+    use diesel::sql_types;
+
+    let ok_to_attach_disk_states = [
+        api::external::DiskState::Creating,
+        api::external::DiskState::Detached,
+    ];
+    let ok_to_attach_disk_state_labels: Vec<_> =
+        ok_to_attach_disk_states.iter().map(|s| s.label()).collect();
+
+    // TODO(https://github.com/oxidecomputer/omicron/issues/811):
+    // This list of instance attach states is more restrictive than it
+    // plausibly could be.
+    //
+    // We currently only permit attaching disks to stopped instances.
+    let ok_to_attach_instance_states =
+        [db::model::InstanceState::Creating, db::model::InstanceState::NoVmm];
+
+    collection_attach::AttachQueryTemplate::new(
+        collection_attach::Collection::new(
+            "instance",     // collection table
+            "id",           // collection id column
+            "time_deleted", // collection time_deleted
+        ),
+        collection_attach::Resource::new(
+            "disk",               // resource table
+            "id",                 // resource id column
+            "attach_instance_id", // resource FK column
+            "time_deleted",       // resource time_deleted
+        ),
+        false, // allow_from_attached
+        |builder: &mut QueryBuilder, instance_id: Uuid| {
+            // Build SET clause: attach_instance_id, disk_state, slot
+            let attached_label =
+                api::external::DiskState::Attached(instance_id).label();
+            builder.sql("attach_instance_id = ");
+            builder.param().bind::<sql_types::Uuid, _>(instance_id);
+            builder.sql(", disk_state = ");
+            builder.param().bind::<sql_types::Text, _>(attached_label);
+            builder.sql(", slot = (");
+            build_next_disk_slot_subquery(builder, instance_id);
+            builder.sql(")");
+        },
+        Some(move |builder: &mut QueryBuilder| {
+            // Collection (instance) filter: state and active_propolis_id
+            builder.sql(" AND state = ANY(");
+            builder.param().bind::<sql_types::Array<
+                nexus_db_schema::enums::InstanceStateEnum,
+            >, _>(
+                ok_to_attach_instance_states.to_vec()
+            );
+            builder.sql(") AND active_propolis_id IS NULL");
+        }),
+        Some(move |builder: &mut QueryBuilder| {
+            // Resource (disk) filter: disk_state
+            builder.sql(" AND disk_state = ANY(");
+            builder.param().bind::<sql_types::Array<sql_types::Text>, _>(
+                ok_to_attach_disk_state_labels,
+            );
+            builder.sql(")");
+        }),
+    )
+}
+
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub enum Disk {
     Crucible(CrucibleDisk),
@@ -503,44 +573,18 @@ impl DataStore {
         authz_disk: &authz::Disk,
         max_disks: u32,
     ) -> Result<(Instance, Disk), Error> {
-        use nexus_db_schema::schema::{disk, instance};
-
         opctx.authorize(authz::Action::Modify, authz_instance).await?;
         opctx.authorize(authz::Action::Modify, authz_disk).await?;
 
-        let ok_to_attach_disk_states = [
-            api::external::DiskState::Creating,
-            api::external::DiskState::Detached,
-        ];
-        let ok_to_attach_disk_state_labels: Vec<_> =
-            ok_to_attach_disk_states.iter().map(|s| s.label()).collect();
-
-        // TODO(https://github.com/oxidecomputer/omicron/issues/811):
-        // This list of instance attach states is more restrictive than it
-        // plausibly could be.
-        //
-        // We currently only permit attaching disks to stopped instances.
-        let ok_to_attach_instance_states = vec![
-            db::model::InstanceState::Creating,
-            db::model::InstanceState::NoVmm,
-        ];
+        let instance_id = authz_instance.id();
+        let disk_id = authz_disk.id();
 
-        let attach_update = DiskSetClauseForAttach::new(authz_instance.id());
+        // Create the query template with all SQL structure
+        let template = disk_attach_query_template();
 
-        let query = Instance::attach_resource(
-            authz_instance.id(),
-            authz_disk.id(),
-            instance::table.into_boxed().filter(
-                instance::dsl::state
-                    .eq_any(ok_to_attach_instance_states)
-                    .and(instance::dsl::active_propolis_id.is_null()),
-            ),
-            disk::table.into_boxed().filter(
-                disk::dsl::disk_state.eq_any(ok_to_attach_disk_state_labels),
-            ),
-            max_disks,
-            diesel::update(disk::dsl::disk).set(attach_update),
-        );
+        // Build the query with runtime parameter values
+        let query: AttachQuery<model::Disk, model::Instance> =
+            template.build(instance_id, disk_id, max_disks);
 
         let conn = self.pool_connection_authorized(opctx).await?;
 

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

@@ -8,8 +8,9 @@ use super::DataStore;
 use super::SQL_BATCH_SIZE;
 use crate::authz;
 use crate::context::OpContext;
+use crate::db::collection_attach;
 use crate::db::collection_attach::AttachError;
-use crate::db::collection_attach::DatastoreAttachTarget;
+use crate::db::collection_attach::AttachQuery;
 use crate::db::collection_detach::DatastoreDetachTarget;
 use crate::db::collection_detach::DetachError;
 use crate::db::model::ExternalIp;
@@ -25,6 +26,7 @@ use crate::db::queries::external_ip::MAX_EXTERNAL_IPS_PER_INSTANCE;
 use crate::db::queries::external_ip::NextExternalIp;
 use crate::db::queries::external_ip::SAFE_TO_ATTACH_INSTANCE_STATES;
 use crate::db::queries::external_ip::SAFE_TO_ATTACH_INSTANCE_STATES_CREATING;
+use crate::db::raw_query_builder::QueryBuilder;
 use crate::db::update_and_check::UpdateAndCheck;
 use crate::db::update_and_check::UpdateStatus;
 use async_bb8_diesel::AsyncRunQueryDsl;
@@ -62,6 +64,72 @@ use uuid::Uuid;
 
 const MAX_EXTERNAL_IPS_PLUS_SNAT: u32 = MAX_EXTERNAL_IPS_PER_INSTANCE + 1;
 
+/// Creates a database query template for attaching external IPs to instances.
+///
+/// This defines the complete SQL structure for the attach operation,
+/// including the SET clause and filters.
+fn external_ip_attach_query_template(
+    creating_instance: bool,
+    kind: IpKind,
+) -> collection_attach::AttachQueryTemplate {
+    use diesel::sql_types;
+    use nexus_db_schema::enums::{
+        InstanceStateEnum, IpAttachStateEnum, IpKindEnum,
+    };
+
+    let safe_states = if creating_instance {
+        SAFE_TO_ATTACH_INSTANCE_STATES_CREATING.to_vec()
+    } else {
+        SAFE_TO_ATTACH_INSTANCE_STATES.to_vec()
+    };
+
+    collection_attach::AttachQueryTemplate::new(
+        collection_attach::Collection::new(
+            "instance",     // collection table
+            "id",           // collection id column
+            "time_deleted", // collection time_deleted
+        ),
+        collection_attach::Resource::new(
+            "external_ip",  // resource table
+            "id",           // resource id column
+            "parent_id",    // resource FK column
+            "time_deleted", // resource time_deleted
+        ),
+        false, // allow_from_attached
+        |builder: &mut QueryBuilder, instance_id: uuid::Uuid| {
+            // Build SET clause: parent_id, time_modified, state
+            builder.sql("parent_id = ");
+            builder.param().bind::<sql_types::Nullable<sql_types::Uuid>, _>(
+                Some(instance_id),
+            );
+            builder.sql(", time_modified = ");
+            builder.param().bind::<sql_types::Timestamptz, _>(Utc::now());
+            builder.sql(", state = ");
+            builder
+                .param()
+                .bind::<IpAttachStateEnum, _>(IpAttachState::Attaching);
+        },
+        Some(move |builder: &mut QueryBuilder| {
+            // Collection (instance) filter: state and migration_id
+            builder.sql(" AND state = ANY(");
+            builder
+                .param()
+                .bind::<sql_types::Array<InstanceStateEnum>, _>(safe_states);
+            builder.sql(") AND migration_id IS NULL");
+        }),
+        Some(move |builder: &mut QueryBuilder| {
+            // Resource (external_ip) filter: state, parent_id, and kind
+            builder.sql(" AND state = ");
+            builder
+                .param()
+                .bind::<IpAttachStateEnum, _>(IpAttachState::Detached);
+            builder.sql(" AND kind = ");
+            builder.param().bind::<IpKindEnum, _>(kind);
+            builder.sql(" AND parent_id IS NULL");
+        }),
+    )
+}
+
 impl DataStore {
     /// Create an external IP address for source NAT for an instance.
     pub async fn allocate_instance_snat_ip(
@@ -430,39 +498,21 @@ impl DataStore {
     ) -> Result<Option<(ExternalIp, bool)>, Error> {
         use diesel::result::DatabaseErrorKind::UniqueViolation;
         use diesel::result::Error::DatabaseError;
-        use nexus_db_schema::schema::external_ip::dsl;
-        use nexus_db_schema::schema::external_ip::table;
-        use nexus_db_schema::schema::instance::dsl as inst_dsl;
-        use nexus_db_schema::schema::instance::table as inst_table;
 
-        let safe_states = if creating_instance {
-            &SAFE_TO_ATTACH_INSTANCE_STATES_CREATING[..]
-        } else {
-            &SAFE_TO_ATTACH_INSTANCE_STATES[..]
-        };
+        // Create the query template with all SQL structure
+        let template =
+            external_ip_attach_query_template(creating_instance, kind);
 
-        let query = Instance::attach_resource(
+        // Build the query with runtime parameter values
+        let query: AttachQuery<ExternalIp, Instance> = template.build(
             instance_id.into_untyped_uuid(),
             ip_id,
-            inst_table
-                .into_boxed()
-                .filter(inst_dsl::state.eq_any(safe_states))
-                .filter(inst_dsl::migration_id.is_null()),
-            table
-                .into_boxed()
-                .filter(dsl::state.eq(IpAttachState::Detached))
-                .filter(dsl::kind.eq(kind))
-                .filter(dsl::parent_id.is_null()),
             MAX_EXTERNAL_IPS_PLUS_SNAT,
-            diesel::update(dsl::external_ip).set((
-                dsl::parent_id.eq(Some(instance_id.into_untyped_uuid())),
-                dsl::time_modified.eq(Utc::now()),
-                dsl::state.eq(IpAttachState::Attaching),
-            )),
         );
 
+        let conn = self.pool_connection_authorized(opctx).await?;
         let mut do_saga = true;
-        query.attach_and_get_result_async(&*self.pool_connection_authorized(opctx).await?)
+        query.attach_and_get_result_async(&conn)
         .await
         .map(|(_, resource)| Some(resource))
         .or_else(|e: AttachError<ExternalIp, _, _>| match e {

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

@@ -9,8 +9,6 @@ use super::SQL_BATCH_SIZE;
 use crate::authz;
 use crate::context::OpContext;
 use crate::db;
-use crate::db::collection_insert::AsyncInsertError;
-use crate::db::collection_insert::DatastoreCollection;
 use crate::db::cte_utils::BoxedQuery;
 use crate::db::model::IncompleteNetworkInterface;
 use crate::db::model::Instance;
@@ -20,7 +18,6 @@ use crate::db::model::NetworkInterface;
 use crate::db::model::NetworkInterfaceKind;
 use crate::db::model::NetworkInterfaceUpdate;
 use crate::db::model::SqlU8;
-use crate::db::model::VpcSubnet;
 use crate::db::pagination::Paginator;
 use crate::db::pagination::paginated;
 use crate::db::queries::network_interface;
@@ -404,27 +401,10 @@ impl DataStore {
         conn: &async_bb8_diesel::Connection<DbConnection>,
         interface: IncompleteNetworkInterface,
     ) -> Result<NetworkInterface, network_interface::InsertError> {
-        use nexus_db_schema::schema::network_interface::dsl;
-        let subnet_id = interface.subnet.identity.id;
-        let query = network_interface::InsertQuery::new(interface.clone());
-        VpcSubnet::insert_resource(
-            subnet_id,
-            diesel::insert_into(dsl::network_interface).values(query),
-        )
-        .insert_and_get_result_async(conn)
-        .await
-        .map_err(|e| match e {
-            AsyncInsertError::CollectionNotFound => {
-                network_interface::InsertError::External(
-                    Error::ObjectNotFound {
-                        type_name: ResourceType::VpcSubnet,
-                        lookup_type: LookupType::ById(subnet_id),
-                    },
-                )
-            }
-            AsyncInsertError::DatabaseError(e) => {
-                network_interface::InsertError::from_diesel(e, &interface)
-            }
+        let query = network_interface::InsertQuery::new(interface.clone())
+            .to_insert_query();
+        query.get_result_async(conn).await.map_err(|e| {
+            network_interface::InsertError::from_diesel(e, &interface)
         })
     }