[core] feat: add support for proto message parser

Author: AndreyMarkinPPC

libs/core/garf_core/__init__.py

@@ -26,4 +26,4 @@
   'ApiReportFetcher',
 ]
 
-__version__ = '0.5.0'
+__version__ = '0.6.0'

libs/core/garf_core/parsers.py

@@ -65,6 +65,10 @@ def parse_response(
       span.set_attribute('num_results', len(results))
       return results
 
+  @abc.abstractmethod
+  def get_row_element(self, row, key):
+    """Defines how to get a single element from a row."""
+
   def _evalute_virtual_column(
     self,
     fields: list[str],
@@ -138,20 +142,31 @@ def process_customizer(
 
   def _process_customizer_slice(self, row, customizer, field):
     slice_object = customizer.value.slice_literal
-    return [r.get(customizer.value.value) for r in row.get(field)[slice_object]]
+    elements = self.get_row_element(row, field)
+    results = []
+    for element in elements[slice_object]:
+      results.append(self.get_row_element(element, customizer.value.value))
+    return results
 
   def _process_nested_field(self, row, customizer, field):
-    nested_field = row.get(field)
+    nested_field = self.get_row_element(row, field)
+    if isinstance(nested_field, MutableSequence):
+      return list(
+        {
+          self.parse_row_element(field, customizer.value)
+          for field in nested_field
+        }
+      )
     try:
-      return operator.attrgetter(customizer.value)(nested_field)
+      return self.parse_row_element(field, customizer.value)
     except AttributeError as e:
       raise query_parser.GarfCustomizerError(
         f'nested field {customizer.value} is missing in row {row}'
       ) from e
 
   def _process_resource_index(self, row, customizer, field):
-    resource = row.get(field, '/')
-    _, *elements = row.get(field, '/').split('/')
+    resource = self.get_row_element(row, field)
+    _, *elements = resource.split('/')
     if not elements:
       raise query_parser.GarfCustomizerError(
         f'Not a valid resource: {resource}'
@@ -213,6 +228,10 @@ def parse_row_element(
     except (TypeError, KeyError):
       return None
 
+  def get_row_element(self, row, key):
+    """Gets element from a dict by key."""
+    return row.get(key)
+
 
 class NumericConverterDictParser(DictParser):
   """Extracts nested dict elements with numerical conversions."""
@@ -241,5 +260,24 @@ def convert_field(value):
       return None
 
 
+class ProtoParser(BaseParser):
+  """Extracts attribute from Protobuf messages."""
+
+  def parse_row_element(
+    self, row: api_clients.ApiResponseRow, key: str
+  ) -> api_clients.ApiRowElement:
+    """Returns attributes from a Protobuf message based on a key."""
+    try:
+      return operator.attrgetter(key)(row)
+    except AttributeError as e:
+      raise query_parser.GarfFieldError(
+        f'field {key} is missing in row {row}'
+      ) from e
+
+  def get_row_element(self, row, key):
+    """Gets nested attribute from a Protobuf message."""
+    return operator.attrgetter(key)(row)
+
+
 class GarfParserError(exceptions.GarfError):
   """Incorrect data format for parser."""

libs/core/tests/unit/test_parsers.py

@@ -25,6 +25,18 @@ class NestedResource(pydantic.BaseModel):
   nested_element: int
 
 
+class ArrayElement(pydantic.BaseModel):
+  element: int
+
+
+class FakeProtoMessage(pydantic.BaseModel):
+  resource: str
+  resource_id: int
+  resource_name: str
+  resource_data: NestedResource
+  array_data: list[ArrayElement]
+
+
 class TestDictParser:
   @pytest.fixture
   def test_parser(self):
@@ -174,38 +186,6 @@ def test_parse_response_raises_customizer_error_on_invalid_resource(self):
     ):
       test_parser.parse_response(test_response)
 
-  def test_parse_response_returns_correct_nested_attribute(self):
-    spec = query_editor.QuerySpecification(
-      text='SELECT resource:nested_element AS column FROM test'
-    ).generate()
-    test_parser = parsers.DictParser(spec)
-    test_response = api_clients.GarfApiResponse(
-      results=[
-        {'resource': NestedResource(nested_element=1)},
-        {'resource': NestedResource(nested_element=2)},
-      ]
-    )
-    parsed_row = test_parser.parse_response(test_response)
-    assert parsed_row == [[1], [2]]
-
-  def test_parse_response_raises_customizer_error_on_missing_nested_attribute(
-    self,
-  ):
-    spec = query_editor.QuerySpecification(
-      text='SELECT resource:missing_element AS column FROM test'
-    ).generate()
-    test_parser = parsers.DictParser(spec)
-    test_response = api_clients.GarfApiResponse(
-      results=[
-        {'resource': NestedResource(nested_element=1)},
-      ]
-    )
-    with pytest.raises(
-      query_parser.GarfCustomizerError,
-      match='nested field missing_element is missing in row',
-    ):
-      test_parser.parse_response(test_response)
-
   def test_parse_response_skips_omitted_columns(self):
     test_specification = query_editor.QuerySpecification(
       'SELECT test_column_1 AS _, test_column_2 FROM test'
@@ -237,3 +217,58 @@ def test_parse_row_returns_converted_numeric_values(self, test_parser):
     expected_row = [1]
 
     assert parsed_row == expected_row
+
+
+class TestProtoParser:
+  @pytest.fixture
+  def test_parser(self):
+    spec = query_editor.QuerySpecification(
+      """
+      SELECT
+        resource_id,
+        resource_id + 1 AS next_resource_id,
+        resource_name,
+        resource_data.nested_element,
+        array_data[0].element AS slice_element
+      FROM test
+        """
+    ).generate()
+    return parsers.ProtoParser(spec)
+
+  def test_parse_row_returns_converted_numeric_values(self, test_parser):
+    test_row = FakeProtoMessage(
+      resource='resources/1/resource/99',
+      resource_id=1,
+      resource_name='test',
+      resource_data=NestedResource(nested_element=10),
+      array_data=[ArrayElement(element=100)],
+    )
+
+    parsed_row = test_parser.parse_row(test_row)
+    expected_row = [1, 2, 'test', 10, [100]]
+
+    assert parsed_row == expected_row
+
+  def test_parse_response_raises_customizer_error_on_missing_nested_attribute(
+    self,
+  ):
+    spec = query_editor.QuerySpecification(
+      text='SELECT resource_data:missing_element AS column FROM test'
+    ).generate()
+    test_parser = parsers.ProtoParser(spec)
+    test_response = api_clients.GarfApiResponse(
+      results=[
+        FakeProtoMessage(
+          resource='resources/1/resource/99',
+          resource_id=1,
+          resource_name='test',
+          resource_data=NestedResource(nested_element=10),
+          array_data=[ArrayElement(element=100)],
+        )
+      ]
+    )
+    with pytest.raises(
+      query_parser.GarfFieldError,
+      match='field missing_element is missing in row resource_data',
+    ):
+      test_parser.parse_response(test_response)