[authz-macros] accept an optional input_key argument

Cargo.toml

@@ -290,6 +290,7 @@ postgres-protocol = "0.6.6"
 predicates = "3.0.4"
 pretty_assertions = "1.4.0"
 pretty-hex = "0.4.0"
+prettyplease = "0.2.15"
 proc-macro2 = "1.0"
 progenitor = { git = "https://github.com/oxidecomputer/progenitor", branch = "main" }
 progenitor-client = { git = "https://github.com/oxidecomputer/progenitor", branch = "main" }

nexus/authz-macros/Cargo.toml

@@ -15,3 +15,6 @@ serde.workspace = true
 serde_tokenstream.workspace = true
 syn.workspace = true
 omicron-workspace-hack.workspace = true
+
+[dev-dependencies]
+prettyplease.workspace = true

nexus/authz-macros/src/lib.rs

@@ -95,6 +95,33 @@ use serde_tokenstream::ParseWrapper;
 ///     polar_snippet = FleetChild,
 /// }
 /// ```
+///
+/// In some cases, it may be more convenient to identify a composite key with a
+/// struct rather than relying on tuples. This is supported too:
+///
+/// ```ignore
+/// struct SomeCompositeId {
+///     foo: String,
+///     bar: String,
+/// }
+///
+/// // There needs to be a `From` impl from the composite ID to the primary key.
+/// impl From<SomeCompositeId> for (String, String) {
+///     fn from(id: SomeCompositeId) -> Self {
+///         (id.foo, id.bar)
+///     }
+/// }
+///
+/// authz_resource! {
+///     name = "MyResource",
+///     parent = "Fleet",
+///     primary_key = (String, String),
+///     input_key = SomeCompositeId,
+///     roles_allowed = false,
+///     polar_snippet = FleetChild,
+/// }
+/// ```
+
 // Allow private intra-doc links.  This is useful because the `Input` struct
 // cannot be exported (since we're a proc macro crate, and we can't expose
 // a struct), but its documentation is very useful.
@@ -121,6 +148,12 @@ struct Input {
     parent: String,
     /// Rust type for the primary key for this resource
     primary_key: ParseWrapper<syn::Type>,
+    /// Rust type for the input key for this resource (the key users specify
+    /// for this resource, convertible to `primary_key`).
+    ///
+    /// This is the same as primary_key if not specified.
+    #[serde(default)]
+    input_key: Option<ParseWrapper<syn::Type>>,
     /// Whether roles may be attached directly to this resource
     roles_allowed: bool,
     /// How to generate the Polar snippet for this resource
@@ -153,6 +186,9 @@ fn do_authz_resource(
     let parent_resource_name = format_ident!("{}", input.parent);
     let parent_as_snake = heck::AsSnakeCase(&input.parent).to_string();
     let primary_key_type = &*input.primary_key;
+    let input_key_type =
+        &**input.input_key.as_ref().unwrap_or(&input.primary_key);
+
     let (has_role_body, as_roles_body, api_resource_roles_trait) =
         if input.roles_allowed {
             (
@@ -334,6 +370,21 @@ fn do_authz_resource(
             /// `parent`, unique key `key`, looked up as described by
             /// `lookup_type`
             pub fn new(
+                parent: #parent_resource_name,
+                key: #input_key_type,
+                lookup_type: LookupType,
+            ) -> #resource_name {
+                #resource_name {
+                    parent,
+                    key: key.into(),
+                    lookup_type,
+                }
+            }
+
+            /// A version of `new` that takes the primary key type directly.
+            /// This is only different from [`Self::new`] if this resource
+            /// uses a different input key type.
+            pub fn with_primary_key(
                 parent: #parent_resource_name,
                 key: #primary_key_type,
                 lookup_type: LookupType,
@@ -346,7 +397,7 @@ fn do_authz_resource(
             }
 
             pub fn id(&self) -> #primary_key_type {
-                self.key.clone()
+                self.key.clone().into()
             }
 
             /// Describes how to register this type with Oso
@@ -411,15 +462,36 @@ fn do_authz_resource(
 
 // See the test for lookup_resource.
 #[cfg(test)]
-#[test]
-fn test_authz_dump() {
-    let output = do_authz_resource(quote! {
-        name = "Organization",
-        parent = "Fleet",
-        primary_key = Uuid,
-        roles_allowed = false,
-        polar_snippet = Custom,
-    })
-    .unwrap();
-    println!("{}", output);
+mod tests {
+    use super::*;
+    #[test]
+    fn test_authz_dump() {
+        let output = do_authz_resource(quote! {
+            name = "Organization",
+            parent = "Fleet",
+            primary_key = Uuid,
+            roles_allowed = false,
+            polar_snippet = Custom,
+        })
+        .unwrap();
+        println!("{}", pretty_format(output));
+
+        let output = do_authz_resource(quote! {
+            name = "Instance",
+            parent = "Project",
+            primary_key = (String, String),
+            // The SomeCompositeId type doesn't exist, but that's okay because
+            // this code is never compiled, just printed out.
+            input_key = SomeCompositeId,
+            roles_allowed = false,
+            polar_snippet = InProject,
+        })
+        .unwrap();
+        println!("{}", pretty_format(output));
+    }
+
+    fn pretty_format(input: TokenStream) -> String {
+        let parsed = syn::parse2(input).unwrap();
+        prettyplease::unparse(&parsed)
+    }
 }

nexus/db-macros/Cargo.toml

@@ -18,4 +18,4 @@ syn = { workspace = true, features = ["extra-traits"] }
 omicron-workspace-hack.workspace = true
 
 [dev-dependencies]
-rustfmt-wrapper.workspace = true
+prettyplease.workspace = true

nexus/db-macros/src/lookup.rs

@@ -406,7 +406,7 @@ fn generate_misc_helpers(config: &Config) -> TokenStream {
             db_row: &nexus_db_model::#resource_name,
             lookup_type: LookupType,
         ) -> authz::#resource_name {
-            authz::#resource_name::new(
+            authz::#resource_name::with_primary_key(
                 authz_parent.clone(),
                 db_row.id(),
                 lookup_type
@@ -923,8 +923,8 @@ fn generate_database_functions(config: &Config) -> TokenStream {
 #[cfg(test)]
 mod test {
     use super::lookup_resource;
+    use proc_macro2::TokenStream;
     use quote::quote;
-    use rustfmt_wrapper::rustfmt;
 
     #[test]
     #[ignore]
@@ -938,7 +938,7 @@ mod test {
             primary_key_columns = [ { column_name = "id", rust_type = Uuid } ]
         })
         .unwrap();
-        println!("{}", rustfmt(output).unwrap());
+        println!("{}", pretty_format(output));
 
         let output = lookup_resource(quote! {
             name = "SiloUser",
@@ -949,7 +949,7 @@ mod test {
             primary_key_columns = [ { column_name = "id", rust_type = Uuid } ]
         })
         .unwrap();
-        println!("{}", rustfmt(output).unwrap());
+        println!("{}", pretty_format(output));
 
         let output = lookup_resource(quote! {
             name = "UpdateArtifact",
@@ -964,6 +964,11 @@ mod test {
             ]
         })
         .unwrap();
-        println!("{}", rustfmt(output).unwrap());
+        println!("{}", pretty_format(output));
+    }
+
+    fn pretty_format(input: TokenStream) -> String {
+        let parsed = syn::parse2(input).unwrap();
+        prettyplease::unparse(&parsed)
     }
 }