feat(protobuf): add source tracking for query types and flattened fields

Signed-off-by: Ahmed Mohamed <[email protected]>

Author: ahmdmhd

src/s2dm/cli.py

@@ -587,7 +587,7 @@ def jsonschema(
 def protobuf(
     ctx: click.Context,
     schemas: list[Path],
-    selection_query: Path | None,
+    selection_query: Path,
     output: Path,
     root_type: str | None,
     flatten_naming: bool,
@@ -598,13 +598,11 @@ def protobuf(
     naming_config = ctx.obj.get("naming_config")
     graphql_schema = load_schema_with_naming(schemas, naming_config)
 
-    query_document = None
-    if selection_query:
-        query_document = parse(selection_query.read_text())
-        graphql_schema = prune_schema_using_query_selection(graphql_schema, query_document)
+    query_document = parse(selection_query.read_text())
+    graphql_schema = prune_schema_using_query_selection(graphql_schema, query_document)
 
     result = translate_to_protobuf(
-        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, query_document
+        graphql_schema, query_document, root_type, flatten_naming, package_name, naming_config, expanded_instances
     )
     _ = output.write_text(result)
 

src/s2dm/exporters/protobuf/models.py

@@ -74,4 +74,3 @@ class ProtoSchema(BaseModel):
     messages: list[ProtoMessage] = Field(default_factory=list)
     unions: list[ProtoUnion] = Field(default_factory=list)
     flatten_mode: bool = False
-    flattened_fields: list[ProtoField] = Field(default_factory=list)

src/s2dm/exporters/protobuf/protobuf.py

@@ -9,12 +9,12 @@
 
 def transform(
     graphql_schema: GraphQLSchema,
+    selection_query: DocumentNode,
     root_type: str | None = None,
     flatten_naming: bool = False,
     package_name: str | None = None,
     naming_config: dict[str, Any] | None = None,
     expanded_instances: bool = False,
-    selection_query: DocumentNode | None = None,
 ) -> str:
     """
     Transform a GraphQL schema object to Protocol Buffers format.
@@ -26,10 +26,13 @@ def transform(
         package_name: Optional package name for the .proto file
         naming_config: Optional naming configuration
         expanded_instances: If True, expand instance tags into nested structures
-        selection_query: Optional selection query document to determine root-level types
+        selection_query: Required selection query document to determine root-level types
 
     Returns:
         str: Protocol Buffers representation as a string
+
+    Raises:
+        ValueError: If selection_query is not provided
     """
     log.info(f"Transforming GraphQL schema to Protobuf with {len(graphql_schema.type_map)} types")
 
@@ -39,7 +42,7 @@ def transform(
         log.info(f"Using root type: {root_type}")
 
     transformer = ProtobufTransformer(
-        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, selection_query
+        graphql_schema, selection_query, root_type, flatten_naming, package_name, naming_config, expanded_instances
     )
     proto_content = transformer.transform()
 
@@ -50,12 +53,12 @@ def transform(
 
 def translate_to_protobuf(
     schema: GraphQLSchema,
+    selection_query: DocumentNode,
     root_type: str | None = None,
     flatten_naming: bool = False,
     package_name: str | None = None,
     naming_config: dict[str, Any] | None = None,
     expanded_instances: bool = False,
-    selection_query: DocumentNode | None = None,
 ) -> str:
     """
     Translate a GraphQL schema to Protocol Buffers format.
@@ -67,11 +70,14 @@ def translate_to_protobuf(
         package_name: Optional package name for the .proto file
         naming_config: Optional naming configuration
         expanded_instances: If True, expand instance tags into nested structures
-        selection_query: Optional selection query document to determine root-level types
+        selection_query: Required selection query document to determine root-level types
 
     Returns:
         str: Protocol Buffers (.proto) representation as a string
+
+    Raises:
+        ValueError: If selection_query is not provided
     """
     return transform(
-        schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, selection_query
+        schema, selection_query, root_type, flatten_naming, package_name, naming_config, expanded_instances
     )

src/s2dm/exporters/protobuf/templates/proto_flattened.j2

@@ -82,13 +82,16 @@ class ProtobufTransformer:
     def __init__(
         self,
         graphql_schema: GraphQLSchema,
+        selection_query: DocumentNode,
         root_type: str | None = None,
         flatten_naming: bool = False,
         package_name: str | None = None,
         naming_config: dict[str, Any] | None = None,
         expanded_instances: bool = False,
-        selection_query: DocumentNode | None = None,
     ):
+        if selection_query is None:
+            raise ValueError("selection_query is required")
+
         self.graphql_schema = graphql_schema
         self.root_type = root_type
         self.flatten_naming = flatten_naming
@@ -151,7 +154,7 @@ def transform(self) -> str:
             # If no fields reference that object type directly (non-flattened), the type definition is no longer needed.
             # However, unions and enums cannot be flattened and must remain as separate type definitions.
             (
-                proto_schema.flattened_fields,
+                flattened_fields,
                 referenced_type_names,
                 flattened_root_types,
             ) = self._build_flattened_fields(message_types)
@@ -167,7 +170,17 @@ def transform(self) -> str:
         proto_schema.unions = self._build_unions(union_types)
         proto_schema.messages = self._build_messages(message_types)
 
-        template_name = "proto_flattened.j2" if self.flatten_naming else "proto_standard.j2"
+        if self.flatten_naming:
+            root_message_name = self._get_query_operation_name()
+            root_message_source = f"query: {root_message_name}"
+            root_message = ProtoMessage(
+                name=root_message_name,
+                fields=flattened_fields,
+                source=root_message_source,
+            )
+            proto_schema.messages.append(root_message)
+
+        template_name = "proto_standard.j2"
         template = self.env.get_template(template_name)
 
         template_vars = self._build_template_vars(proto_schema)
@@ -185,9 +198,7 @@ def check_message(message: ProtoMessage) -> bool:
                 return True
             return any(check_message(nested) for nested in message.nested_messages)
 
-        return any(field.validation_rules for field in proto_schema.flattened_fields) or any(
-            check_message(message) for message in proto_schema.messages
-        )
+        return any(check_message(message) for message in proto_schema.messages)
 
     def _has_source_option(self, proto_schema: ProtoSchema) -> bool:
         """Check if any type in the schema has a source option."""
@@ -199,9 +210,6 @@ def _get_query_operation_name(self) -> str:
         """Extract the operation name from the selection query, defaulting to appropriate fallback."""
         default_name = "Message" if self.flatten_naming else "Query"
 
-        if not self.selection_query:
-            return default_name
-
         for definition in self.selection_query.definitions:
             if not isinstance(definition, OperationDefinitionNode) or definition.operation != OperationType.QUERY:
                 continue
@@ -226,7 +234,6 @@ def _build_template_vars(self, proto_schema: ProtoSchema) -> dict[str, Any]:
         template_vars = proto_schema.model_dump()
         template_vars["imports"] = imports
         template_vars["has_source_option"] = has_source_option
-        template_vars["message_name"] = self._get_query_operation_name()
 
         return template_vars
 
@@ -259,15 +266,18 @@ def _build_messages(self, message_types: list[GraphQLObjectType | GraphQLInterfa
             fields, nested_messages = self._build_message_fields(message_type)
 
             message_name = message_type.name
+            source = message_type.name
+
             if message_type.name == "Query":
                 message_name = self._get_query_operation_name()
+                source = f"query: {message_name}"
 
             messages.append(
                 ProtoMessage(
                     name=message_name,
                     fields=fields,
                     description=message_type.description,
-                    source=message_type.name,
+                    source=source,
                     nested_messages=nested_messages,
                 )
             )