add correct mcp schema

src/fgai/graylog_source.py

@@ -17,9 +17,19 @@ DEFAULT_FIELD_MAP = {
     "action": ("action", "event_action", "disposition"),
 }
 
+DEFAULT_FIELDS = ["timestamp", "source", "srcip", "dstip", "srcport", "dstport", "service", "action", "severity", "policyid", "subtype", "type", "sentbyte", "rcvdbyte", "hitcount"]
+
 
 def _records(value: object) -> Iterable[dict[str, object]]:
     if isinstance(value, dict):
+        schema = value.get("schema")
+        datarows = value.get("datarows")
+        if isinstance(schema, list) and isinstance(datarows, list):
+            fields = [str(column.get("field", "")) for column in schema if isinstance(column, dict)]
+            for row in datarows:
+                if isinstance(row, list):
+                    yield {field: row[index] for index, field in enumerate(fields) if field and index < len(row)}
+            return
         for key in ("messages", "results", "events", "data"):
             if isinstance(value.get(key), list):
                 yield from (item for item in value[key] if isinstance(item, dict))
@@ -44,7 +54,13 @@ class GraylogStreamSource:
 
     def fetch(self) -> tuple[list[LogEvent], dict[str, object]]:
         status = self.client.probe()
-        arguments: dict[str, object] = {"query": self.query, "limit": 1000, "range_seconds": 300}
+        mapping_fields = [str(value) for value in self.mapping.values() if isinstance(value, str)]
+        arguments: dict[str, object] = {
+            "query": self.query,
+            "limit": 1000,
+            "range_seconds": 300,
+            "fields": list(dict.fromkeys([*DEFAULT_FIELDS, *mapping_fields])),
+        }
         if self.stream:
             arguments["streams"] = [self.stream]
         result = self.client.call_tool("search_messages", arguments)

tests/test_graylog_source.py

@@ -12,7 +12,7 @@ class _Client:
 
     def call_tool(self, _name, arguments):
         self.arguments = arguments
-        return {"result": {"content": [{"type": "text", "text": '{"messages":[{"client":"10.0.0.1","server":"8.8.8.8","result":"deny"}]}' }]}}
+        return {"result": {"content": [{"type": "text", "text": '{"schema":[{"field":"client"},{"field":"server"},{"field":"result"}],"datarows":[["10.0.0.1","8.8.8.8","deny"],["10.0.0.2","8.8.8.8","accept"]]}' }]}}
 
 
 class GraylogSourceTests(unittest.TestCase):
@@ -21,12 +21,13 @@ class GraylogSourceTests(unittest.TestCase):
         events, status = GraylogStreamSource(
             client, "vpn", field_mapping='{"srcip":"client","dstip":"server","action":"result"}'
         ).fetch()
-        self.assertEqual(status["events_fetched"], 1)
+        self.assertEqual(status["events_fetched"], 2)
         self.assertEqual(events[0].src_ip, "10.0.0.1")
         self.assertEqual(events[0].dst_ip, "8.8.8.8")
         self.assertEqual(events[0].action, "deny")
         self.assertEqual(client.arguments["streams"], ["vpn"])
         self.assertEqual(client.arguments["range_seconds"], 300)
+        self.assertIn("client", client.arguments["fields"])
 
 
 if __name__ == "__main__":