From c868f715ca284b6d3d61a1cfe09c82447ab67b02 Mon Sep 17 00:00:00 2001 From: Sean Klein Date: Wed, 14 Feb 2024 16:58:14 -0800 Subject: [PATCH] arcane chicanery which bends the type system to our will --- nexus/db-macros/src/lib.rs | 4 +- nexus/db-model/src/dataset.rs | 1 - nexus/db-model/src/region.rs | 1 - nexus/db-queries/src/db/datastore/region.rs | 21 +++------ .../src/db/queries/region_allocation.rs | 13 +++--- nexus/db-queries/src/db/raw_query_builder.rs | 44 ++++++++++++++++++- 6 files changed, 56 insertions(+), 28 deletions(-) diff --git a/nexus/db-macros/src/lib.rs b/nexus/db-macros/src/lib.rs index 30208412a1..fd15b59128 100644 --- a/nexus/db-macros/src/lib.rs +++ b/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, diff --git a/nexus/db-model/src/dataset.rs b/nexus/db-model/src/dataset.rs index 0ea465bebc..65c0070509 100644 --- a/nexus/db-model/src/dataset.rs +++ b/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, diff --git a/nexus/db-model/src/region.rs b/nexus/db-model/src/region.rs index 0ea6f82e65..fefc4f4fce 100644 --- a/nexus/db-model/src/region.rs +++ b/nexus/db-model/src/region.rs @@ -15,7 +15,6 @@ use uuid::Uuid; /// allocated within a volume. #[derive( Queryable, - QueryableByName, Insertable, Debug, Clone, diff --git a/nexus/db-queries/src/db/datastore/region.rs b/nexus/db-queries/src/db/datastore/region.rs index b334ac1085..c7b98b9a82 100644 --- a/nexus/db-queries/src/db/datastore/region.rs +++ b/nexus/db-queries/src/db/datastore/region.rs @@ -128,15 +128,7 @@ 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, @@ -144,14 +136,11 @@ impl DataStore { extent_count, allocation_strategy, ) - .get_results_async::( - &*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::>(); + .map_err(|e| { + crate::db::queries::region_allocation::from_diesel(e) + })?; info!( self.log, diff --git a/nexus/db-queries/src/db/queries/region_allocation.rs b/nexus/db-queries/src/db/queries/region_allocation.rs index 407f2df865..3d39808831 100644 --- a/nexus/db-queries/src/db/queries/region_allocation.rs +++ b/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 = < + >::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, SelectableSql)> { let (seed, distinct_sleds) = { let (input_seed, distinct_sleds) = match allocation_strategy { RegionAllocationStrategy::Random { seed } => (seed, false), diff --git a/nexus/db-queries/src/db/raw_query_builder.rs b/nexus/db-queries/src/db/raw_query_builder.rs index b0094f09ce..65ff22f86f 100644 --- a/nexus/db-queries/src/db/raw_query_builder.rs +++ b/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(self) -> TypedSqlQuery { + 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 { + inner: diesel::query_builder::BoxedSqlQuery< + 'static, + Pg, + diesel::query_builder::SqlQuery, + >, + _phantom: PhantomData, +} + +impl QueryFragment for TypedSqlQuery { + 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 RunQueryDsl for TypedSqlQuery {} + +impl Query for TypedSqlQuery { + type SqlType = T; +}