feat(protobuf): flatten all root-level types from selection query

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

Author: ahmdmhd

src/s2dm/cli.py

@@ -574,7 +574,7 @@ def jsonschema(
     "-f",
     is_flag=True,
     default=False,
-    help="Flatten nested field names. Requires --root-type to be set.",
+    help="Flatten nested field names.",
 )
 @click.option(
     "--package-name",
@@ -598,12 +598,13 @@ 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)
 
     result = translate_to_protobuf(
-        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances
+        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, query_document
     )
     _ = output.write_text(result)
 

src/s2dm/exporters/protobuf/protobuf.py

@@ -1,6 +1,6 @@
 from typing import Any
 
-from graphql import GraphQLSchema
+from graphql import DocumentNode, GraphQLSchema
 
 from s2dm import log
 
@@ -14,6 +14,7 @@ def transform(
     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.
@@ -25,6 +26,7 @@ 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
 
     Returns:
         str: Protocol Buffers representation as a string
@@ -36,12 +38,8 @@ def transform(
             raise ValueError(f"Root type '{root_type}' not found in schema")
         log.info(f"Using root type: {root_type}")
 
-    if flatten_naming and not root_type:
-        log.warning("Flatten naming mode requires a root type, falling back to standard mode")
-        flatten_naming = False
-
     transformer = ProtobufTransformer(
-        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances
+        graphql_schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, selection_query
     )
     proto_content = transformer.transform()
 
@@ -57,6 +55,7 @@ def translate_to_protobuf(
     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.
@@ -68,8 +67,11 @@ 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
 
     Returns:
         str: Protocol Buffers (.proto) representation as a string
     """
-    return transform(schema, root_type, flatten_naming, package_name, naming_config, expanded_instances)
+    return transform(
+        schema, root_type, flatten_naming, package_name, naming_config, expanded_instances, selection_query
+    )

src/s2dm/exporters/protobuf/transformer.py

@@ -1,6 +1,7 @@
 from typing import Any, cast
 
 from graphql import (
+    DocumentNode,
     GraphQLEnumType,
     GraphQLField,
     GraphQLInterfaceType,
@@ -26,7 +27,7 @@
 from s2dm import log
 from s2dm.exporters.protobuf.models import ProtoEnum, ProtoEnumValue, ProtoField, ProtoMessage, ProtoSchema, ProtoUnion
 from s2dm.exporters.utils.directive import get_directive_arguments, has_given_directive
-from s2dm.exporters.utils.extraction import get_all_named_types
+from s2dm.exporters.utils.extraction import get_all_named_types, get_root_level_types_from_query
 from s2dm.exporters.utils.field import get_cardinality
 from s2dm.exporters.utils.instance_tag import expand_instance_tag, get_instance_tag_object, is_instance_tag_field
 from s2dm.exporters.utils.naming import convert_name, get_target_case_for_element
@@ -84,13 +85,15 @@ def __init__(
         package_name: str | None = None,
         naming_config: dict[str, Any] | None = None,
         expanded_instances: bool = False,
+        selection_query: DocumentNode | None = None,
     ):
         self.graphql_schema = graphql_schema
         self.root_type = root_type
         self.flatten_naming = flatten_naming
         self.package_name = package_name
         self.naming_config = naming_config
         self.expanded_instances = expanded_instances
+        self.selection_query = selection_query
 
         self.env = Environment(
             loader=PackageLoader("s2dm.exporters.protobuf", "templates"),
@@ -137,17 +140,23 @@ def transform(self) -> str:
         proto_schema = ProtoSchema(
             package=self.package_name,
             enums=[],
-            flatten_mode=self.flatten_naming and self.root_type is not None,
+            flatten_mode=self.flatten_naming,
         )
 
-        if self.flatten_naming and self.root_type:
+        if self.flatten_naming:
             # In flatten mode, we need a second filtering pass to remove types that were completely flattened.
             # When object fields are flattened, they become prefixed fields in the parent (e.g., parent_child_field).
             # 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, referenced_type_names = self._build_flattened_fields(message_types)
+            (
+                proto_schema.flattened_fields,
+                referenced_type_names,
+                flattened_root_types,
+            ) = self._build_flattened_fields(message_types)
             message_types = [
-                message_type for message_type in message_types if message_type.name in referenced_type_names
+                message_type
+                for message_type in message_types
+                if message_type.name in referenced_type_names and message_type.name not in flattened_root_types
             ]
             union_types = [union_type for union_type in union_types if union_type.name in referenced_type_names]
             enum_types = [enum_type for enum_type in enum_types if enum_type.name in referenced_type_names]
@@ -156,7 +165,7 @@ 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 and self.root_type else "proto_standard.j2"
+        template_name = "proto_flattened.j2" if self.flatten_naming else "proto_standard.j2"
         template = self.env.get_template(template_name)
 
         template_vars = self._build_template_vars(proto_schema)
@@ -312,20 +321,47 @@ def _build_unions(self, union_types: list[GraphQLUnionType]) -> list[ProtoUnion]
 
     def _build_flattened_fields(
         self, message_types: list[GraphQLObjectType | GraphQLInterfaceType]
-    ) -> tuple[list[ProtoField], set[str]]:
-        """Build flattened fields for flatten_naming mode."""
-        root_object = None
-        for message_type in message_types:
-            if message_type.name == self.root_type:
-                root_object = message_type
-                break
+    ) -> tuple[list[ProtoField], set[str], set[str]]:
+        """Build flattened fields for flatten_naming mode.
+
+        Returns:
+            tuple: (flattened_fields, referenced_types, flattened_root_types)
+        """
+        type_cache = {type_def.name: type_def for type_def in message_types}
+
+        if self.root_type:
+            root_object = type_cache.get(self.root_type)
+            if not root_object:
+                log.warning(f"Root type '{self.root_type}' not found, creating empty message")
+                return [], set(), set()
+
+            fields, referenced_types, _ = self._flatten_fields(root_object, root_object.name, message_types, 1)
+            return fields, referenced_types, {self.root_type}
+
+        root_level_type_names = get_root_level_types_from_query(self.graphql_schema, self.selection_query)
+        if not root_level_type_names:
+            log.warning("No root-level types found in selection query, creating empty message")
+            return [], set(), set()
+
+        all_fields: list[ProtoField] = []
+        all_referenced_types: set[str] = set()
+        flattened_root_types: set[str] = set()
+        field_counter = 1
 
-        if not root_object:
-            log.warning(f"Root type '{self.root_type}' not found, creating empty message")
-            return [], set()
+        for type_name in root_level_type_names:
+            root_object = type_cache.get(type_name)
+            if not root_object:
+                log.warning(f"Root-level type '{type_name}' not found in message types")
+                continue
+
+            flattened_root_types.add(type_name)
+            fields, referenced_types, field_counter = self._flatten_fields(
+                root_object, type_name, message_types, field_counter, type_cache
+            )
+            all_fields.extend(fields)
+            all_referenced_types.update(referenced_types)
 
-        fields, referenced_types, _ = self._flatten_fields(root_object, root_object.name, message_types, 1)
-        return fields, referenced_types
+        return all_fields, all_referenced_types, flattened_root_types
 
     def _create_proto_field_with_validation(
         self, field: GraphQLField, field_name: str, proto_type: str, field_number: int

src/s2dm/exporters/utils/extraction.py

@@ -1,4 +1,15 @@
-from graphql import GraphQLNamedType, GraphQLObjectType, GraphQLSchema
+from graphql import (
+    DocumentNode,
+    FieldNode,
+    GraphQLNamedType,
+    GraphQLObjectType,
+    GraphQLSchema,
+    OperationDefinitionNode,
+    OperationType,
+    get_named_type,
+    is_interface_type,
+    is_object_type,
+)
 
 from s2dm.exporters.utils.directive import has_given_directive
 from s2dm.exporters.utils.graphql_type import is_introspection_type
@@ -34,3 +45,40 @@ def get_all_object_types(
 def get_all_objects_with_directive(objects: list[GraphQLObjectType], directive_name: str) -> list[GraphQLObjectType]:
     # TODO: Extend this function to return all objects that have any directive is directive_name is None
     return [o for o in objects if has_given_directive(o, directive_name)]
+
+
+def get_root_level_types_from_query(schema: GraphQLSchema, selection_query: DocumentNode | None) -> list[str]:
+    """Extract root-level type names from the selection query.
+
+    Args:
+        schema: The GraphQL schema
+        selection_query: The selection query document
+
+    Returns:
+        List of type names that are selected at the root level of the query
+    """
+    query_type = schema.query_type
+    if not selection_query or not query_type:
+        return []
+
+    root_type_names: list[str] = []
+
+    for definition in selection_query.definitions:
+        if not isinstance(definition, OperationDefinitionNode) or definition.operation != OperationType.QUERY:
+            continue
+
+        for selection in definition.selection_set.selections:
+            if not isinstance(selection, FieldNode):
+                continue
+
+            field_name = selection.name.value
+            if field_name not in query_type.fields:
+                continue
+
+            field = query_type.fields[field_name]
+            field_type = get_named_type(field.type)
+
+            if is_object_type(field_type) or is_interface_type(field_type):
+                root_type_names.append(field_type.name)
+
+    return root_type_names

tests/test_protobuf.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 
 import pytest
-from graphql import build_schema
+from graphql import build_schema, parse
 
 from s2dm.exporters.protobuf import translate_to_protobuf
 from s2dm.exporters.utils.schema import load_schema_with_naming
@@ -1040,3 +1040,56 @@ def test_complete_proto_file(self) -> None:
             result,
             re.DOTALL,
         ), "Transmission message with enum field and validation"
+
+    def test_flatten_naming_multiple_root_types(self, test_schema_path: list[Path]) -> None:
+        """Test that flatten mode without root_type flattens all root-level types."""
+        graphql_schema = load_schema_with_naming(test_schema_path, None)
+
+        # Selection query that selects vehicle, cabin, and door at the top level
+        query_str = """
+        query {
+            vehicle {
+                doors { isLocked }
+                model
+            }
+            cabin {
+                seats { isOccupied }
+                temperature
+            }
+            door {
+                isLocked
+                position
+                instanceTag { row side }
+            }
+        }
+        """
+        selection_query = parse(query_str)
+
+        result = translate_to_protobuf(
+            graphql_schema, flatten_naming=True, expanded_instances=False, selection_query=selection_query
+        )
+
+        assert re.search(
+            r"message Message \{.*?"
+            r"optional repeated Door Vehicle_doors = 1;.*?"
+            r"optional string Vehicle_model = 2;.*?"
+            r"optional int32 Vehicle_year = 3;.*?"
+            r"optional repeated string Vehicle_features = 4 "
+            r"\[\(buf\.validate\.field\)\.repeated = \{unique: true, min_items: 1, max_items: 10\}\];.*?"
+            r"optional repeated Seat Cabin_seats = 5;.*?"
+            r"optional repeated Door Cabin_doors = 6;.*?"
+            r"optional float Cabin_temperature = 7 \[\(buf\.validate\.field\)\.float = \{gte: -100, lte: 100\}\];.*?"
+            r"optional bool Door_isLocked = 8;.*?"
+            r"optional int32 Door_position = 9 \[\(buf\.validate\.field\)\.int32 = \{gte: 0, lte: 100\}\];.*?"
+            r"RowEnum\.Enum Door_instanceTag_row = 10 \[\(buf\.validate\.field\)\.required = true\];.*?"
+            r"SideEnum\.Enum Door_instanceTag_side = 11 \[\(buf\.validate\.field\)\.required = true\];.*?"
+            r"\}",
+            result,
+            re.DOTALL,
+        ), "Message with flattened fields from all root-level types (Vehicle, Cabin, Door)"
+
+        assert "message Seat {" in result, "Seat message should be included as it's referenced by arrays"
+
+        assert "message Vehicle {" not in result, "Vehicle should be completely flattened"
+        assert "message Cabin {" not in result, "Cabin should be completely flattened"
+        assert "message Door {" not in result, "Door should be completely flattened"