Add EXPECTORATE tests for virtual_provisioning_collection CTE (#5081)

This is a direct follow-up to
#5063, focused on the
`virtual_provisioning_collection` CTE.

This PR adds a test which validates the current SQL query, before
re-structuring it to use `TypedSqlQuery`.

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

@@ -369,6 +369,7 @@ mod test {
     use super::*;
     use crate::db::datastore::REGION_REDUNDANCY_THRESHOLD;
     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;
@@ -395,15 +396,11 @@ mod test {
             },
             REGION_REDUNDANCY_THRESHOLD,
         );
-        let s = dev::db::format_sql(
-            &diesel::debug_query::<Pg, _>(&region_allocate).to_string(),
-        )
-        .await
-        .unwrap();
-        expectorate::assert_contents(
+        expectorate_query_contents(
+            &region_allocate,
             "tests/output/region_allocate_distinct_sleds.sql",
-            &s,
-        );
+        )
+        .await;
 
         // Second structure: "Random"
 
@@ -415,15 +412,11 @@ mod test {
             &RegionAllocationStrategy::Random { seed: Some(1) },
             REGION_REDUNDANCY_THRESHOLD,
         );
-        let s = dev::db::format_sql(
-            &diesel::debug_query::<Pg, _>(&region_allocate).to_string(),
-        )
-        .await
-        .unwrap();
-        expectorate::assert_contents(
+        expectorate_query_contents(
+            &region_allocate,
             "tests/output/region_allocate_random_sleds.sql",
-            &s,
-        );
+        )
+        .await;
     }
 
     // Explain the possible forms of the SQL query to ensure that it

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

@@ -646,3 +646,90 @@ impl Query for VirtualProvisioningCollectionUpdate {
 }
 
 impl RunQueryDsl<DbConnection> for VirtualProvisioningCollectionUpdate {}
+
+#[cfg(test)]
+mod test {
+    use super::*;
+    use crate::db::raw_query_builder::expectorate_query_contents;
+    use uuid::Uuid;
+
+    // These tests are a bit of a "change detector", but they're here to help
+    // with debugging too. If you change this query, it can be useful to see
+    // exactly how the output SQL has been altered.
+
+    #[tokio::test]
+    async fn expectorate_query_insert_storage() {
+        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,
+        );
+        expectorate_query_contents(
+            &query,
+            "tests/output/virtual_provisioning_collection_update_insert_storage.sql",
+        ).await;
+    }
+
+    #[tokio::test]
+    async fn expectorate_query_delete_storage() {
+        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,
+        );
+
+        expectorate_query_contents(
+            &query,
+            "tests/output/virtual_provisioning_collection_update_delete_storage.sql",
+        ).await;
+    }
+
+    #[tokio::test]
+    async fn expectorate_query_insert_instance() {
+        let id = Uuid::nil();
+        let project_id = Uuid::nil();
+        let cpus_diff = 4;
+        let ram_diff = 2048.try_into().unwrap();
+
+        let query = VirtualProvisioningCollectionUpdate::new_insert_instance(
+            id, cpus_diff, ram_diff, project_id,
+        );
+
+        expectorate_query_contents(
+            &query,
+            "tests/output/virtual_provisioning_collection_update_insert_instance.sql",
+        ).await;
+    }
+
+    #[tokio::test]
+    async fn expectorate_query_delete_instance() {
+        let id = Uuid::nil();
+        let project_id = Uuid::nil();
+        let cpus_diff = 4;
+        let ram_diff = 2048.try_into().unwrap();
+        let max_instance_gen = 0;
+
+        let query = VirtualProvisioningCollectionUpdate::new_delete_instance(
+            id,
+            max_instance_gen,
+            cpus_diff,
+            ram_diff,
+            project_id,
+        );
+
+        expectorate_query_contents(
+            &query,
+            "tests/output/virtual_provisioning_collection_update_delete_instance.sql",
+        ).await;
+    }
+}

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

@@ -181,3 +181,18 @@ impl<T> RunQueryDsl<DbConnection> for TypedSqlQuery<T> {}
 impl<T> Query for TypedSqlQuery<T> {
     type SqlType = T;
 }
+
+#[cfg(test)]
+pub async fn expectorate_query_contents<T: QueryFragment<Pg>>(
+    query: T,
+    path: &str,
+) {
+    use omicron_test_utils::dev;
+
+    let s =
+        dev::db::format_sql(&diesel::debug_query::<Pg, _>(&query).to_string())
+            .await
+            .expect("Failed to format SQL");
+
+    expectorate::assert_contents(path, &s);
+}

nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_instance.sql

@@ -0,0 +1,97 @@
+WITH
+  parent_silo AS (SELECT project.silo_id AS id FROM project WHERE project.id = $1),
+  all_collections
+    AS (
+      ((SELECT $2 AS id) UNION (SELECT parent_silo.id AS id FROM parent_silo))
+      UNION (SELECT $3 AS id)
+    ),
+  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
+    ),
+  do_update
+    AS (
+      SELECT
+        (
+          SELECT
+            count(*)
+          FROM
+            virtual_provisioning_resource
+          WHERE
+            virtual_provisioning_resource.id = $4
+          LIMIT
+            $5
+        )
+        = $6
+          AS update
+    ),
+  unused_cte_arm
+    AS (
+      DELETE FROM
+        virtual_provisioning_resource
+      WHERE
+        virtual_provisioning_resource.id = $7
+        AND virtual_provisioning_resource.id
+          = (
+              SELECT
+                instance.id
+              FROM
+                instance
+              WHERE
+                instance.id = $8 AND instance.state_generation < $9
+              LIMIT
+                $10
+            )
+      RETURNING
+        virtual_provisioning_resource.id,
+        virtual_provisioning_resource.time_modified,
+        virtual_provisioning_resource.resource_type,
+        virtual_provisioning_resource.virtual_disk_bytes_provisioned,
+        virtual_provisioning_resource.cpus_provisioned,
+        virtual_provisioning_resource.ram_provisioned
+    ),
+  virtual_provisioning_collection
+    AS (
+      UPDATE
+        virtual_provisioning_collection
+      SET
+        time_modified = current_timestamp(),
+        cpus_provisioned = virtual_provisioning_collection.cpus_provisioned - $11,
+        ram_provisioned = virtual_provisioning_collection.ram_provisioned - $12
+      WHERE
+        virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections)
+        AND (SELECT do_update.update FROM do_update LIMIT $13)
+      RETURNING
+        virtual_provisioning_collection.id,
+        virtual_provisioning_collection.time_modified,
+        virtual_provisioning_collection.collection_type,
+        virtual_provisioning_collection.virtual_disk_bytes_provisioned,
+        virtual_provisioning_collection.cpus_provisioned,
+        virtual_provisioning_collection.ram_provisioned
+    )
+SELECT
+  virtual_provisioning_collection.id,
+  virtual_provisioning_collection.time_modified,
+  virtual_provisioning_collection.collection_type,
+  virtual_provisioning_collection.virtual_disk_bytes_provisioned,
+  virtual_provisioning_collection.cpus_provisioned,
+  virtual_provisioning_collection.ram_provisioned
+FROM
+  virtual_provisioning_collection

nexus/db-queries/tests/output/virtual_provisioning_collection_update_delete_storage.sql

@@ -0,0 +1,86 @@
+WITH
+  parent_silo AS (SELECT project.silo_id AS id FROM project WHERE project.id = $1),
+  all_collections
+    AS (
+      ((SELECT $2 AS id) UNION (SELECT parent_silo.id AS id FROM parent_silo))
+      UNION (SELECT $3 AS id)
+    ),
+  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
+    ),
+  do_update
+    AS (
+      SELECT
+        (
+          SELECT
+            count(*)
+          FROM
+            virtual_provisioning_resource
+          WHERE
+            virtual_provisioning_resource.id = $4
+          LIMIT
+            $5
+        )
+        = $6
+          AS update
+    ),
+  unused_cte_arm
+    AS (
+      DELETE FROM
+        virtual_provisioning_resource
+      WHERE
+        virtual_provisioning_resource.id = $7
+      RETURNING
+        virtual_provisioning_resource.id,
+        virtual_provisioning_resource.time_modified,
+        virtual_provisioning_resource.resource_type,
+        virtual_provisioning_resource.virtual_disk_bytes_provisioned,
+        virtual_provisioning_resource.cpus_provisioned,
+        virtual_provisioning_resource.ram_provisioned
+    ),
+  virtual_provisioning_collection
+    AS (
+      UPDATE
+        virtual_provisioning_collection
+      SET
+        time_modified = current_timestamp(),
+        virtual_disk_bytes_provisioned
+          = virtual_provisioning_collection.virtual_disk_bytes_provisioned - $8
+      WHERE
+        virtual_provisioning_collection.id = ANY (SELECT all_collections.id FROM all_collections)
+        AND (SELECT do_update.update FROM do_update LIMIT $9)
+      RETURNING
+        virtual_provisioning_collection.id,
+        virtual_provisioning_collection.time_modified,
+        virtual_provisioning_collection.collection_type,
+        virtual_provisioning_collection.virtual_disk_bytes_provisioned,
+        virtual_provisioning_collection.cpus_provisioned,
+        virtual_provisioning_collection.ram_provisioned
+    )
+SELECT
+  virtual_provisioning_collection.id,
+  virtual_provisioning_collection.time_modified,
+  virtual_provisioning_collection.collection_type,
+  virtual_provisioning_collection.virtual_disk_bytes_provisioned,
+  virtual_provisioning_collection.cpus_provisioned,
+  virtual_provisioning_collection.ram_provisioned
+FROM
+  virtual_provisioning_collection