Merge pull request #155 from grafbase/tomhoule-tlsluvktrpnv

protoc-gen-grafbase-subgraph: proper @derive support

Author: tomhoule

Cargo.lock

@@ -63,6 +63,7 @@ clap = "4.5.36"
 ctor = "0.4"
 duration-str = "0.17"
 enumflags2 = "0.7.11"
+field-selection-map = { git = "https://github.com/grafbase/grafbase", package = "engine-field-selection-map" }
 futures = "0.3"
 futures-util = "0.3.31"
 fxhash = "0.2.1"

Cargo.toml

@@ -2,6 +2,17 @@
 
 ### Added
 
+- **Input argument directives** added. You can now add GraphQL directives to RPC method input arguments using the `input_argument_directives` option on methods.
+
+- **Composite schema entity references with @derive** added. You can now create federation-style entity references using the `derive` option on messages:
+  - Use `option (grafbase.graphql.derive) = {entity: "User", is: "{ id: user_id }"};` on fields
+  - Automatically generates reference fields with `@derive` and `@is` directives
+  - Creates stub entity types with `@key` directives if the type is not already defined
+  - Supports custom relation field names with the `field` parameter
+  - The `is` parameter defines the field mapping using format `"{ <entity_key_field>: <proto_field> }"`
+  - The `@is` directive uses the value from the `is` parameter directly
+  - Enables cross-subgraph entity references in federated schemas
+
 - **Multiple subgraphs support** added. Support for generating multiple GraphQL files based on service annotations:
 
   - Services can now have a `subgraph_name` option that maps them to different subgraph files

cli/protoc-gen-grafbase-subgraph/CHANGELOG.md

@@ -8,11 +8,16 @@ keywords.workspace = true
 repository.workspace = true
 
 [dependencies]
+anyhow.workspace = true
+field-selection-map.workspace = true
 indexmap.workspace = true
 paste.workspace = true
 protobuf.workspace = true
 protobuf-support.workspace = true
 
+[build-dependencies]
+protobuf-codegen = "3"
+
 [dev-dependencies]
 insta.workspace = true
 tempfile.workspace = true

cli/protoc-gen-grafbase-subgraph/Cargo.toml

@@ -106,13 +106,13 @@ import "grafbase/options.proto";
 
 service MyService {
   option (grafbase.graphql.schema_directives) = "@contact(name: \"API Support\", url: \"https://api.example.com/support\")";
-  
+
   rpc GetItem(GetItemRequest) returns (Item);
 }
 
 service AnotherService {
   option (grafbase.graphql.schema_directives) = "@tag(name: \"backend\") @auth(rules: [{allow: \"admin\"}])";
-  
+
   rpc UpdateItem(UpdateItemRequest) returns (Item);
 }
 ```
@@ -133,6 +133,69 @@ extend schema
 - **Multi-file mode**: When using `subgraph_name`, only directives from services in that subgraph are included
 - **Multiple directives**: You can include multiple directives in a single string, separated by spaces
 
+### Composite schema entity references
+
+You can create federation-style entity references using the `derive` option on message fields. This allows you to reference entities from other subgraphs:
+
+```protobuf
+import "grafbase/options.proto";
+
+message Product {
+  // Basic usage: creates a user field that references User entity by id
+  (grafbase.graphql.derive) = {
+    entity: "User",
+    is: "{ id: user_id }"
+  };
+
+  // Custom relation field name: creates an owner field instead of user
+  (grafbase.graphql.derive) = {
+    entity: "User",
+    field: "owner"
+    is: "{ id: owner_id }"
+  };
+
+  // Reference by non-id field
+  (grafbase.graphql.composite_schemas_entity) = {
+    entity: "Shop",
+    is: "{ slug: shop_slug }"
+  };
+
+  string id = 1;
+  string name = 2;
+
+  string user_id = 3;
+  string owner_id = 4;
+
+  string shop_slug = 6;
+}
+```
+
+This generates GraphQL with entity references:
+
+```graphql
+type Product {
+  id: String
+  name: String
+  user_id: String
+  user: User @derive @is(field: "{ id: user_id }")  # Automatically added reference field
+  owner_id: String
+  owner: User @derive @is(field: "{ id: owner_id }")  # Custom name for the reference
+  shop_slug: String
+  shop: Shop @derive @is(field: "{ slug: shop_slug }")
+}
+
+# Stub entities are automatically created
+type User @key(fields: "id") {
+  id: String
+}
+
+type Shop @key(fields: "slug") {
+  slug: String
+}
+```
+
+Composite (multiple fields) and list derives are also supported.
+
 ### Mapping specific services to different subgraphs
 
 By default, the plugin generates a single `schema.graphql` file containing all services. However, you can map different services to different subgraph files using the `subgraph_name` option:

cli/protoc-gen-grafbase-subgraph/README.md

@@ -0,0 +1,12 @@
+use std::io::Result;
+
+fn main() -> Result<()> {
+    protobuf_codegen::Codegen::new()
+        .pure()
+        .cargo_out_dir("protos")
+        .include("proto")
+        .input("proto/grafbase/options.proto")
+        .run_from_script();
+
+    Ok(())
+}

cli/protoc-gen-grafbase-subgraph/build.rs

@@ -4,9 +4,16 @@ import "google/protobuf/descriptor.proto";
 
 package grafbase.graphql;
 
+message CompositeSchemaEntity {
+  optional string entity = 1;
+  optional string field = 2;
+  optional string is = 3;
+}
+
 extend google.protobuf.MessageOptions {
   optional string object_directives = 58301;
   optional string input_object_directives = 58302;
+  repeated CompositeSchemaEntity derive = 58303;
 }
 
 extend google.protobuf.FieldOptions {

cli/protoc-gen-grafbase-subgraph/proto/grafbase/options.proto

@@ -1,4 +1,5 @@
 mod display_utils;
+mod options_proto;
 mod render_graphql_sdl;
 mod schema;
 mod translate_schema;
@@ -70,7 +71,8 @@ fn generate(raw_request: &[u8]) -> CodeGeneratorResponse {
         // Generate a file for each subgraph
         for (subgraph_name, service_ids) in services_by_subgraph {
             let mut graphql_schema = String::new();
-            render_graphql_sdl_filtered(&translated_schema, Some(&service_ids), &mut graphql_schema).unwrap();
+            render_graphql_sdl_filtered(&translated_schema, Some(&service_ids), &mut graphql_schema)
+                .expect("Failed to render GraphQL schema");
 
             let mut file = File::new();
             file.set_name(format!("{}.graphql", subgraph_name));
@@ -80,7 +82,7 @@ fn generate(raw_request: &[u8]) -> CodeGeneratorResponse {
     } else {
         // Single-file mode: generate schema.graphql
         let mut graphql_schema = String::new();
-        render_graphql_sdl(&translated_schema, &mut graphql_schema).unwrap();
+        render_graphql_sdl(&translated_schema, &mut graphql_schema).expect("Failed to render GraphQL schema");
 
         let mut file = File::new();
         file.set_name("schema.graphql".to_owned());

cli/protoc-gen-grafbase-subgraph/src/main.rs

@@ -0,0 +1 @@
+include!(concat!(env!("OUT_DIR"), "/protos/mod.rs"));

cli/protoc-gen-grafbase-subgraph/src/options_proto.rs

@@ -1,6 +1,5 @@
-use crate::schema::ScalarType;
-
 use super::*;
+use crate::schema::{ProtoMessage, ScalarType, View};
 
 pub(super) fn render_graphql_types(
     schema: &GrpcSchema,
@@ -28,6 +27,8 @@ pub(super) fn render_graphql_types(
         render_enum_definition(schema, *enum_id, f)?;
     }
 
+    render_entity_types(schema, messages_to_render_as_output, f)?;
+
     Ok(())
 }
 
@@ -95,7 +96,7 @@ fn render_message(
         if input {
             render_input_field_type(schema, &field.r#type, field.repeated, f)?;
         } else {
-            render_output_field_type(schema, &field.r#type, field.repeated, f)?;
+            render_output_field_type(schema, &field.r#type, field.repeated, true, f)?;
         }
 
         let field_directives = if input {
@@ -112,13 +113,54 @@ fn render_message(
         f.write_str("\n")?;
     }
 
+    if !input {
+        render_derives(message, f)?;
+    }
+
     f.write_str("}\n")
 }
 
+fn render_derives(message: View<'_, ProtoMessageId, ProtoMessage>, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+    for entity_info in &message.derives {
+        f.write_str(INDENT)?;
+
+        let derived_field_name = if let Some(name) = &entity_info.field {
+            name.clone()
+        } else {
+            entity_info
+                .is
+                .as_ref()
+                .and_then(|is| is.fields.first())
+                .map(|first_field| {
+                    first_field
+                        .input_field_name
+                        .trim_end_matches("_id")
+                        .trim_end_matches("Id")
+                        .to_owned()
+                })
+                .unwrap_or_else(|| entity_info.entity.to_lowercase())
+        };
+
+        f.write_str(&derived_field_name)?;
+        f.write_str(": ")?;
+        f.write_str(&entity_info.entity)?;
+        f.write_str(" @derive")?;
+
+        if let Some(is) = &entity_info.is {
+            write!(f, " @is(field: \"{is}\")")?;
+        }
+
+        f.write_str("\n")?;
+    }
+
+    Ok(())
+}
+
 pub(super) fn render_output_field_type(
     schema: &GrpcSchema,
     ty: &FieldType,
     repeated: bool,
+    optional: bool,
     f: &mut fmt::Formatter<'_>,
 ) -> fmt::Result {
     if repeated {
@@ -133,6 +175,10 @@ pub(super) fn render_output_field_type(
 
     if repeated {
         f.write_str("!]")?;
+    } else if !optional && matches!(ty, FieldType::Scalar(_) | FieldType::Enum(_)) {
+        // Only scalar and enum types can be non-null based on optional flag
+        // Message types are always nullable
+        f.write_str("!")?;
     }
 
     Ok(())
@@ -235,3 +281,104 @@ fn render_message_type_name(
         }
     }
 }
+
+fn render_entity_types(
+    schema: &GrpcSchema,
+    messages_to_render_as_output: &BTreeSet<ProtoMessageId>,
+    f: &mut fmt::Formatter<'_>,
+) -> fmt::Result {
+    use std::collections::{HashMap, HashSet};
+
+    // entity name -> (list of key fields, message ID for field type lookup)
+    let mut entities: HashMap<String, (Vec<&str>, ProtoMessageId)> = HashMap::new();
+
+    let mut rendered_output_types: HashSet<String> = HashSet::new();
+    for message_id in messages_to_render_as_output {
+        let message = &schema[*message_id];
+        rendered_output_types.insert(message.graphql_output_name().to_string());
+    }
+
+    for message_id in messages_to_render_as_output {
+        let message = &schema[*message_id];
+        for entity_info in &message.derives {
+            if rendered_output_types.contains(&entity_info.entity) {
+                continue;
+            }
+
+            match &entity_info.is {
+                Some(simple_is) => {
+                    let key_fields: Vec<&str> = simple_is.fields.iter().map(|f| f.output_field_name.as_str()).collect();
+
+                    entities.insert(entity_info.entity.clone(), (key_fields, *message_id));
+                }
+                None => {
+                    entities.insert(entity_info.entity.clone(), (vec!["id"], *message_id));
+                }
+            }
+        }
+    }
+
+    let mut sorted_entities: Vec<_> = entities.into_iter().collect();
+    sorted_entities.sort_by_key(|(entity_name, _)| entity_name.clone());
+
+    for (entity_name, (key_fields, message_id)) in sorted_entities {
+        f.write_str("\n")?;
+
+        // Format the @key directive
+        if key_fields.len() == 1 {
+            write!(f, "type {} @key(fields: \"{}\")", entity_name, key_fields[0])?;
+        } else {
+            // For composite keys, use space-separated format
+            write!(f, "type {} @key(fields: \"", entity_name)?;
+            for (i, field) in key_fields.iter().enumerate() {
+                if i > 0 {
+                    f.write_str(" ")?;
+                }
+                f.write_str(field)?;
+            }
+            f.write_str("\")")?;
+        }
+
+        f.write_str(" {\n")?;
+
+        let message = &schema[message_id];
+        let matching_entity_info = message.derives.iter().find(|info| info.entity == entity_name);
+
+        if let Some(entity_info) = matching_entity_info {
+            if let Some(simple_is) = &entity_info.is {
+                // Render each key field with its corresponding type from the message
+                for is_field in &simple_is.fields {
+                    f.write_str(INDENT)?;
+                    f.write_str(&is_field.output_field_name)?;
+                    f.write_str(": ")?;
+
+                    // Find the corresponding field in the message
+                    let field_found = message_id.fields(schema).find(|f| f.name == is_field.input_field_name);
+
+                    if let Some(field) = field_found {
+                        render_output_field_type(schema, &field.r#type, false, true, f)?;
+                    } else {
+                        // Default to String if field not found
+                        f.write_str("String")?;
+                    }
+                    f.write_str("\n")?;
+                }
+            } else {
+                // No is mapping, default to id: String
+                f.write_str(INDENT)?;
+                f.write_str("id: String\n")?;
+            }
+        } else {
+            // No entity info, render all key fields as String
+            for field in &key_fields {
+                f.write_str(INDENT)?;
+                f.write_str(field)?;
+                f.write_str(": String\n")?;
+            }
+        }
+
+        f.write_str("}\n")?;
+    }
+
+    Ok(())
+}