Add grains collector and state aggregate processor

Author: MKLeb

setup.cfg

@@ -53,7 +53,7 @@ saf.collect =
   event-bus = saf.collect.event_bus
   file = saf.collect.file
   salt_exec = saf.collect.salt_exec
-  salt_cmd = saf.collect.salt_cmd
+  grains = saf.collect.grains
   test = saf.collect.test
 saf.process =
   regex_mask = saf.process.regex_mask

src/saf/collect/salt_cmd.py renamed to src/saf/collect/grains.py

@@ -26,12 +26,9 @@ class SaltCommandConfig(CollectConfigBase):
     Configuration schema for the beacons collect plugin.
     """
 
-    targets: str
-    cmd: str
-    args: List[str] | None
-    kwargs: Dict[str, str] | None
-    interval: float = 5
-    cache_flag: str | None = None
+    targets: str = "*"
+    grains: List[str]
+    interval: float = 20
 
 
 def get_config_schema() -> Type[SaltCommandConfig]:
@@ -41,16 +38,27 @@ def get_config_schema() -> Type[SaltCommandConfig]:
     return SaltCommandConfig
 
 
-async def collect(*, ctx: PipelineRunContext[SaltCommandConfig]) -> AsyncIterator[CollectedEvent]:
+class GrainsCollectedEvent(CollectedEvent):
+    """
+    A collected event surrounding a SaltEvent.
+    """
+
+    minion: str
+    grains: Dict[str, str]
+
+
+async def collect(
+    *, ctx: PipelineRunContext[SaltCommandConfig]
+) -> AsyncIterator[GrainsCollectedEvent]:
     """
     Method called to collect events.
     """
     config = ctx.config
     client = LocalClient(mopts=ctx.salt_config.copy())
 
     while True:
-        ret = client.cmd(config.targets, config.cmd, arg=config.args, kwarg=config.kwargs)
-        event = CollectedEvent(data={config._name: ret})  # noqa: SLF001
-        log.debug("CollectedEvent: %s", event)
-        yield event
+        ret = client.cmd(config.targets, "grains.item", arg=config.grains)
+        for minion, grains in ret.items():
+            event = GrainsCollectedEvent(data=ret, minion=minion, grains=grains)
+            yield event
         await asyncio.sleep(config.interval)

src/saf/process/state_aggregate.py

@@ -9,9 +9,11 @@
 import logging
 from typing import TYPE_CHECKING
 from typing import AsyncIterator
+from typing import Dict
 from typing import Type
 
 from saf.collect.event_bus import EventBusCollectedEvent
+from saf.collect.grains import GrainsCollectedEvent
 from saf.models import CollectedEvent
 from saf.models import PipelineRunContext
 from saf.models import ProcessConfigBase
@@ -45,6 +47,7 @@ class StateAggregateCollectedEvent(CollectedEvent):
     end_time: datetime
     duration: timedelta
     minion_id: str
+    grains: Dict[str, str]
 
 
 async def process(
@@ -80,10 +83,29 @@ async def process(
                 start_time = job_start_event.stamp
                 end_time = salt_event.stamp
                 duration = end_time - start_time
-                yield StateAggregateCollectedEvent.construct(
+                grains = ctx.cache.get("grains", {}).get(minion_id, {})
+                ret = StateAggregateCollectedEvent.construct(
                     data=data,
                     start_time=start_time,
                     end_time=end_time,
                     duration=duration,
                     minion_id=minion_id,
+                    grains=grains,
                 )
+                if grains:
+                    yield ret
+                else:
+                    if "waiting_for_grains" not in ctx.cache:
+                        ctx.cache["waiting_for_grains"] = set()
+                    ctx.cache["waiting_for_grains"].add(ret)
+    elif isinstance(event, GrainsCollectedEvent):
+        if "grains" not in ctx.cache:
+            ctx.cache["grains"] = {}
+        ctx.cache["grains"][event.minion] = event.grains
+        waiting = ctx.cache.get("waiting_for_grains")
+        if waiting:
+            to_remove = [agg_event for agg_event in waiting if agg_event.minion_id == event.minion]
+            for event_with_grains in to_remove:
+                event_with_grains.grains = event.grains
+                waiting.remove(event_with_grains)
+                yield event_with_grains