diff --git a/nexus/db-model/src/lib.rs b/nexus/db-model/src/lib.rs index c57836a567..51fd0f6c9e 100644 --- a/nexus/db-model/src/lib.rs +++ b/nexus/db-model/src/lib.rs @@ -62,7 +62,6 @@ mod v2p_mapping; mod deployment; mod ipv4_nat_entry; mod omicron_zone_config; -pub mod queries; mod quota; mod rack; mod region; diff --git a/nexus/db-model/src/queries/mod.rs b/nexus/db-model/src/queries/mod.rs deleted file mode 100644 index e138508f84..0000000000 --- a/nexus/db-model/src/queries/mod.rs +++ /dev/null @@ -1,7 +0,0 @@ -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at https://mozilla.org/MPL/2.0/. - -//! Subqueries used in CTEs. - -pub mod virtual_provisioning_collection_update; diff --git a/nexus/db-model/src/queries/virtual_provisioning_collection_update.rs b/nexus/db-model/src/queries/virtual_provisioning_collection_update.rs deleted file mode 100644 index 124ffe4db6..0000000000 --- a/nexus/db-model/src/queries/virtual_provisioning_collection_update.rs +++ /dev/null @@ -1,60 +0,0 @@ -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at https://mozilla.org/MPL/2.0/. - -//! Describes the resource provisioning update CTE -//! -//! Refer to -//! for the construction of this query. - -use crate::schema::silo; -use crate::schema::silo_quotas; -use crate::schema::virtual_provisioning_collection; - -table! { - parent_silo { - id -> Uuid, - } -} - -table! { - all_collections { - id -> Uuid, - } -} - -table! { - do_update (update) { - update -> Bool, - } -} - -table! { - quotas (silo_id) { - silo_id -> Uuid, - cpus -> Int8, - memory -> Int8, - storage -> Int8, - } -} - -table! { - silo_provisioned { - id -> Uuid, - virtual_disk_bytes_provisioned -> Int8, - cpus_provisioned -> Int8, - ram_provisioned -> Int8, - } -} - -diesel::allow_tables_to_appear_in_same_query!(silo, parent_silo,); - -diesel::allow_tables_to_appear_in_same_query!( - virtual_provisioning_collection, - silo_quotas, - parent_silo, - all_collections, - do_update, - quotas, - silo_provisioned -); diff --git a/nexus/db-queries/src/db/alias.rs b/nexus/db-queries/src/db/alias.rs deleted file mode 100644 index 0a5bcca743..0000000000 --- a/nexus/db-queries/src/db/alias.rs +++ /dev/null @@ -1,84 +0,0 @@ -// This Source Code Form is subject to the terms of the Mozilla Public -// License, v. 2.0. If a copy of the MPL was not distributed with this -// file, You can obtain one at https://mozilla.org/MPL/2.0/. - -//! Tools for creating aliases in diesel. - -use diesel::pg::Pg; -use diesel::query_builder::AstPass; -use diesel::query_builder::QueryFragment; -use diesel::Expression; -use diesel::SelectableExpression; - -/// Allows an [`diesel::Expression`] to be referenced by a new name. -/// -/// This generates an ` AS ` SQL fragment. -/// -/// -/// For example: -/// -/// ```ignore -/// diesel::sql_function!(fn gen_random_uuid() -> Uuid); -/// -/// let query = sleds.select( -/// ( -/// ExpressionAlias::(gen_random_uuid()), -/// ExpressionAlias::(gen_random_uuid()), -/// ), -/// ); -/// ``` -/// -/// Produces the following SQL: -/// -/// ```sql -/// SELECT -/// gen_random_uuid() as id, -/// gen_random_uuid() as sled_id, -/// FROM sleds -/// ``` -#[derive(diesel::expression::ValidGrouping, diesel::query_builder::QueryId)] -pub struct ExpressionAlias { - expr: E, - name: &'static str, -} - -impl ExpressionAlias -where - E: Expression, -{ - pub fn new(expr: E) -> Self { - Self { expr, name: C::NAME } - } -} - -impl Expression for ExpressionAlias -where - E: Expression, -{ - type SqlType = E::SqlType; -} - -impl diesel::AppearsOnTable for ExpressionAlias where - E: diesel::AppearsOnTable -{ -} - -impl SelectableExpression for ExpressionAlias where - E: SelectableExpression -{ -} - -impl QueryFragment for ExpressionAlias -where - E: QueryFragment, -{ - fn walk_ast<'a>( - &'a self, - mut out: AstPass<'_, 'a, Pg>, - ) -> diesel::QueryResult<()> { - self.expr.walk_ast(out.reborrow())?; - out.push_sql(" AS "); - out.push_sql(&self.name); - Ok(()) - } -} diff --git a/nexus/db-queries/src/db/mod.rs b/nexus/db-queries/src/db/mod.rs index 9b3d71970c..7ce6890a4d 100644 --- a/nexus/db-queries/src/db/mod.rs +++ b/nexus/db-queries/src/db/mod.rs @@ -4,7 +4,6 @@ //! Facilities for working with the Omicron database -pub(crate) mod alias; // This is not intended to be public, but this is necessary to use it from // doctests pub mod collection_attach; diff --git a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs index 09798e4e5d..156691866e 100644 --- a/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs +++ b/nexus/db-queries/src/db/queries/virtual_provisioning_collection_update.rs @@ -4,32 +4,27 @@ //! Implementation of queries for updating resource provisioning info. -use crate::db::alias::ExpressionAlias; +use crate::db::column_walker::AllColumnsOf; use crate::db::model::ByteCount; use crate::db::model::ResourceTypeProvisioned; use crate::db::model::VirtualProvisioningCollection; use crate::db::model::VirtualProvisioningResource; -use crate::db::pool::DbConnection; +use crate::db::raw_query_builder::{QueryBuilder, TypedSqlQuery}; use crate::db::schema::virtual_provisioning_collection; use crate::db::schema::virtual_provisioning_resource; -use crate::db::subquery::{AsQuerySource, Cte, CteBuilder, CteQuery}; use crate::db::true_or_cast_error::matches_sentinel; -use crate::db::true_or_cast_error::TrueOrCastError; -use db_macros::Subquery; +use const_format::concatcp; use diesel::pg::Pg; -use diesel::query_builder::{AstPass, Query, QueryFragment, QueryId}; use diesel::result::Error as DieselError; -use diesel::{ - sql_types, BoolExpressionMethods, CombineDsl, ExpressionMethods, IntoSql, - JoinOnDsl, NullableExpressionMethods, QueryDsl, RunQueryDsl, - SelectableHelper, -}; -use nexus_db_model::queries::virtual_provisioning_collection_update::{ - all_collections, do_update, parent_silo, quotas, silo_provisioned, -}; +use diesel::sql_types; use omicron_common::api::external; use omicron_common::api::external::MessagePair; +type AllColumnsOfVirtualResource = + AllColumnsOf; +type AllColumnsOfVirtualCollection = + AllColumnsOf; + const NOT_ENOUGH_CPUS_SENTINEL: &'static str = "Not enough cpus"; const NOT_ENOUGH_MEMORY_SENTINEL: &'static str = "Not enough memory"; const NOT_ENOUGH_STORAGE_SENTINEL: &'static str = "Not enough storage"; @@ -77,319 +72,33 @@ pub fn from_diesel(e: DieselError) -> external::Error { error::public_error_from_diesel(e, error::ErrorHandler::Server) } -#[derive(Subquery, QueryId)] -#[subquery(name = parent_silo)] -struct ParentSilo { - query: Box>, -} - -impl ParentSilo { - fn new(project_id: uuid::Uuid) -> Self { - use crate::db::schema::project::dsl; - Self { - query: Box::new( - dsl::project.filter(dsl::id.eq(project_id)).select(( - ExpressionAlias::new::(dsl::silo_id), - )), - ), - } - } -} - -#[derive(Subquery, QueryId)] -#[subquery(name = all_collections)] -struct AllCollections { - query: Box>, -} - -impl AllCollections { - fn new( - project_id: uuid::Uuid, - parent_silo: &ParentSilo, - fleet_id: uuid::Uuid, - ) -> Self { - let project_id = project_id.into_sql::(); - let fleet_id = fleet_id.into_sql::(); - Self { - query: Box::new( - diesel::select((ExpressionAlias::new::< - all_collections::dsl::id, - >(project_id),)) - .union(parent_silo.query_source().select(( - ExpressionAlias::new::( - parent_silo::id, - ), - ))) - .union(diesel::select((ExpressionAlias::new::< - all_collections::dsl::id, - >(fleet_id),))), - ), - } - } -} - -#[derive(Subquery, QueryId)] -#[subquery(name = do_update)] -struct DoUpdate { - query: Box>, -} - -impl DoUpdate { - fn new_for_insert( - silo_provisioned: &SiloProvisioned, - quotas: &Quotas, - resource: VirtualProvisioningResource, - ) -> Self { - use virtual_provisioning_resource::dsl; - - let cpus_provisioned_delta = - resource.cpus_provisioned.into_sql::(); - let memory_provisioned_delta = - i64::from(resource.ram_provisioned).into_sql::(); - let storage_provisioned_delta = - i64::from(resource.virtual_disk_bytes_provisioned) - .into_sql::(); - - let not_allocted = dsl::virtual_provisioning_resource - .find(resource.id) - .count() - .single_value() - .assume_not_null() - .eq(0); - - let has_sufficient_cpus = quotas - .query_source() - .select(quotas::cpus) - .single_value() - .assume_not_null() - .ge(silo_provisioned - .query_source() - .select(silo_provisioned::cpus_provisioned) - .single_value() - .assume_not_null() - + cpus_provisioned_delta); - - let has_sufficient_memory = quotas - .query_source() - .select(quotas::memory) - .single_value() - .assume_not_null() - .ge(silo_provisioned - .query_source() - .select(silo_provisioned::ram_provisioned) - .single_value() - .assume_not_null() - + memory_provisioned_delta); - - let has_sufficient_storage = quotas - .query_source() - .select(quotas::storage) - .single_value() - .assume_not_null() - .ge(silo_provisioned - .query_source() - .select(silo_provisioned::virtual_disk_bytes_provisioned) - .single_value() - .assume_not_null() - + storage_provisioned_delta); - - Self { - query: Box::new(diesel::select((ExpressionAlias::new::< - do_update::update, - >( - not_allocted - .and(TrueOrCastError::new( - cpus_provisioned_delta.eq(0).or(has_sufficient_cpus), - NOT_ENOUGH_CPUS_SENTINEL, - )) - .and(TrueOrCastError::new( - memory_provisioned_delta - .eq(0) - .or(has_sufficient_memory), - NOT_ENOUGH_MEMORY_SENTINEL, - )) - .and(TrueOrCastError::new( - storage_provisioned_delta - .eq(0) - .or(has_sufficient_storage), - NOT_ENOUGH_STORAGE_SENTINEL, - )), - ),))), - } - } - - fn new_for_delete(id: uuid::Uuid) -> Self { - use virtual_provisioning_resource::dsl; - - let already_allocated = dsl::virtual_provisioning_resource - .find(id) - .count() - .single_value() - .assume_not_null() - .eq(1); - - Self { - query: Box::new(diesel::select((ExpressionAlias::new::< - do_update::update, - >(already_allocated),))), - } - } -} - -#[derive(Subquery, QueryId)] -#[subquery(name = virtual_provisioning_collection)] -struct UpdatedProvisions { - query: - Box>, -} - -impl UpdatedProvisions { - fn new( - all_collections: &AllCollections, - do_update: &DoUpdate, - values: V, - ) -> Self - where - V: diesel::AsChangeset, - ::Changeset: - QueryFragment + Send + 'static, - { - use virtual_provisioning_collection::dsl; - - Self { - query: Box::new( - diesel::update(dsl::virtual_provisioning_collection) - .set(values) - .filter( - dsl::id.eq_any( - all_collections - .query_source() - .select(all_collections::id), - ), - ) - .filter( - do_update - .query_source() - .select(do_update::update) - .single_value() - .assume_not_null(), - ) - .returning(virtual_provisioning_collection::all_columns), - ), - } - } -} - -#[derive(Subquery, QueryId)] -#[subquery(name = quotas)] -struct Quotas { - query: Box>, -} - -impl Quotas { - // TODO: We could potentially skip this in cases where we know we're removing a resource instead of inserting - fn new(parent_silo: &ParentSilo) -> Self { - use crate::db::schema::silo_quotas::dsl; - Self { - query: Box::new( - dsl::silo_quotas - .inner_join( - parent_silo - .query_source() - .on(dsl::silo_id.eq(parent_silo::id)), - ) - .select(( - dsl::silo_id, - dsl::cpus, - ExpressionAlias::new::( - dsl::memory_bytes, - ), - ExpressionAlias::new::( - dsl::storage_bytes, - ), - )), - ), - } - } -} - -#[derive(Subquery, QueryId)] -#[subquery(name = silo_provisioned)] -struct SiloProvisioned { - query: Box>, -} - -impl SiloProvisioned { - fn new(parent_silo: &ParentSilo) -> Self { - use virtual_provisioning_collection::dsl; - Self { - query: Box::new( - dsl::virtual_provisioning_collection - .inner_join( - parent_silo - .query_source() - .on(dsl::id.eq(parent_silo::id)), - ) - .select(( - dsl::id, - dsl::cpus_provisioned, - dsl::ram_provisioned, - dsl::virtual_disk_bytes_provisioned, - )), - ), - } - } -} - -// This structure wraps a query, such that it can be used within a CTE. -// -// It generates a name that can be used by the "CteBuilder", but does not -// implement "AsQuerySource". This basically means: -// - It can be used to add data-modifying statements to the CTE -// - The result of the query cannot be referenced by subsequent queries -// -// NOTE: The name for each CTE arm should be unique, so this shouldn't be used -// multiple times within a single CTE. This restriction could be removed by -// generating unique identifiers. -struct UnreferenceableSubquery(Q); - -impl QueryFragment for UnreferenceableSubquery -where - Q: QueryFragment + Send + 'static, -{ - fn walk_ast<'a>( - &'a self, - mut out: diesel::query_builder::AstPass<'_, 'a, Pg>, - ) -> diesel::QueryResult<()> { - out.push_identifier("unused_cte_arm")?; - Ok(()) - } -} - -impl crate::db::subquery::Subquery for UnreferenceableSubquery -where - Q: QueryFragment + Send + 'static, -{ - fn query(&self) -> &dyn QueryFragment { - &self.0 - } -} - /// The virtual resource collection is only updated when a resource is inserted /// or deleted from the resource provisioning table. By probing for the presence /// or absence of a resource, we can update collections at the same time as we /// create or destroy the resource, which helps make the operation idempotent. +#[derive(Clone)] enum UpdateKind { - Insert(VirtualProvisioningResource), - Delete(uuid::Uuid), + InsertStorage(VirtualProvisioningResource), + DeleteStorage { + id: uuid::Uuid, + disk_byte_diff: ByteCount, + }, + InsertInstance(VirtualProvisioningResource), + DeleteInstance { + id: uuid::Uuid, + max_instance_gen: i64, + cpus_diff: i64, + ram_diff: ByteCount, + }, } +type SelectableSql = < + >::SelectExpression as diesel::Expression +>::SqlType; + /// Constructs a CTE for updating resource provisioning information in all /// collections for a particular object. -#[derive(QueryId)] -pub struct VirtualProvisioningCollectionUpdate { - cte: Cte, -} +pub struct VirtualProvisioningCollectionUpdate {} impl VirtualProvisioningCollectionUpdate { // Generic utility for updating all collections including this resource, @@ -399,66 +108,336 @@ impl VirtualProvisioningCollectionUpdate { // - Project // - Silo // - Fleet - // - // Arguments: - // - do_update: A boolean SQL query to answer the question: "Should this update - // be applied"? This query is necessary for idempotency. - // - update: A SQL query to actually modify the resource record. Generally - // this is an "INSERT", "UPDATE", or "DELETE". - // - project_id: The project to which the resource belongs. - // - values: The updated values to propagate through collections (iff - // "do_update" evaluates to "true"). - fn apply_update( + fn apply_update( update_kind: UpdateKind, - update: U, project_id: uuid::Uuid, - values: V, - ) -> Self - where - U: QueryFragment + crate::db::subquery::Subquery + Send + 'static, - V: diesel::AsChangeset, - ::Changeset: - QueryFragment + Send + 'static, - { - let parent_silo = ParentSilo::new(project_id); - let all_collections = AllCollections::new( - project_id, - &parent_silo, - *crate::db::fixed_data::FLEET_ID, - ); - - let quotas = Quotas::new(&parent_silo); - let silo_provisioned = SiloProvisioned::new(&parent_silo); + ) -> TypedSqlQuery> { + let query = QueryBuilder::new().sql(" +WITH + parent_silo AS (SELECT project.silo_id AS id FROM project WHERE project.id = ").param().sql("),") + .bind::(project_id).sql(" + all_collections + AS ( + ((SELECT ").param().sql(" AS id) UNION (SELECT parent_silo.id AS id FROM parent_silo)) + UNION (SELECT ").param().sql(" AS id) + ),") + .bind::(project_id) + .bind::(*crate::db::fixed_data::FLEET_ID) + .sql(" + quotas + AS ( + SELECT + silo_quotas.silo_id, + silo_quotas.cpus, + silo_quotas.memory_bytes AS memory, + silo_quotas.storage_bytes AS storage + FROM + silo_quotas INNER JOIN parent_silo ON silo_quotas.silo_id = parent_silo.id + ), + silo_provisioned + AS ( + SELECT + virtual_provisioning_collection.id, + virtual_provisioning_collection.cpus_provisioned, + virtual_provisioning_collection.ram_provisioned, + virtual_provisioning_collection.virtual_disk_bytes_provisioned + FROM + virtual_provisioning_collection + INNER JOIN parent_silo ON virtual_provisioning_collection.id = parent_silo.id + ),"); + + let query = match update_kind.clone() { + UpdateKind::InsertInstance(resource) | UpdateKind::InsertStorage(resource) => { + query.sql(" + do_update + AS ( + SELECT + ( + ( + ( + SELECT count(*) + FROM virtual_provisioning_resource + WHERE virtual_provisioning_resource.id = ").param().sql(" + LIMIT 1 + ) + = 0 + AND CAST( + IF( + ( + ").param().sql(" = 0 + OR (SELECT quotas.cpus FROM quotas LIMIT 1) + >= ( + (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT 1) + + ").param().sql(concatcp!(" + ) + ), + 'TRUE', + '", NOT_ENOUGH_CPUS_SENTINEL, "' + ) + AS BOOL + ) + ) + AND CAST( + IF( + ( + ")).param().sql(" = 0 + OR (SELECT quotas.memory FROM quotas LIMIT 1) + >= ( + (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT 1) + + ").param().sql(concatcp!(" + ) + ), + 'TRUE', + '", NOT_ENOUGH_MEMORY_SENTINEL, "' + ) + AS BOOL + ) + ) + AND CAST( + IF( + ( + ")).param().sql(" = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) + >= ( + ( + SELECT + silo_provisioned.virtual_disk_bytes_provisioned + FROM + silo_provisioned + LIMIT + 1 + ) + + ").param().sql(concatcp!(" + ) + ), + 'TRUE', + '", NOT_ENOUGH_STORAGE_SENTINEL, "' + ) + AS BOOL + ) + AS update + ),")) + .bind::(resource.id) + .bind::(resource.cpus_provisioned) + .bind::(resource.cpus_provisioned) + .bind::(resource.ram_provisioned) + .bind::(resource.ram_provisioned) + .bind::(resource.virtual_disk_bytes_provisioned) + .bind::(resource.virtual_disk_bytes_provisioned) + }, + UpdateKind::DeleteInstance { id, .. } | UpdateKind::DeleteStorage { id, .. } => { + query.sql(" + do_update + AS ( + SELECT + ( + SELECT + count(*) + FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = ").param().sql(" + LIMIT + 1 + ) + = 1 + AS update + ),") + .bind::(id) + }, + }; - let do_update = match update_kind { - UpdateKind::Insert(resource) => { - DoUpdate::new_for_insert(&silo_provisioned, "as, resource) + let query = match update_kind.clone() { + UpdateKind::InsertInstance(resource) + | UpdateKind::InsertStorage(resource) => query + .sql( + " + unused_cte_arm + AS ( + INSERT + INTO + virtual_provisioning_resource + ( + id, + time_modified, + resource_type, + virtual_disk_bytes_provisioned, + cpus_provisioned, + ram_provisioned + ) + VALUES + (", + ) + .param() + .sql(", DEFAULT, ") + .param() + .sql(", ") + .param() + .sql(", ") + .param() + .sql(", ") + .param() + .sql( + ") + ON CONFLICT + DO + NOTHING + RETURNING ", + ) + .sql(AllColumnsOfVirtualResource::with_prefix( + "virtual_provisioning_resource", + )) + .sql("),") + .bind::(resource.id) + .bind::(resource.resource_type) + .bind::( + resource.virtual_disk_bytes_provisioned, + ) + .bind::(resource.cpus_provisioned) + .bind::(resource.ram_provisioned), + UpdateKind::DeleteInstance { id, max_instance_gen, .. } => { + // The filter condition here ensures that the provisioning record is + // only deleted if the corresponding instance has a generation + // number less than the supplied `max_instance_gen`. This allows a + // caller that is about to apply an instance update that will stop + // the instance and that bears generation G to avoid deleting + // resources if the instance generation was already advanced to or + // past G. + // + // If the relevant instance ID is not in the database, then some + // other operation must have ensured the instance was previously + // stopped (because that's the only way it could have been deleted), + // and that operation should have cleaned up the resources already, + // in which case there's nothing to do here. + query + .sql( + " + unused_cte_arm + AS ( + DELETE FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = ", + ) + .param() + .sql( + " + AND + virtual_provisioning_resource.id = ( + SELECT instance.id FROM instance WHERE + instance.id = ", + ) + .param() + .sql( + " AND + instance.state_generation < ", + ) + .param() + .sql( + " LIMIT 1) + RETURNING ", + ) + .sql(AllColumnsOfVirtualResource::with_prefix( + "virtual_provisioning_resource", + )) + .sql("),") + .bind::(id) + .bind::(id) + .bind::(max_instance_gen) } - UpdateKind::Delete(id) => DoUpdate::new_for_delete(id), + UpdateKind::DeleteStorage { id, .. } => query + .sql( + " + unused_cte_arm + AS ( + DELETE FROM + virtual_provisioning_resource + WHERE + virtual_provisioning_resource.id = ", + ) + .param() + .sql( + " + RETURNING ", + ) + .sql(AllColumnsOfVirtualResource::with_prefix( + "virtual_provisioning_resource", + )) + .sql("),") + .bind::(id), }; - let updated_collections = - UpdatedProvisions::new(&all_collections, &do_update, values); - - // TODO: Do we want to select from "all_collections" instead? Seems more - // idempotent; it'll work even when we don't update anything... - let final_select = Box::new( - updated_collections - .query_source() - .select(VirtualProvisioningCollection::as_select()), + let query = query.sql( + " + virtual_provisioning_collection + AS ( + UPDATE + virtual_provisioning_collection + SET", ); + let query = match update_kind.clone() { + UpdateKind::InsertInstance(resource) => query + .sql( + " + time_modified = current_timestamp(), + cpus_provisioned = virtual_provisioning_collection.cpus_provisioned + ", + ) + .param() + .sql( + ", + ram_provisioned = virtual_provisioning_collection.ram_provisioned + ", + ) + .param() + .bind::(resource.cpus_provisioned) + .bind::(resource.ram_provisioned), + UpdateKind::InsertStorage(resource) => query + .sql( + " + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned + ", + ) + .param() + .bind::( + resource.virtual_disk_bytes_provisioned, + ), + UpdateKind::DeleteInstance { cpus_diff, ram_diff, .. } => query + .sql( + " + time_modified = current_timestamp(), + cpus_provisioned = virtual_provisioning_collection.cpus_provisioned - ", + ) + .param() + .sql( + ", + ram_provisioned = virtual_provisioning_collection.ram_provisioned - ", + ) + .param() + .bind::(cpus_diff) + .bind::(ram_diff), + UpdateKind::DeleteStorage { disk_byte_diff, .. } => query + .sql( + " + time_modified = current_timestamp(), + virtual_disk_bytes_provisioned + = virtual_provisioning_collection.virtual_disk_bytes_provisioned - ", + ) + .param() + .bind::(disk_byte_diff), + }; - let cte = CteBuilder::new() - .add_subquery(parent_silo) - .add_subquery(all_collections) - .add_subquery(quotas) - .add_subquery(silo_provisioned) - .add_subquery(do_update) - .add_subquery(update) - .add_subquery(updated_collections) - .build(final_select); - - Self { cte } + query.sql(" + WHERE + virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) + AND (SELECT do_update.update FROM do_update LIMIT 1) + RETURNING " + ).sql(AllColumnsOfVirtualCollection::with_prefix("virtual_provisioning_collection")).sql(" + ) +SELECT " + ).sql(AllColumnsOfVirtualCollection::with_prefix("virtual_provisioning_collection")).sql(" +FROM + virtual_provisioning_collection +").query() } pub fn new_insert_storage( @@ -466,62 +445,22 @@ impl VirtualProvisioningCollectionUpdate { disk_byte_diff: ByteCount, project_id: uuid::Uuid, storage_type: crate::db::datastore::StorageType, - ) -> Self { - use virtual_provisioning_collection::dsl as collection_dsl; - use virtual_provisioning_resource::dsl as resource_dsl; - + ) -> TypedSqlQuery> { let mut provision = VirtualProvisioningResource::new(id, storage_type.into()); provision.virtual_disk_bytes_provisioned = disk_byte_diff; - Self::apply_update( - UpdateKind::Insert(provision.clone()), - // The query to actually insert the record. - UnreferenceableSubquery( - diesel::insert_into( - resource_dsl::virtual_provisioning_resource, - ) - .values(provision) - .on_conflict_do_nothing() - .returning(virtual_provisioning_resource::all_columns), - ), - // Within this project, silo, fleet... - project_id, - // ... We add the disk usage. - ( - collection_dsl::time_modified.eq(diesel::dsl::now), - collection_dsl::virtual_disk_bytes_provisioned - .eq(collection_dsl::virtual_disk_bytes_provisioned - + disk_byte_diff), - ), - ) + Self::apply_update(UpdateKind::InsertStorage(provision), project_id) } pub fn new_delete_storage( id: uuid::Uuid, disk_byte_diff: ByteCount, project_id: uuid::Uuid, - ) -> Self { - use virtual_provisioning_collection::dsl as collection_dsl; - use virtual_provisioning_resource::dsl as resource_dsl; - + ) -> TypedSqlQuery> { Self::apply_update( - UpdateKind::Delete(id), - // The query to actually delete the record. - UnreferenceableSubquery( - diesel::delete(resource_dsl::virtual_provisioning_resource) - .filter(resource_dsl::id.eq(id)) - .returning(virtual_provisioning_resource::all_columns), - ), - // Within this project, silo, fleet... + UpdateKind::DeleteStorage { id, disk_byte_diff }, project_id, - // ... We subtract the disk usage. - ( - collection_dsl::time_modified.eq(diesel::dsl::now), - collection_dsl::virtual_disk_bytes_provisioned - .eq(collection_dsl::virtual_disk_bytes_provisioned - - disk_byte_diff), - ), ) } @@ -530,10 +469,7 @@ impl VirtualProvisioningCollectionUpdate { cpus_diff: i64, ram_diff: ByteCount, project_id: uuid::Uuid, - ) -> Self { - use virtual_provisioning_collection::dsl as collection_dsl; - use virtual_provisioning_resource::dsl as resource_dsl; - + ) -> TypedSqlQuery> { let mut provision = VirtualProvisioningResource::new( id, ResourceTypeProvisioned::Instance, @@ -541,28 +477,7 @@ impl VirtualProvisioningCollectionUpdate { provision.cpus_provisioned = cpus_diff; provision.ram_provisioned = ram_diff; - Self::apply_update( - UpdateKind::Insert(provision.clone()), - // The query to actually insert the record. - UnreferenceableSubquery( - diesel::insert_into( - resource_dsl::virtual_provisioning_resource, - ) - .values(provision) - .on_conflict_do_nothing() - .returning(virtual_provisioning_resource::all_columns), - ), - // Within this project, silo, fleet... - project_id, - // ... We update the resource usage. - ( - collection_dsl::time_modified.eq(diesel::dsl::now), - collection_dsl::cpus_provisioned - .eq(collection_dsl::cpus_provisioned + cpus_diff), - collection_dsl::ram_provisioned - .eq(collection_dsl::ram_provisioned + ram_diff), - ), - ) + Self::apply_update(UpdateKind::InsertInstance(provision), project_id) } pub fn new_delete_instance( @@ -571,86 +486,26 @@ impl VirtualProvisioningCollectionUpdate { cpus_diff: i64, ram_diff: ByteCount, project_id: uuid::Uuid, - ) -> Self { - use crate::db::schema::instance::dsl as instance_dsl; - use virtual_provisioning_collection::dsl as collection_dsl; - use virtual_provisioning_resource::dsl as resource_dsl; - + ) -> TypedSqlQuery> { Self::apply_update( - UpdateKind::Delete(id), - // The query to actually delete the record. - // - // The filter condition here ensures that the provisioning record is - // only deleted if the corresponding instance has a generation - // number less than the supplied `max_instance_gen`. This allows a - // caller that is about to apply an instance update that will stop - // the instance and that bears generation G to avoid deleting - // resources if the instance generation was already advanced to or - // past G. - // - // If the relevant instance ID is not in the database, then some - // other operation must have ensured the instance was previously - // stopped (because that's the only way it could have been deleted), - // and that operation should have cleaned up the resources already, - // in which case there's nothing to do here. - // - // There is an additional "direct" filter on the target resource ID - // to avoid a full scan of the resource table. - UnreferenceableSubquery( - diesel::delete(resource_dsl::virtual_provisioning_resource) - .filter(resource_dsl::id.eq(id)) - .filter( - resource_dsl::id.nullable().eq(instance_dsl::instance - .filter(instance_dsl::id.eq(id)) - .filter( - instance_dsl::state_generation - .lt(max_instance_gen), - ) - .select(instance_dsl::id) - .single_value()), - ) - .returning(virtual_provisioning_resource::all_columns), - ), - // Within this project, silo, fleet... + UpdateKind::DeleteInstance { + id, + max_instance_gen, + cpus_diff, + ram_diff, + }, project_id, - // ... We update the resource usage. - ( - collection_dsl::time_modified.eq(diesel::dsl::now), - collection_dsl::cpus_provisioned - .eq(collection_dsl::cpus_provisioned - cpus_diff), - collection_dsl::ram_provisioned - .eq(collection_dsl::ram_provisioned - ram_diff), - ), ) } } -impl QueryFragment for VirtualProvisioningCollectionUpdate { - fn walk_ast<'a>( - &'a self, - mut out: AstPass<'_, 'a, Pg>, - ) -> diesel::QueryResult<()> { - out.unsafe_to_cache_prepared(); - - self.cte.walk_ast(out.reborrow())?; - Ok(()) - } -} - -type SelectableSql = < - >::SelectExpression as diesel::Expression ->::SqlType; - -impl Query for VirtualProvisioningCollectionUpdate { - type SqlType = SelectableSql; -} - -impl RunQueryDsl for VirtualProvisioningCollectionUpdate {} - #[cfg(test)] mod test { use super::*; + use crate::db::explain::ExplainableAsync; use crate::db::raw_query_builder::expectorate_query_contents; + use nexus_test_utils::db::test_setup_database; + use omicron_test_utils::dev; use uuid::Uuid; // These tests are a bit of a "change detector", but they're here to help @@ -732,4 +587,120 @@ mod test { "tests/output/virtual_provisioning_collection_update_delete_instance.sql", ).await; } + + // Explain the possible forms of the SQL query to ensure that it + // creates a valid SQL string. + + #[tokio::test] + async fn explain_insert_storage() { + let logctx = dev::test_setup_log("explain_insert_storage"); + let log = logctx.log.new(o!()); + let mut db = test_setup_database(&log).await; + let cfg = crate::db::Config { url: db.pg_config().clone() }; + let pool = crate::db::Pool::new(&logctx.log, &cfg); + let conn = pool.pool().get().await.unwrap(); + + let id = Uuid::nil(); + let project_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + let storage_type = crate::db::datastore::StorageType::Disk; + + let query = VirtualProvisioningCollectionUpdate::new_insert_storage( + id, + disk_byte_diff, + project_id, + storage_type, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.cleanup().await.unwrap(); + logctx.cleanup_successful(); + } + + #[tokio::test] + async fn explain_delete_storage() { + let logctx = dev::test_setup_log("explain_delete_storage"); + let log = logctx.log.new(o!()); + let mut db = test_setup_database(&log).await; + let cfg = crate::db::Config { url: db.pg_config().clone() }; + let pool = crate::db::Pool::new(&logctx.log, &cfg); + let conn = pool.pool().get().await.unwrap(); + + let id = Uuid::nil(); + let project_id = Uuid::nil(); + let disk_byte_diff = 2048.try_into().unwrap(); + + let query = VirtualProvisioningCollectionUpdate::new_delete_storage( + id, + disk_byte_diff, + project_id, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.cleanup().await.unwrap(); + logctx.cleanup_successful(); + } + + #[tokio::test] + async fn explain_insert_instance() { + let logctx = dev::test_setup_log("explain_insert_instance"); + let log = logctx.log.new(o!()); + let mut db = test_setup_database(&log).await; + let cfg = crate::db::Config { url: db.pg_config().clone() }; + let pool = crate::db::Pool::new(&logctx.log, &cfg); + let conn = pool.pool().get().await.unwrap(); + + let id = Uuid::nil(); + let project_id = Uuid::nil(); + let cpus_diff = 16.try_into().unwrap(); + let ram_diff = 2048.try_into().unwrap(); + + let query = VirtualProvisioningCollectionUpdate::new_insert_instance( + id, cpus_diff, ram_diff, project_id, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.cleanup().await.unwrap(); + logctx.cleanup_successful(); + } + + #[tokio::test] + async fn explain_delete_instance() { + let logctx = dev::test_setup_log("explain_delete_instance"); + let log = logctx.log.new(o!()); + let mut db = test_setup_database(&log).await; + let cfg = crate::db::Config { url: db.pg_config().clone() }; + let pool = crate::db::Pool::new(&logctx.log, &cfg); + let conn = pool.pool().get().await.unwrap(); + + let id = Uuid::nil(); + let max_instance_gen = 0; + let project_id = Uuid::nil(); + let cpus_diff = 16.try_into().unwrap(); + let ram_diff = 2048.try_into().unwrap(); + + let query = VirtualProvisioningCollectionUpdate::new_delete_instance( + id, + max_instance_gen, + cpus_diff, + ram_diff, + project_id, + ); + let _ = query + .explain_async(&conn) + .await + .expect("Failed to explain query - is it valid SQL?"); + + db.cleanup().await.unwrap(); + logctx.cleanup_successful(); + } } diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_instance.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_instance.sql index fcabefef26..48094a8371 100644 --- a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_instance.sql +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_instance.sql @@ -37,9 +37,9 @@ WITH WHERE virtual_provisioning_resource.id = $4 LIMIT - $5 + 1 ) - = $6 + = 1 AS update ), unused_cte_arm @@ -47,7 +47,7 @@ WITH DELETE FROM virtual_provisioning_resource WHERE - virtual_provisioning_resource.id = $7 + virtual_provisioning_resource.id = $5 AND virtual_provisioning_resource.id = ( SELECT @@ -55,9 +55,9 @@ WITH FROM instance WHERE - instance.id = $8 AND instance.state_generation < $9 + instance.id = $6 AND instance.state_generation < $7 LIMIT - $10 + 1 ) RETURNING virtual_provisioning_resource.id, @@ -73,11 +73,11 @@ WITH virtual_provisioning_collection SET time_modified = current_timestamp(), - cpus_provisioned = virtual_provisioning_collection.cpus_provisioned - $11, - ram_provisioned = virtual_provisioning_collection.ram_provisioned - $12 + cpus_provisioned = virtual_provisioning_collection.cpus_provisioned - $8, + ram_provisioned = virtual_provisioning_collection.ram_provisioned - $9 WHERE virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) - AND (SELECT do_update.update FROM do_update LIMIT $13) + AND (SELECT do_update.update FROM do_update LIMIT 1) RETURNING virtual_provisioning_collection.id, virtual_provisioning_collection.time_modified, diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_storage.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_storage.sql index 72c0b81e15..b607ac4185 100644 --- a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_storage.sql +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_storage.sql @@ -37,9 +37,9 @@ WITH WHERE virtual_provisioning_resource.id = $4 LIMIT - $5 + 1 ) - = $6 + = 1 AS update ), unused_cte_arm @@ -47,7 +47,7 @@ WITH DELETE FROM virtual_provisioning_resource WHERE - virtual_provisioning_resource.id = $7 + virtual_provisioning_resource.id = $5 RETURNING virtual_provisioning_resource.id, virtual_provisioning_resource.time_modified, @@ -63,10 +63,10 @@ WITH SET time_modified = current_timestamp(), virtual_disk_bytes_provisioned - = virtual_provisioning_collection.virtual_disk_bytes_provisioned - $8 + = virtual_provisioning_collection.virtual_disk_bytes_provisioned - $6 WHERE virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) - AND (SELECT do_update.update FROM do_update LIMIT $9) + AND (SELECT do_update.update FROM do_update LIMIT 1) RETURNING virtual_provisioning_collection.id, virtual_provisioning_collection.time_modified, diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_instance.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_instance.sql index 753b7f09f3..38f10a7148 100644 --- a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_instance.sql +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_instance.sql @@ -39,17 +39,17 @@ WITH WHERE virtual_provisioning_resource.id = $4 LIMIT - $5 + 1 ) - = $6 + = 0 AND CAST( IF( ( - $7 = $8 - OR (SELECT quotas.cpus FROM quotas LIMIT $9) + $5 = 0 + OR (SELECT quotas.cpus FROM quotas LIMIT 1) >= ( - (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT $10) - + $11 + (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT 1) + + $6 ) ), 'TRUE', @@ -61,11 +61,10 @@ WITH AND CAST( IF( ( - $12 = $13 - OR (SELECT quotas.memory FROM quotas LIMIT $14) + $7 = 0 + OR (SELECT quotas.memory FROM quotas LIMIT 1) >= ( - (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT $15) - + $16 + (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT 1) + $8 ) ), 'TRUE', @@ -77,8 +76,8 @@ WITH AND CAST( IF( ( - $17 = $18 - OR (SELECT quotas.storage FROM quotas LIMIT $19) + $9 = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) >= ( ( SELECT @@ -86,9 +85,9 @@ WITH FROM silo_provisioned LIMIT - $20 + 1 ) - + $21 + + $10 ) ), 'TRUE', @@ -112,7 +111,7 @@ WITH ram_provisioned ) VALUES - ($22, DEFAULT, $23, $24, $25, $26) + ($11, DEFAULT, $12, $13, $14, $15) ON CONFLICT DO NOTHING @@ -130,11 +129,11 @@ WITH virtual_provisioning_collection SET time_modified = current_timestamp(), - cpus_provisioned = virtual_provisioning_collection.cpus_provisioned + $27, - ram_provisioned = virtual_provisioning_collection.ram_provisioned + $28 + cpus_provisioned = virtual_provisioning_collection.cpus_provisioned + $16, + ram_provisioned = virtual_provisioning_collection.ram_provisioned + $17 WHERE virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) - AND (SELECT do_update.update FROM do_update LIMIT $29) + AND (SELECT do_update.update FROM do_update LIMIT 1) RETURNING virtual_provisioning_collection.id, virtual_provisioning_collection.time_modified, diff --git a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_storage.sql b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_storage.sql index 040a5dc20c..87cd227ed9 100644 --- a/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_storage.sql +++ b/nexus/db-queries/tests/output/virtual_provisioning_collection_update_insert_storage.sql @@ -39,17 +39,17 @@ WITH WHERE virtual_provisioning_resource.id = $4 LIMIT - $5 + 1 ) - = $6 + = 0 AND CAST( IF( ( - $7 = $8 - OR (SELECT quotas.cpus FROM quotas LIMIT $9) + $5 = 0 + OR (SELECT quotas.cpus FROM quotas LIMIT 1) >= ( - (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT $10) - + $11 + (SELECT silo_provisioned.cpus_provisioned FROM silo_provisioned LIMIT 1) + + $6 ) ), 'TRUE', @@ -61,11 +61,10 @@ WITH AND CAST( IF( ( - $12 = $13 - OR (SELECT quotas.memory FROM quotas LIMIT $14) + $7 = 0 + OR (SELECT quotas.memory FROM quotas LIMIT 1) >= ( - (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT $15) - + $16 + (SELECT silo_provisioned.ram_provisioned FROM silo_provisioned LIMIT 1) + $8 ) ), 'TRUE', @@ -77,8 +76,8 @@ WITH AND CAST( IF( ( - $17 = $18 - OR (SELECT quotas.storage FROM quotas LIMIT $19) + $9 = 0 + OR (SELECT quotas.storage FROM quotas LIMIT 1) >= ( ( SELECT @@ -86,9 +85,9 @@ WITH FROM silo_provisioned LIMIT - $20 + 1 ) - + $21 + + $10 ) ), 'TRUE', @@ -112,7 +111,7 @@ WITH ram_provisioned ) VALUES - ($22, DEFAULT, $23, $24, $25, $26) + ($11, DEFAULT, $12, $13, $14, $15) ON CONFLICT DO NOTHING @@ -131,10 +130,10 @@ WITH SET time_modified = current_timestamp(), virtual_disk_bytes_provisioned - = virtual_provisioning_collection.virtual_disk_bytes_provisioned + $27 + = virtual_provisioning_collection.virtual_disk_bytes_provisioned + $16 WHERE virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections) - AND (SELECT do_update.update FROM do_update LIMIT $28) + AND (SELECT do_update.update FROM do_update LIMIT 1) RETURNING virtual_provisioning_collection.id, virtual_provisioning_collection.time_modified,