Use a typed UUID for Dataset (#7048)

clients/nexus-client/src/lib.rs

@@ -48,6 +48,7 @@ progenitor::generate_api!(
         RecoverySiloConfig = nexus_sled_agent_shared::recovery_silo::RecoverySiloConfig,
         Srv = nexus_types::internal_api::params::Srv,
         TypedUuidForCollectionKind = omicron_uuid_kinds::CollectionUuid,
+        TypedUuidForDatasetKind = omicron_uuid_kinds::TypedUuid<omicron_uuid_kinds::DatasetKind>,
         TypedUuidForDemoSagaKind = omicron_uuid_kinds::DemoSagaUuid,
         TypedUuidForDownstairsKind = omicron_uuid_kinds::TypedUuid<omicron_uuid_kinds::DownstairsKind>,
         TypedUuidForPropolisKind = omicron_uuid_kinds::TypedUuid<omicron_uuid_kinds::PropolisKind>,

dev-tools/omdb/src/bin/omdb/db.rs

@@ -52,6 +52,7 @@ use internal_dns_types::names::ServiceName;
 use ipnetwork::IpNetwork;
 use itertools::Itertools;
 use nexus_config::PostgresConfigWithUrl;
+use nexus_db_model::to_db_typed_uuid;
 use nexus_db_model::Dataset;
 use nexus_db_model::Disk;
 use nexus_db_model::DnsGroup;
@@ -126,6 +127,7 @@ use omicron_common::api::external::Generation;
 use omicron_common::api::external::InstanceState;
 use omicron_common::api::external::MacAddr;
 use omicron_uuid_kinds::CollectionUuid;
+use omicron_uuid_kinds::DatasetUuid;
 use omicron_uuid_kinds::GenericUuid;
 use omicron_uuid_kinds::InstanceUuid;
 use omicron_uuid_kinds::PropolisUuid;
@@ -772,7 +774,7 @@ struct RegionSnapshotReplacementInfoArgs {
 #[derive(Debug, Args)]
 struct RegionSnapshotReplacementRequestArgs {
     /// The dataset id for a given region snapshot
-    dataset_id: Uuid,
+    dataset_id: DatasetUuid,
 
     /// The region id for a given region snapshot
     region_id: Uuid,
@@ -1634,7 +1636,7 @@ async fn cmd_db_disk_physical(
         .context("loading zpool from pysical disk id")?;
 
     let mut sled_ids = HashSet::new();
-    let mut dataset_ids = HashSet::new();
+    let mut dataset_ids: HashSet<DatasetUuid> = HashSet::new();
 
     if zpools.is_empty() {
         println!("Found no zpools on physical disk UUID {}", args.uuid);
@@ -1692,7 +1694,7 @@ async fn cmd_db_disk_physical(
     for did in dataset_ids.clone().into_iter() {
         use db::schema::region::dsl as region_dsl;
         let regions = region_dsl::region
-            .filter(region_dsl::dataset_id.eq(did))
+            .filter(region_dsl::dataset_id.eq(to_db_typed_uuid(did)))
             .select(Region::as_select())
             .load_async(&*conn)
             .await
@@ -1774,6 +1776,9 @@ async fn cmd_db_disk_physical(
     println!("{}", table);
 
     // Collect the region_snapshots associated with the dataset IDs
+    let dataset_ids: Vec<_> =
+        dataset_ids.into_iter().map(|did| to_db_typed_uuid(did)).collect();
+
     let limit = fetch_opts.fetch_limit;
     use db::schema::region_snapshot::dsl as region_snapshot_dsl;
     let region_snapshots = region_snapshot_dsl::region_snapshot
@@ -2328,7 +2333,7 @@ async fn cmd_db_region_list(
         #[derive(Tabled)]
         struct RegionRow {
             id: Uuid,
-            dataset_id: Uuid,
+            dataset_id: DatasetUuid,
             volume_id: Uuid,
             block_size: i64,
             blocks_per_extent: u64,
@@ -2563,7 +2568,7 @@ async fn cmd_db_region_find_deleted(
 
     #[derive(Tabled)]
     struct Row {
-        dataset_id: Uuid,
+        dataset_id: DatasetUuid,
         region_id: Uuid,
         volume_id: String,
     }
@@ -4418,7 +4423,7 @@ async fn cmd_db_validate_volume_references(
 
     #[derive(Tabled)]
     struct Row {
-        dataset_id: Uuid,
+        dataset_id: DatasetUuid,
         region_id: Uuid,
         snapshot_id: Uuid,
         error: String,
@@ -4480,7 +4485,7 @@ async fn cmd_db_validate_volume_references(
 
         if matching_volumes != region_snapshot.volume_references as usize {
             rows.push(Row {
-                dataset_id: region_snapshot.dataset_id,
+                dataset_id: region_snapshot.dataset_id.into(),
                 region_id: region_snapshot.region_id,
                 snapshot_id: region_snapshot.snapshot_id,
                 error: format!(
@@ -4498,7 +4503,7 @@ async fn cmd_db_validate_volume_references(
 
             if matching_volumes == 0 && !region_snapshot.deleting {
                 rows.push(Row {
-                    dataset_id: region_snapshot.dataset_id,
+                    dataset_id: region_snapshot.dataset_id.into(),
                     region_id: region_snapshot.region_id,
                     snapshot_id: region_snapshot.snapshot_id,
                     error: String::from(
@@ -4556,7 +4561,7 @@ async fn cmd_db_validate_region_snapshots(
 
     #[derive(Tabled)]
     struct Row {
-        dataset_id: Uuid,
+        dataset_id: DatasetUuid,
         region_id: Uuid,
         snapshot_id: Uuid,
         dataset_addr: std::net::SocketAddrV6,
@@ -4640,7 +4645,7 @@ async fn cmd_db_validate_region_snapshots(
                                 // This is ok - Nexus currently soft-deletes its
                                 // resource records.
                                 rows.push(Row {
-                                    dataset_id: region_snapshot.dataset_id,
+                                    dataset_id: region_snapshot.dataset_id.into(),
                                     region_id: region_snapshot.region_id,
                                     snapshot_id: region_snapshot.snapshot_id,
                                     dataset_addr,
@@ -4655,7 +4660,7 @@ async fn cmd_db_validate_region_snapshots(
                                 // higher level Snapshot was not deleted!
 
                                 rows.push(Row {
-                                    dataset_id: region_snapshot.dataset_id,
+                                    dataset_id: region_snapshot.dataset_id.into(),
                                     region_id: region_snapshot.region_id,
                                     snapshot_id: region_snapshot.snapshot_id,
                                     dataset_addr,
@@ -4684,7 +4689,7 @@ async fn cmd_db_validate_region_snapshots(
                             // the Agent, so it's a bug.
 
                             rows.push(Row {
-                                dataset_id: region_snapshot.dataset_id,
+                                dataset_id: region_snapshot.dataset_id.into(),
                                 region_id: region_snapshot.region_id,
                                 snapshot_id: region_snapshot.snapshot_id,
                                 dataset_addr,

nexus/db-model/src/dataset.rs

@@ -2,7 +2,8 @@
 // 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/.
 
-use super::{ByteCount, DatasetKind, Generation, Region, SqlU16};
+use super::DatasetKind;
+use super::{ByteCount, Generation, Region, SqlU16};
 use crate::collection::DatastoreCollectionConfig;
 use crate::ipv6;
 use crate::schema::{dataset, region};
@@ -11,7 +12,6 @@ use db_macros::Asset;
 use nexus_types::deployment::BlueprintDatasetConfig;
 use omicron_common::api::external::Error;
 use omicron_common::api::internal::shared::DatasetKind as ApiDatasetKind;
-use omicron_uuid_kinds::DatasetUuid;
 use omicron_uuid_kinds::GenericUuid;
 use omicron_uuid_kinds::ZpoolUuid;
 use serde::{Deserialize, Serialize};
@@ -34,6 +34,7 @@ use uuid::Uuid;
     PartialEq,
 )]
 #[diesel(table_name = dataset)]
+#[asset(uuid_kind = DatasetKind)]
 pub struct Dataset {
     #[diesel(embed)]
     identity: DatasetIdentity,
@@ -62,7 +63,7 @@ pub struct Dataset {
 
 impl Dataset {
     pub fn new(
-        id: Uuid,
+        id: omicron_uuid_kinds::DatasetUuid,
         pool_id: Uuid,
         addr: Option<SocketAddrV6>,
         api_kind: ApiDatasetKind,
@@ -119,7 +120,7 @@ impl From<BlueprintDatasetConfig> for Dataset {
         };
         let addr = bp.address;
         Self {
-            identity: DatasetIdentity::new(bp.id.into_untyped_uuid()),
+            identity: DatasetIdentity::new(bp.id),
             time_deleted: None,
             rcgen: Generation::new(),
             pool_id: bp.pool.id().into_untyped_uuid(),
@@ -148,7 +149,7 @@ impl TryFrom<Dataset> for omicron_common::disk::DatasetConfig {
         };
 
         Ok(Self {
-            id: DatasetUuid::from_untyped_uuid(dataset.identity.id),
+            id: dataset.identity.id.into(),
             name: omicron_common::disk::DatasetName::new(
                 omicron_common::zpool_name::ZpoolName::new_external(
                     ZpoolUuid::from_untyped_uuid(dataset.pool_id),

nexus/db-model/src/region.rs

@@ -4,9 +4,12 @@
 
 use super::ByteCount;
 use crate::schema::region;
+use crate::typed_uuid::DbTypedUuid;
 use crate::SqlU16;
 use db_macros::Asset;
 use omicron_common::api::external;
+use omicron_uuid_kinds::DatasetKind;
+use omicron_uuid_kinds::DatasetUuid;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -30,7 +33,7 @@ pub struct Region {
     #[diesel(embed)]
     identity: RegionIdentity,
 
-    dataset_id: Uuid,
+    dataset_id: DbTypedUuid<DatasetKind>,
     volume_id: Uuid,
 
     block_size: ByteCount,
@@ -54,7 +57,7 @@ pub struct Region {
 
 impl Region {
     pub fn new(
-        dataset_id: Uuid,
+        dataset_id: DatasetUuid,
         volume_id: Uuid,
         block_size: ByteCount,
         blocks_per_extent: u64,
@@ -64,7 +67,7 @@ impl Region {
     ) -> Self {
         Self {
             identity: RegionIdentity::new(Uuid::new_v4()),
-            dataset_id,
+            dataset_id: dataset_id.into(),
             volume_id,
             block_size,
             blocks_per_extent: blocks_per_extent as i64,
@@ -81,8 +84,8 @@ impl Region {
     pub fn volume_id(&self) -> Uuid {
         self.volume_id
     }
-    pub fn dataset_id(&self) -> Uuid {
-        self.dataset_id
+    pub fn dataset_id(&self) -> DatasetUuid {
+        self.dataset_id.into()
     }
     pub fn block_size(&self) -> external::ByteCount {
         self.block_size.0

nexus/db-model/src/region_snapshot.rs

@@ -3,6 +3,9 @@
 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
 
 use crate::schema::region_snapshot;
+use crate::typed_uuid::DbTypedUuid;
+use omicron_uuid_kinds::DatasetKind;
+use omicron_uuid_kinds::DatasetUuid;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -23,7 +26,7 @@ use uuid::Uuid;
 #[diesel(table_name = region_snapshot)]
 pub struct RegionSnapshot {
     // unique identifier of this region snapshot
-    pub dataset_id: Uuid,
+    pub dataset_id: DbTypedUuid<DatasetKind>,
     pub region_id: Uuid,
     pub snapshot_id: Uuid,
 
@@ -43,13 +46,13 @@ pub struct RegionSnapshot {
 
 impl RegionSnapshot {
     pub fn new(
-        dataset_id: Uuid,
+        dataset_id: DatasetUuid,
         region_id: Uuid,
         snapshot_id: Uuid,
         snapshot_addr: String,
     ) -> Self {
         RegionSnapshot {
-            dataset_id,
+            dataset_id: dataset_id.into(),
             region_id,
             snapshot_id,
             snapshot_addr,
@@ -58,4 +61,8 @@ impl RegionSnapshot {
             deleting: false,
         }
     }
+
+    pub fn dataset_id(&self) -> DatasetUuid {
+        self.dataset_id.into()
+    }
 }

nexus/db-model/src/region_snapshot_replacement.rs

@@ -4,9 +4,12 @@
 
 use super::impl_enum_type;
 use crate::schema::region_snapshot_replacement;
+use crate::typed_uuid::DbTypedUuid;
 use crate::RegionSnapshot;
 use chrono::DateTime;
 use chrono::Utc;
+use omicron_uuid_kinds::DatasetKind;
+use omicron_uuid_kinds::DatasetUuid;
 use serde::{Deserialize, Serialize};
 use uuid::Uuid;
 
@@ -118,7 +121,7 @@ pub struct RegionSnapshotReplacement {
     pub request_time: DateTime<Utc>,
 
     // These are a copy of fields from the corresponding region snapshot record
-    pub old_dataset_id: Uuid,
+    pub old_dataset_id: DbTypedUuid<DatasetKind>,
     pub old_region_id: Uuid,
     pub old_snapshot_id: Uuid,
 
@@ -135,21 +138,21 @@ pub struct RegionSnapshotReplacement {
 impl RegionSnapshotReplacement {
     pub fn for_region_snapshot(region_snapshot: &RegionSnapshot) -> Self {
         Self::new(
-            region_snapshot.dataset_id,
+            region_snapshot.dataset_id(),
             region_snapshot.region_id,
             region_snapshot.snapshot_id,
         )
     }
 
     pub fn new(
-        old_dataset_id: Uuid,
+        old_dataset_id: DatasetUuid,
         old_region_id: Uuid,
         old_snapshot_id: Uuid,
     ) -> Self {
         Self {
             id: Uuid::new_v4(),
             request_time: Utc::now(),
-            old_dataset_id,
+            old_dataset_id: old_dataset_id.into(),
             old_region_id,
             old_snapshot_id,
             old_snapshot_volume_id: None,

nexus/db-model/src/volume_resource_usage.rs

@@ -4,6 +4,9 @@
 
 use super::impl_enum_type;
 use crate::schema::volume_resource_usage;
+use crate::typed_uuid::DbTypedUuid;
+use omicron_uuid_kinds::DatasetKind;
+use omicron_uuid_kinds::DatasetUuid;
 use uuid::Uuid;
 
 impl_enum_type!(
@@ -54,16 +57,22 @@ pub struct VolumeResourceUsageRecord {
 
     pub region_id: Option<Uuid>,
 
-    pub region_snapshot_dataset_id: Option<Uuid>,
+    pub region_snapshot_dataset_id: Option<DbTypedUuid<DatasetKind>>,
     pub region_snapshot_region_id: Option<Uuid>,
     pub region_snapshot_snapshot_id: Option<Uuid>,
 }
 
 #[derive(Debug, Clone)]
 pub enum VolumeResourceUsage {
-    ReadOnlyRegion { region_id: Uuid },
-
-    RegionSnapshot { dataset_id: Uuid, region_id: Uuid, snapshot_id: Uuid },
+    ReadOnlyRegion {
+        region_id: Uuid,
+    },
+
+    RegionSnapshot {
+        dataset_id: DatasetUuid,
+        region_id: Uuid,
+        snapshot_id: Uuid,
+    },
 }
 
 impl VolumeResourceUsageRecord {
@@ -94,7 +103,7 @@ impl VolumeResourceUsageRecord {
 
                 region_id: None,
 
-                region_snapshot_dataset_id: Some(dataset_id),
+                region_snapshot_dataset_id: Some(dataset_id.into()),
                 region_snapshot_region_id: Some(region_id),
                 region_snapshot_snapshot_id: Some(snapshot_id),
             },
@@ -132,7 +141,7 @@ impl TryFrom<VolumeResourceUsageRecord> for VolumeResourceUsage {
                 };
 
                 Ok(VolumeResourceUsage::RegionSnapshot {
-                    dataset_id,
+                    dataset_id: dataset_id.into(),
                     region_id,
                     snapshot_id,
                 })