Initial check in for job aggregations

Author: MKLeb

setup.cfg

@@ -50,13 +50,16 @@ salt.loader =
   salt-analytics-framework = saf.saltext
 saf.collect =
   beacons = saf.collect.beacons
+  event-bus = saf.collect.event_bus
   file = saf.collect.file
   salt_exec = saf.collect.salt_exec
+  salt_cmd = saf.collect.salt_cmd
   test = saf.collect.test
 saf.process =
   regex_mask = saf.process.regex_mask
   shannon_mask = saf.process.shannon_mask
   jupyter_notebook = saf.process.jupyter_notebook
+  state-aggregate = saf.process.state_aggregate
   test = saf.process.test
 saf.forward =
   disk = saf.forward.disk

src/saf/collect/event_bus.py

@@ -0,0 +1,55 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Collect events from the event bus.
+"""
+from __future__ import annotations
+
+import logging
+from typing import AsyncIterator
+from typing import Set
+from typing import Type
+
+from saf.models import CollectConfigBase
+from saf.models import CollectedEvent
+from saf.models import PipelineRunContext
+from saf.models import SaltEvent
+from saf.utils import eventbus
+
+log = logging.getLogger(__name__)
+
+
+class EventBusConfig(CollectConfigBase):
+    """
+    Configuration schema for the beacons collect plugin.
+    """
+
+    tags: Set[str]
+
+
+def get_config_schema() -> Type[EventBusConfig]:
+    """
+    Get the event bus collect plugin configuration schema.
+    """
+    return EventBusConfig
+
+
+class EventBusCollectedEvent(CollectedEvent):
+    """
+    A collected event surrounding a SaltEvent.
+    """
+
+    salt_event: SaltEvent
+
+
+async def collect(
+    *, ctx: PipelineRunContext[EventBusConfig]
+) -> AsyncIterator[EventBusCollectedEvent]:
+    """
+    Method called to collect events.
+    """
+    config = ctx.config
+    salt_event: SaltEvent
+    log.info("The event bus collect plugin is configured to listen to tags: %s", config.tags)
+    async for salt_event in eventbus.iter_events(opts=ctx.salt_config.copy(), tags=config.tags):
+        yield EventBusCollectedEvent(salt_event=salt_event, data={"tag": salt_event.tag})

src/saf/collect/salt_cmd.py

@@ -0,0 +1,56 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Collect events from the event bus.
+"""
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import AsyncIterator
+from typing import Dict
+from typing import List
+from typing import Type
+
+from salt.client import LocalClient
+
+from saf.models import CollectConfigBase
+from saf.models import CollectedEvent
+from saf.models import PipelineRunContext
+
+log = logging.getLogger(__name__)
+
+
+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
+
+
+def get_config_schema() -> Type[SaltCommandConfig]:
+    """
+    Get the event bus collect plugin configuration schema.
+    """
+    return SaltCommandConfig
+
+
+async def collect(*, ctx: PipelineRunContext[SaltCommandConfig]) -> AsyncIterator[CollectedEvent]:
+    """
+    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
+        await asyncio.sleep(config.interval)

src/saf/process/state_aggregate.py

@@ -0,0 +1,89 @@
+# Copyright 2021-2023 VMware, Inc.
+# SPDX-License-Identifier: Apache-2.0
+"""
+Aggregate the necessary job info into one event to be forwarded.
+"""
+from __future__ import annotations
+
+import fnmatch
+import logging
+from typing import TYPE_CHECKING
+from typing import AsyncIterator
+from typing import Type
+
+from saf.collect.event_bus import EventBusCollectedEvent
+from saf.models import CollectedEvent
+from saf.models import PipelineRunContext
+from saf.models import ProcessConfigBase
+
+if TYPE_CHECKING:
+    from datetime import datetime
+    from datetime import timedelta
+
+log = logging.getLogger(__name__)
+
+
+class StateAggregateConfig(ProcessConfigBase):
+    """
+    Job aggregate collector configuration.
+    """
+
+
+def get_config_schema() -> Type[StateAggregateConfig]:
+    """
+    Get the job aggregate collect plugin configuration schema.
+    """
+    return StateAggregateConfig
+
+
+class StateAggregateCollectedEvent(CollectedEvent):
+    """
+    A collected event with aggregated state run information.
+    """
+
+    start_time: datetime
+    end_time: datetime
+    duration: timedelta
+    minion_id: str
+
+
+async def process(
+    *,
+    ctx: PipelineRunContext[StateAggregateConfig],
+    event: CollectedEvent,
+) -> AsyncIterator[CollectedEvent]:
+    """
+    Aggregate received events, otherwise store in cache.
+    """
+    if isinstance(event, EventBusCollectedEvent):
+        salt_event = event.salt_event
+        tag = salt_event.tag
+        data = salt_event.data
+        if fnmatch.fnmatch(tag, "salt/job/*/new"):
+            # We will probably want to make this condition configurable
+            if TYPE_CHECKING:
+                assert isinstance(salt_event.data, dict)
+            if data.get("fun") == "state.apply":
+                jid = tag.split("/")[2]
+                if "watched_jids" not in ctx.cache:
+                    ctx.cache["watched_jids"] = {}
+                # We are going to want a TTL at some point for the watched jids
+                ctx.cache["watched_jids"][jid] = salt_event
+        elif fnmatch.fnmatch(tag, "salt/job/*/ret/*"):
+            split_tag = tag.split("/")
+            jid = split_tag[2]
+            if "watched_jids" not in ctx.cache:
+                ctx.cache["watched_jids"] = {}
+            if jid in ctx.cache["watched_jids"]:
+                job_start_event = ctx.cache["watched_jids"][jid]
+                minion_id = split_tag[-1]
+                start_time = job_start_event.stamp
+                end_time = salt_event.stamp
+                duration = end_time - start_time
+                yield StateAggregateCollectedEvent.construct(
+                    data=data,
+                    start_time=start_time,
+                    end_time=end_time,
+                    duration=duration,
+                    minion_id=minion_id,
+                )

src/saf/utils/eventbus.py

@@ -37,10 +37,15 @@ def _construct_event(tag: str, stamp: datetime, event_data: dict[str, Any]) -> S
         for key in list(event_data):
             if key.startswith("_"):
                 event_data.pop(key)
+        event_data_value = event_data.get("data") or event_data
+        if isinstance(event_data_value, str):
+            data = {"return": event_data_value}
+        else:
+            data = event_data_value
         salt_event = SaltEvent(
             tag=tag,
             stamp=stamp,
-            data=event_data.get("data") or event_data,
+            data=data,
             raw_data=event_raw_data,
         )
         log.debug("Constructed SaltEvent: %s", salt_event)