Add EXPECTORATE tests for virtual_provisioning_collection CTE

oxidecomputer · Feb 15, 2024 · 7f46a0a · 7f46a0a
1 parent e7c71d7
commit 7f46a0a
Show file tree

Hide file tree

Showing 7 changed files with 583 additions and 16 deletions.
diff --git a/nexus/db-queries/src/db/queries/region_allocation.rs b/nexus/db-queries/src/db/queries/region_allocation.rs
@@ -301,6 +301,7 @@ UNION
 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;
@@ -326,15 +327,11 @@ mod test {
                 seed: Some(1),
             },
         );
-        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"
 
@@ -345,15 +342,11 @@ mod test {
             extent_count,
             &RegionAllocationStrategy::Random { seed: Some(1) },
         );
-        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

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
@@ -646,3 +646,71 @@ 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;
+
+    // This test is a bit of a "change detector", but it's 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() {
+        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;
+
+        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;
+
+        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;
+
+        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;
+    }
+}
diff --git a/nexus/db-queries/src/db/raw_query_builder.rs b/nexus/db-queries/src/db/raw_query_builder.rs
@@ -193,3 +193,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);
+}
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
@@ -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
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
@@ -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