Added provenance_manager extension point

This commit adds the possibility to upload custom `ProvenanceManager`
classes through StreamFlow plugins and to visualize them through the
`streamflow ext list` and `streamflow ext show` commands

Author: GlassOfWhiskey

pyproject.toml

@@ -72,6 +72,7 @@ zip-safe = true
 "streamflow.deployment.connector" = ["schemas/base/*.json", "schemas/*.json"]
 "streamflow.deployment.filter" = ["schemas/*.json"]
 "streamflow.persistence" = ["schemas/*.sql", "schemas/*.json"]
+"streamflow.provenance" = ["schemas/**/*.json"]
 "streamflow.recovery" = ["schemas/*.json"]
 "streamflow.scheduling" = ["schemas/*.json"]
 "streamflow.scheduling.policy" = ["schemas/*.json"]

streamflow/ext/plugin.py

@@ -7,6 +7,7 @@
 from streamflow.core.data import DataManager
 from streamflow.core.deployment import BindingFilter, Connector, DeploymentManager
 from streamflow.core.persistence import Database
+from streamflow.core.provenance import ProvenanceManager
 from streamflow.core.recovery import CheckpointManager, FailureManager
 from streamflow.core.scheduling import Policy, Scheduler
 from streamflow.cwl.requirement.docker import cwl_docker_translator_classes
@@ -17,6 +18,7 @@
 from streamflow.deployment.filter import binding_filter_classes
 from streamflow.log_handler import logger
 from streamflow.persistence import database_classes
+from streamflow.provenance import provenance_manager_classes
 from streamflow.recovery import checkpoint_manager_classes, failure_manager_classes
 from streamflow.scheduling import scheduler_classes
 from streamflow.scheduling.policy import policy_classes
@@ -31,6 +33,7 @@
     "deployment_manager": deployment_manager_classes,
     "failure_manager": failure_manager_classes,
     "policy": policy_classes,
+    "provenance_manager": provenance_manager_classes,
     "scheduler": scheduler_classes,
 }
 
@@ -85,6 +88,28 @@ def register_failure_manager(self, name: str, cls: type[FailureManager]):
     def register_policy(self, name: str, cls: type[Policy]):
         self._register(name, cls, "policy")
 
+    def register_provenance_manager(
+        self, name: str, wf_type: str, cls: type[ProvenanceManager]
+    ):
+        self.classes_.setdefault("provenance_manager", []).append(
+            {
+                "name": f"{name}/{wf_type}",
+                "class": cls,
+            }
+        )
+        if (
+            name in extension_points["provenance_manager"]
+            and wf_type in extension_points["provenance_manager"][name]
+        ):
+            if logger.isEnabledFor(logging.WARNING):
+                logger.warning(
+                    "{} is already installed and will be overridden by {}".format(
+                        f"{name}/{wf_type}",
+                        self.__class__.__module__ + "." + self.__class__.__name__,
+                    )
+                )
+        extension_points["provenance_manager"][name] = {wf_type: cls}
+
     def register_scheduler(self, name: str, cls: type[Scheduler]):
         self._register(name, cls, "scheduler")
 

streamflow/ext/utils.py

@@ -19,6 +19,28 @@
 PLUGIN_ENTRY_POINT = "unito.streamflow.plugin"
 
 
+def _add_extension(
+    extensions: MutableMapping[str, Any],
+    ext_name: str,
+    name: str,
+    plugin_name: str,
+    value: type | MutableMapping[str, type],
+) -> None:
+    if ext_name == "provenance_manager":
+        for wf_type in value:
+            extensions.setdefault(ext_name, {})[name] = {
+                "name": f"{name}/{wf_type}",
+                "class": get_class_fullname(value[wf_type]),
+                "plugin": plugin_name,
+            }
+    else:
+        extensions.setdefault(ext_name, {})[name] = {
+            "name": name,
+            "class": get_class_fullname(value),
+            "plugin": plugin_name,
+        }
+
+
 def _filter_by_name(classes: MutableMapping[str, Any], name: str):
     filtered_classes = {}
     for class_ in classes:
@@ -39,6 +61,31 @@ def _flatten_all_of(entity_schema):
     return dict(sorted(entity_schema["properties"].items()))
 
 
+def _get_plugin_and_class(name: str, type_: str) -> tuple[str, type | None]:
+    if type_ == "provenance_manager":
+        if "/" not in name:
+            return "-", None
+        prov_name, wf_type = name.split("/", 1)
+        if (
+            prov_name in extension_points[type_]
+            and wf_type in extension_points[type_][prov_name]
+        ):
+            return "-", extension_points[type_][prov_name][wf_type]
+    elif name in extension_points[type_]:
+        return "-", extension_points[type_][name]
+    for plugin_obj in entry_points(group=PLUGIN_ENTRY_POINT):
+        plugin_class = (plugin_obj.load())()
+        if isinstance(plugin_class, StreamFlowPlugin):
+            plugin_class.register()
+            plugin_classes = _filter_by_name(
+                {k: v for k, v in plugin_class.classes_.items() if k == type_}, name
+            )
+            for item in plugin_classes.get(type_, []):
+                if item["name"] == name:
+                    return plugin_obj.name, item["class"]
+    return "-", None
+
+
 def _get_property_desc(
     k: str, obj: MutableMapping[str, Any], refs: MutableMapping[str, Any]
 ) -> str:
@@ -151,11 +198,13 @@ def list_extensions(name: str | None, type_: str | None):
         if type_ is None or ext == type_:
             for n, v in classes.items():
                 if name is None or name == n:
-                    extensions.setdefault(ext, {})[n] = {
-                        "name": n,
-                        "class": get_class_fullname(v),
-                        "plugin": "-",
-                    }
+                    _add_extension(
+                        extensions=extensions,
+                        ext_name=ext,
+                        name=n,
+                        plugin_name="-",
+                        value=v,
+                    )
                     for k in max_sizes:
                         max_sizes[k] = max(max_sizes[k], len(extensions[ext][n][k]))
     if plugins := entry_points(group=PLUGIN_ENTRY_POINT):
@@ -172,11 +221,13 @@ def list_extensions(name: str | None, type_: str | None):
                     plugin_classes = _filter_by_name(plugin_classes, name)
                 for ext, items in plugin_classes.items():
                     for item in items:
-                        extensions.setdefault(ext, {})[item["name"]] = {
-                            "name": item["name"],
-                            "class": get_class_fullname(item["class"]),
-                            "plugin": plugin.name,
-                        }
+                        _add_extension(
+                            extensions=extensions,
+                            ext_name=ext,
+                            name=item["name"],
+                            plugin_name=plugin.name,
+                            value=item["class"],
+                        )
                         for k in max_sizes:
                             max_sizes[k] = max(
                                 max_sizes[k], len(extensions[ext][item["name"]][k])
@@ -288,22 +339,7 @@ def load_extensions():
 
 
 def show_extension(name: str, type_: str):
-    plugin = "-"
-    if name in extension_points[type_]:
-        class_ = extension_points[type_][name]
-    else:
-        class_ = None
-        for plugin_obj in entry_points(group=PLUGIN_ENTRY_POINT):
-            plugin_class = (plugin_obj.load())()
-            if isinstance(plugin_class, StreamFlowPlugin):
-                plugin_class.register()
-                plugin_classes = _filter_by_name(
-                    {k: v for k, v in plugin_class.classes_.items() if k == type_}, name
-                )
-                for item in plugin_classes.get(type_, []):
-                    if item["name"] == name:
-                        class_ = item["class"]
-                        plugin = plugin_obj.name
+    plugin, class_ = _get_plugin_and_class(name, type_)
     if class_ is not None:
         class_name = get_class_fullname(class_)
         entity_schema = class_.get_schema()

streamflow/main.py

@@ -31,7 +31,7 @@
 from streamflow.parser import parser
 from streamflow.persistence import database_classes
 from streamflow.persistence.loading_context import DefaultDatabaseLoadingContext
-from streamflow.provenance import prov_classes
+from streamflow.provenance import provenance_manager_classes
 from streamflow.recovery import checkpoint_manager_classes, failure_manager_classes
 from streamflow.scheduling import scheduler_classes
 
@@ -126,20 +126,20 @@ async def _async_prov(args: argparse.Namespace):
                 f"Workflow {args.workflow} is associated to the following types: {','.join(wf_type)}"
             )
         wf_type = list(wf_type)[0]
-        if args.type not in prov_classes:
+        if args.type not in provenance_manager_classes:
             raise WorkflowProvenanceException(
                 f"{args.type} provenance format is not supported."
             )
-        elif wf_type not in prov_classes[args.type]:
+        elif wf_type not in provenance_manager_classes[args.type]:
             raise WorkflowProvenanceException(
                 "{} provenance format is not supported for workflows of type {}.".format(
                     args.type, wf_type
                 )
             )
         else:
-            provenance_manager: ProvenanceManager = prov_classes[args.type][wf_type](
-                context, db_context, workflows
-            )
+            provenance_manager: ProvenanceManager = provenance_manager_classes[
+                args.type
+            ][wf_type](context, db_context, workflows)
             await provenance_manager.create_archive(
                 outdir=args.outdir,
                 filename=args.name,

streamflow/provenance/__init__.py

@@ -1,3 +1,3 @@
 from streamflow.provenance.run_crate import CWLRunCrateProvenanceManager
 
-prov_classes = {"run_crate": {"cwl": CWLRunCrateProvenanceManager}}
+provenance_manager_classes = {"run_crate": {"cwl": CWLRunCrateProvenanceManager}}

streamflow/provenance/run_crate.py

@@ -12,6 +12,7 @@
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import MutableMapping, MutableSequence
+from importlib.resources import files
 from typing import Any, cast, get_args
 from zipfile import ZipFile
 
@@ -1424,6 +1425,16 @@ async def get_main_entity(self) -> MutableMapping[str, Any]:
                     )
         return main_entity
 
+    @classmethod
+    def get_schema(cls) -> str:
+        return (
+            files(__package__)
+            .joinpath("schemas")
+            .joinpath("run_crate")
+            .joinpath("cwl.json")
+            .read_text("utf-8")
+        )
+
     async def get_property_value(
         self, name: str, token: Token
     ) -> MutableMapping[str, Any] | None:

streamflow/provenance/schemas/run_crate/cwl.json

@@ -0,0 +1,7 @@
+{
+  "$schema": "https://json-schema.org/draft/2019-09/schema",
+  "$id": "https://streamflow.di.unito.it/schemas/provenance/run_crate/cwl.json",
+  "type": "object",
+  "properties": {},
+  "additionalProperties": false
+}