Playbook generator

.python-version

@@ -0,0 +1 @@
+3.10

pyproject.toml

@@ -21,7 +21,7 @@ robusta = "robusta.cli.main:app"
 
 
 [tool.poetry.dependencies]
-python = "^3.8"
+python = "^3.10"
 typer = "^0.4.1"
 colorlog = "^5.0.1"
 pydantic = "^1.8.1"
@@ -112,6 +112,10 @@ all = ["Flask", "grafana-api", "watchdog", "dulwich", "better-exceptions", "Cair
 [tool.poetry.group.dev.dependencies]
 sphinx-jinja = { git = "https://github.com/robusta-dev/sphinx-jinja.git" }
 sphinx-reredirects = "^0.1.1"
+streamlit = "^1.25.0"
+streamlit-pydantic = "^0.6.0"
+streamlit-extras = "^0.3.4"
+streamlit-antd-components = "^0.2.3"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

scripts/custom_streamlit_pydantic.py

@@ -0,0 +1,90 @@
+# based on code in streamlit_pydantic.pydantic_form but with extra modifications for our use case
+import streamlit as st
+from typing import Type, Optional, TypeVar, Callable
+from pydantic import BaseModel, ValidationError, parse_obj_as
+from streamlit_pydantic import pydantic_input
+from streamlit_pydantic.ui_renderer import GroupOptionalFieldsStrategy
+from streamlit_extras.stylable_container import stylable_container
+
+# Define generic type to allow autocompletion for the model fields
+T = TypeVar("T", bound=BaseModel)
+
+
+def modified_pydantic_form(
+    key: str,
+    model: Type[T],
+    submit_label: str = "Submit",
+    clear_on_submit: bool = False,
+    group_optional_fields: GroupOptionalFieldsStrategy = "no",  # type: ignore
+    lowercase_labels: bool = False,
+    ignore_empty_values: bool = False,
+    title: str = None,
+    on_submit: Callable = None
+) -> Optional[T]:
+    """Auto-generates a Streamlit form based on the given (Pydantic-based) input class.
+
+    Args:
+        key (str): A string that identifies the form. Each form must have its own key.
+        model (Type[BaseModel]): The input model. Either a class or instance based on Pydantic `BaseModel` or Python `dataclass`.
+        submit_label (str): A short label explaining to the user what this button is for. Defaults to “Submit”.
+        clear_on_submit (bool): If True, all widgets inside the form will be reset to their default values after the user presses the Submit button. Defaults to False.
+        group_optional_fields (str, optional): If `sidebar`, optional input elements will be rendered on the sidebar.
+            If `expander`,  optional input elements will be rendered inside an expander element. Defaults to `no`.
+        lowercase_labels (bool): If `True`, all input element labels will be lowercased. Defaults to `False`.
+        ignore_empty_values (bool): If `True`, empty values for strings and numbers will not be stored in the session state. Defaults to `False`.
+
+    Returns:
+        Optional[BaseModel]: An instance of the given input class,
+            if the submit button is used and the input data passes the Pydantic validation.
+    """
+    # we can't use pydantic_form because of https://github.com/LukasMasuch/streamlit-pydantic/issues/39
+    # so we design our own container for all elements, just like a form would
+    with stylable_container(
+            key=key,
+            css_styles="""
+                        {
+                            border: 1px solid rgba(49, 51, 63, 0.2);
+                            border-radius: 0.5rem;
+                            padding: calc(1em - 1px);
+                            box-sizing: content-box;
+                        }
+                        """,
+    ):
+        if title is not None:
+            st.subheader(title)
+
+        input_state = pydantic_input(
+            key,
+            model,
+            group_optional_fields=group_optional_fields,
+            lowercase_labels=lowercase_labels,
+            ignore_empty_values=ignore_empty_values,
+        )
+        submit_button = st.button(label=submit_label)
+
+        try:
+            result = None
+            # check if the model is an instance before parsing
+            if isinstance(model, BaseModel):
+                result = parse_obj_as(model.__class__, input_state)
+            else:
+                result = parse_obj_as(model, input_state)
+
+            if submit_button and on_submit is not None:
+                on_submit()
+            return result
+
+        except ValidationError as ex:
+            if not submit_button:
+                return None
+            error_text = "**Whoops! There were some problems with your input:**"
+            for error in ex.errors():
+                if "loc" in error and "msg" in error:
+                    location = ".".join(error["loc"]).replace("__root__.", "")  # type: ignore
+                    error_msg = f"**{location}:** " + error["msg"]
+                    error_text += "\n\n" + error_msg
+                else:
+                    # Fallback
+                    error_text += "\n\n" + str(error)
+            st.error(error_text)
+            return None

scripts/generate_kubernetes_code.py

@@ -338,7 +338,7 @@ def autogenerate_triggers(f: TextIO):
     f.write(
         textwrap.dedent(
             """\
-        from typing import Optional, Dict
+        from typing import Optional, Dict, Literal
         from pydantic import BaseModel
         from robusta.integrations.kubernetes.base_triggers import K8sBaseTrigger
         from robusta.core.model.k8s_operation_type import K8sOperationType
@@ -357,14 +357,8 @@ def autogenerate_triggers(f: TextIO):
                 textwrap.dedent(
                     f"""\
             class {resource}{get_trigger_class_name(trigger_name)}Trigger(K8sBaseTrigger):
-                def __init__(self, name_prefix: str = None, namespace_prefix: str = None, labels_selector: str = None):
-                    super().__init__(
-                        kind=\"{resource}\",
-                        operation={operation_type},
-                        name_prefix=name_prefix,
-                        namespace_prefix=namespace_prefix,
-                        labels_selector=labels_selector,
-                    )
+                kind: Literal[\"{resource}\"] = \"{resource}\"
+                operation: Literal[{operation_type}] = {operation_type}
 
                 @staticmethod
                 def get_execution_event_type() -> type:
@@ -387,14 +381,8 @@ def get_execution_event_type() -> type:
             textwrap.dedent(
                 f"""\
         class KubernetesAny{get_trigger_class_name(trigger_name)}Trigger(K8sBaseTrigger):
-            def __init__(self, name_prefix: str = None, namespace_prefix: str = None, labels_selector: str = None):
-                super().__init__(
-                    kind=\"Any\",
-                    operation={operation_type},
-                    name_prefix=name_prefix,
-                    namespace_prefix=namespace_prefix,
-                    labels_selector=labels_selector,
-                )
+            kind: Literal[\"{resource}\"] = \"{resource}\"
+            operation: Literal[{operation_type}] = {operation_type}
 
             @staticmethod
             def get_execution_event_type() -> type:
@@ -451,3 +439,6 @@ def main():
 
 if __name__ == "__main__":
     main()
+    print("Done. Before committing, it is recommended to run:")
+    print("poetry run isort src/robusta/integrations/kubernetes/autogenerated")
+    print("poetry run black src/robusta/integrations/kubernetes/autogenerated")

scripts/generate_playbook_descriptions.py

@@ -6,7 +6,9 @@
 from typing import Callable
 
 from pydantic import BaseModel
+from robusta.api import Action
 
+from robusta.core.playbooks.generation import ExamplesGenerator
 
 class PlaybookDescription(BaseModel):
     function_name: str
@@ -32,13 +34,12 @@ def get_params_schema(func):
     return action_params.schema()
 
 
-def load_scripts(scripts_root):
-    # install_requirements(os.path.join(scripts_root, 'requirements.txt'))
-
+def find_playbook_actions(scripts_root):
     python_files = glob.glob(f"{scripts_root}/*.py")
+    actions = []
 
     for script in python_files:
-        print(f"loading playbooks {script}")
+        print(f"found playbook file: {script}")
         filename = os.path.basename(script)
         (module_name, ext) = os.path.splitext(filename)
         spec = importlib.util.spec_from_file_location(module_name, script)
@@ -47,26 +48,35 @@ def load_scripts(scripts_root):
 
         playbooks = inspect.getmembers(
             module,
-            lambda f: inspect.isfunction(f) and getattr(f, "__playbook", None) is not None,
+            lambda f: Action.is_action(f),
         )
         for _, func in playbooks:
-            description = PlaybookDescription(
-                function_name=func.__name__,
-                builtin_trigger_params=func.__playbook["default_trigger_params"],
-                docs=inspect.getdoc(func),
-                src=inspect.getsource(func),
-                src_file=inspect.getsourcefile(func),
-                action_params=get_params_schema(func),
-            )
-            print(description.json(), "\n\n")
+            print("found playbook", func)
+            action = Action(func)
+            actions.append(action)
+
+            #description = PlaybookDescription(
+            #    function_name=func.__name__,
+            #    builtin_trigger_params=func.__playbook["default_trigger_params"],
+            #    docs=inspect.getdoc(func),
+            #    src=inspect.getsource(func),
+            #    src_file=inspect.getsourcefile(func),
+            #    action_params=get_params_schema(func),
+            #)
+            #print(description.json(), "\n\n")
+
+    return actions
 
 
 def main():
-    # TODO Arik - Need to be fixed in order to expose actions schema
     parser = argparse.ArgumentParser(description="Generate playbook descriptions")
-    parser.add_argument("directory", type=str, help="directory containing the playbooks")
+    parser.add_argument("--directory", type=str, help="directory containing the playbooks", default="./playbooks/robusta_playbooks")
     args = parser.parse_args()
-    load_scripts(args.directory)
+    actions = find_playbook_actions(args.directory)
+    generator = ExamplesGenerator()
+    triggers = generator.get_all_triggers()
+    trigger_to_actions = generator.get_triggers_to_actions(actions)
+    print(trigger_to_actions)
 
 
 if __name__ == "__main__":

scripts/main_app.py

@@ -0,0 +1,14 @@
+from pages import demo_playbooks, playbook_builder
+from streamlit import session_state as ss
+
+if "current_page" not in ss:
+    ss.current_page = "demo_playbooks"
+
+if "playbook_chosen" not in ss:
+    ss.playbook_chosen = False
+
+if ss.current_page == "demo_playbooks" and not ss.playbook_chosen:
+    demo_playbooks.display_demo_playbook()
+
+elif ss.current_page == "playbook_builder":
+    playbook_builder.display_playbook_builder()