[authz-macros] accept an optional input_key argument (#4707)

In some cases including composite keys, it can be better to make the outside
representation of the primary key some kind of struct, rather than passing
around tuples of various types.

Enable that in the `authz_resource` macro by allowing users to specify an
optional `input_key` argument, which represents a better view into the primary
key.

I'm not entirely sure that the `From` trait is the right thing to use here, but
it seems like a pretty low-cost decision to change in the future.

As part of this PR I also switched to the prettyplease crate, which as the
README explains is more suitable for generated code than rustfmt:
https://crates.io/crates/prettyplease

Cargo.toml

@@ -289,6 +289,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)
     }
 }