arcane chicanery which bends the type system to our will

nexus/db-macros/src/lib.rs

@@ -280,7 +280,7 @@ fn build_resource_identity(
     let identity_name = format_ident!("{}Identity", struct_name);
     quote! {
         #[doc = #identity_doc]
-        #[derive(Clone, Debug, PartialEq, Eq, Selectable, Queryable, QueryableByName, Insertable, serde::Serialize, serde::Deserialize)]
+        #[derive(Clone, Debug, PartialEq, Eq, Selectable, Queryable, Insertable, serde::Serialize, serde::Deserialize)]
         #[diesel(table_name = #table_name) ]
         pub struct #identity_name {
             pub id: ::uuid::Uuid,
@@ -322,7 +322,7 @@ fn build_asset_identity(
     let identity_name = format_ident!("{}Identity", struct_name);
     quote! {
         #[doc = #identity_doc]
-        #[derive(Clone, Debug, PartialEq, Selectable, Queryable, QueryableByName, Insertable, serde::Serialize, serde::Deserialize)]
+        #[derive(Clone, Debug, PartialEq, Selectable, Queryable, Insertable, serde::Serialize, serde::Deserialize)]
         #[diesel(table_name = #table_name) ]
         pub struct #identity_name {
             pub id: ::uuid::Uuid,

nexus/db-model/src/dataset.rs

@@ -18,7 +18,6 @@ use uuid::Uuid;
 /// available to a service on the Sled.
 #[derive(
     Queryable,
-    QueryableByName,
     Insertable,
     Debug,
     Clone,

nexus/db-model/src/region.rs

@@ -15,7 +15,6 @@ use uuid::Uuid;
 /// allocated within a volume.
 #[derive(
     Queryable,
-    QueryableByName,
     Insertable,
     Debug,
     Clone,

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

@@ -128,30 +128,19 @@ impl DataStore {
         let (blocks_per_extent, extent_count) =
             Self::get_crucible_allocation(&block_size, size);
 
-        #[derive(diesel::QueryableByName, Queryable)]
-        struct DatasetAndRegion {
-            #[diesel(embed)]
-            dataset: Dataset,
-            #[diesel(embed)]
-            region: Region,
-        }
-
-        let dataset_and_regions: Vec<_> =
+        let dataset_and_regions: Vec<(Dataset, Region)> =
             crate::db::queries::region_allocation::allocation_query(
                 volume_id,
                 block_size.to_bytes() as u64,
                 blocks_per_extent,
                 extent_count,
                 allocation_strategy,
             )
-            .get_results_async::<DatasetAndRegion>(
-                &*self.pool_connection_authorized(&opctx).await?,
-            )
+            .get_results_async(&*self.pool_connection_authorized(&opctx).await?)
             .await
-            .map_err(|e| crate::db::queries::region_allocation::from_diesel(e))?
-            .into_iter()
-            .map(|DatasetAndRegion { dataset, region }| (dataset, region))
-            .collect::<Vec<_>>();
+            .map_err(|e| {
+                crate::db::queries::region_allocation::from_diesel(e)
+            })?;
 
         info!(
             self.log,

nexus/db-queries/src/db/queries/region_allocation.rs

@@ -6,7 +6,8 @@
 
 use crate::db::column_walker::AllColumnsOf;
 use crate::db::datastore::REGION_REDUNDANCY_THRESHOLD;
-use crate::db::raw_query_builder::QueryBuilder;
+use crate::db::model::{Dataset, Region};
+use crate::db::raw_query_builder::{QueryBuilder, TypedSqlQuery};
 use crate::db::schema;
 use crate::db::true_or_cast_error::matches_sentinel;
 use const_format::concatcp;
@@ -63,17 +64,17 @@ pub fn from_diesel(e: DieselError) -> external::Error {
     error::public_error_from_diesel(e, error::ErrorHandler::Server)
 }
 
+type SelectableSql<T> = <
+    <T as diesel::Selectable<Pg>>::SelectExpression as diesel::Expression
+>::SqlType;
+
 pub fn allocation_query(
     volume_id: uuid::Uuid,
     block_size: u64,
     blocks_per_extent: u64,
     extent_count: u64,
     allocation_strategy: &RegionAllocationStrategy,
-) -> diesel::query_builder::BoxedSqlQuery<
-    'static,
-    Pg,
-    diesel::query_builder::SqlQuery,
-> {
+) -> TypedSqlQuery<(SelectableSql<Dataset>, SelectableSql<Region>)> {
     let (seed, distinct_sleds) = {
         let (input_seed, distinct_sleds) = match allocation_strategy {
             RegionAllocationStrategy::Random { seed } => (seed, false),

nexus/db-queries/src/db/raw_query_builder.rs

@@ -7,9 +7,13 @@
 //! These largely side-step Diesel's type system,
 //! and are recommended for more complex CTE
 
+use crate::db::pool::DbConnection;
 use diesel::pg::Pg;
+use diesel::query_builder::{AstPass, Query, QueryFragment, QueryId};
 use diesel::sql_types;
+use diesel::RunQueryDsl;
 use std::cell::Cell;
+use std::marker::PhantomData;
 
 // Keeps a counter to "how many bind parameters have been used" to
 // aid in the construction of the query string.
@@ -149,7 +153,43 @@ impl QueryBuilder {
     }
 
     /// Takes the final boxed query
-    pub fn query(self) -> BoxedQuery {
-        self.query
+    pub fn query<T>(self) -> TypedSqlQuery<T> {
+        TypedSqlQuery { inner: self.query, _phantom: PhantomData }
     }
 }
+
+/// Diesel's [diesel::query_builder::BoxedSqlQuery] has a few drawbacks that
+/// make wrapper more palatable:
+///
+/// - It always implements "Query" with SqlType = Untyped, so a caller could try to
+/// execute this query and get back any type.
+/// - It forces the usage of "QueryableByName", which acts wrong if we're
+/// returning multiple columns with the same name (this is normal! If you want
+/// to UNION two objects that both have "id" columns, this happens).
+#[derive(QueryId)]
+pub struct TypedSqlQuery<T> {
+    inner: diesel::query_builder::BoxedSqlQuery<
+        'static,
+        Pg,
+        diesel::query_builder::SqlQuery,
+    >,
+    _phantom: PhantomData<T>,
+}
+
+impl<T> QueryFragment<Pg> for TypedSqlQuery<T> {
+    fn walk_ast<'a>(
+        &'a self,
+        mut out: AstPass<'_, 'a, Pg>,
+    ) -> diesel::QueryResult<()> {
+        out.unsafe_to_cache_prepared();
+
+        self.inner.walk_ast(out.reborrow())?;
+        Ok(())
+    }
+}
+
+impl<T> RunQueryDsl<DbConnection> for TypedSqlQuery<T> {}
+
+impl<T> Query for TypedSqlQuery<T> {
+    type SqlType = T;
+}