Add images to resource accounting

nexus/db-model/src/schema_versions.rs

@@ -16,7 +16,7 @@ use std::{collections::BTreeMap, sync::LazyLock};
 ///
 /// This must be updated when you change the database schema.  Refer to
 /// schema/crdb/README.adoc in the root of this repository for details.
-pub const SCHEMA_VERSION: Version = Version::new(219, 0, 0);
+pub const SCHEMA_VERSION: Version = Version::new(220, 0, 0);
 
 /// List of all past database schema versions, in *reverse* order
 ///
@@ -28,6 +28,7 @@ static KNOWN_VERSIONS: LazyLock<Vec<KnownVersion>> = LazyLock::new(|| {
         // |  leaving the first copy as an example for the next person.
         // v
         // KnownVersion::new(next_int, "unique-dirname-with-the-sql-files"),
+        KnownVersion::new(220, "image-virtual-provisioning"),
         KnownVersion::new(219, "blueprint-sled-last-used-ip"),
         KnownVersion::new(218, "measurements"),
         KnownVersion::new(217, "multiple-default-ip-pools-per-silo"),

nexus/db-model/src/virtual_provisioning_resource.rs

@@ -13,6 +13,7 @@ pub enum ResourceTypeProvisioned {
     Instance,
     Disk,
     Snapshot,
+    Image,
 }
 
 impl std::fmt::Display for ResourceTypeProvisioned {
@@ -21,6 +22,7 @@ impl std::fmt::Display for ResourceTypeProvisioned {
             ResourceTypeProvisioned::Instance => write!(f, "instance"),
             ResourceTypeProvisioned::Disk => write!(f, "disk"),
             ResourceTypeProvisioned::Snapshot => write!(f, "snapshot"),
+            ResourceTypeProvisioned::Image => write!(f, "image"),
         }
     }
 }

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

@@ -24,6 +24,7 @@ use uuid::Uuid;
 pub enum StorageType {
     Disk,
     Snapshot,
+    Image,
 }
 
 impl From<StorageType> for crate::db::model::ResourceTypeProvisioned {
@@ -37,6 +38,9 @@ impl From<StorageType> for crate::db::model::ResourceTypeProvisioned {
             StorageType::Snapshot => {
                 crate::db::model::ResourceTypeProvisioned::Snapshot
             }
+            StorageType::Image => {
+                crate::db::model::ResourceTypeProvisioned::Image
+            }
         }
     }
 }
@@ -178,6 +182,23 @@ impl DataStore {
         .await
     }
 
+    pub async fn virtual_provisioning_collection_insert_project_image(
+        &self,
+        opctx: &OpContext,
+        id: Uuid,
+        project_id: Uuid,
+        disk_byte_diff: ByteCount,
+    ) -> Result<Vec<VirtualProvisioningCollection>, Error> {
+        self.virtual_provisioning_collection_insert_storage(
+            opctx,
+            id,
+            project_id,
+            disk_byte_diff,
+            StorageType::Image,
+        )
+        .await
+    }
+
     /// Transitively updates all provisioned disk provisions from project -> fleet.
     async fn virtual_provisioning_collection_insert_storage(
         &self,
@@ -236,6 +257,22 @@ impl DataStore {
         .await
     }
 
+    pub async fn virtual_provisioning_collection_delete_project_image(
+        &self,
+        opctx: &OpContext,
+        id: Uuid,
+        project_id: Uuid,
+        disk_byte_diff: ByteCount,
+    ) -> Result<Vec<VirtualProvisioningCollection>, Error> {
+        self.virtual_provisioning_collection_delete_storage(
+            opctx,
+            id,
+            project_id,
+            disk_byte_diff,
+        )
+        .await
+    }
+
     // Transitively updates all provisioned disk provisions from project -> fleet.
     async fn virtual_provisioning_collection_delete_storage(
         &self,
@@ -258,6 +295,59 @@ impl DataStore {
         Ok(provisions)
     }
 
+    /// Insert storage accounting for a silo-scoped image.
+    ///
+    /// This updates the Silo and Fleet collections (no Project).
+    pub async fn virtual_provisioning_collection_insert_silo_image(
+        &self,
+        opctx: &OpContext,
+        id: Uuid,
+        silo_id: Uuid,
+        disk_byte_diff: ByteCount,
+    ) -> Result<Vec<VirtualProvisioningCollection>, Error> {
+        let provisions =
+            VirtualProvisioningCollectionUpdate::new_insert_silo_storage(
+                id,
+                disk_byte_diff,
+                silo_id,
+                StorageType::Image,
+            )
+            .get_results_async(&*self.pool_connection_authorized(opctx).await?)
+            .await
+            .map_err(|e| {
+                crate::db::queries::virtual_provisioning_collection_update::from_diesel(e)
+            })?;
+        self.virtual_provisioning_collection_producer
+            .append_disk_metrics(&provisions)?;
+        Ok(provisions)
+    }
+
+    /// Delete storage accounting for a silo-scoped image.
+    ///
+    /// This updates the Silo and Fleet collections (no Project).
+    pub async fn virtual_provisioning_collection_delete_silo_image(
+        &self,
+        opctx: &OpContext,
+        id: Uuid,
+        silo_id: Uuid,
+        disk_byte_diff: ByteCount,
+    ) -> Result<Vec<VirtualProvisioningCollection>, Error> {
+        let provisions =
+            VirtualProvisioningCollectionUpdate::new_delete_silo_storage(
+                id,
+                disk_byte_diff,
+                silo_id,
+            )
+            .get_results_async(&*self.pool_connection_authorized(opctx).await?)
+            .await
+            .map_err(|e| {
+                crate::db::queries::virtual_provisioning_collection_update::from_diesel(e)
+            })?;
+        self.virtual_provisioning_collection_producer
+            .append_disk_metrics(&provisions)?;
+        Ok(provisions)
+    }
+
     /// Transitively updates all CPU/RAM provisions from project -> fleet.
     pub async fn virtual_provisioning_collection_insert_instance(
         &self,