feat: Support document worker plugins

packages/dsw-document-worker/Dockerfile

@@ -45,4 +45,4 @@ RUN python -m pip install --break-system-packages --user --no-cache --no-index /
  && rm -rf /home/user/wheels
 
 # Run
-CMD ["dsw-document-worker"]
+CMD ["dsw-document-worker", "run"]

packages/dsw-document-worker/dsw/document_worker/cli.py

@@ -42,13 +42,18 @@ def validate_config(ctx, param, value: typing.IO | None):
     return load_config_str(content)
 
 
-@click.command(name='docworker')
+@click.group(name='dsw-document-worker')
 @click.version_option(version=VERSION)
+def main():
+    pass
+
+
+@main.command()
 @click.argument('config', envvar=VAR_APP_CONFIG_PATH,
                 required=False, callback=validate_config,
                 type=click.File('r', encoding=DEFAULT_ENCODING))
 @click.argument('workdir', envvar=VAR_WORKDIR_PATH,)
-def main(config: DocumentWorkerConfig, workdir: str):
+def run(config: DocumentWorkerConfig, workdir: str):
     config.log.apply()
     workdir_path = pathlib.Path(workdir)
     workdir_path.mkdir(parents=True, exist_ok=True)
@@ -62,3 +67,13 @@ def main(config: DocumentWorkerConfig, workdir: str):
         SentryReporter.capture_exception(e)
         click.echo(f'Error: {e}', err=True)
         sys.exit(2)
+
+
+@main.command()
+def list_plugins():
+    # pylint: disable-next=import-outside-toplevel
+    from .plugins.manager import create_manager
+
+    pm = create_manager()
+    for plugin in pm.list_name_plugin():
+        click.echo(f'{plugin[0]}: {plugin[1].__name__}')

packages/dsw-document-worker/dsw/document_worker/consts.py

@@ -6,6 +6,8 @@
 DEFAULT_ENCODING = 'utf-8'
 EXIT_SUCCESS = 0
 NULL_UUID = '00000000-0000-0000-0000-000000000000'
+PACKAGE_NAME = 'dsw-document-worker'
+PLUGINS_ENTRYPOINT = 'dsw_document_worker_plugins'
 PROG_NAME = 'docworker'
 VERSION = '4.14.0'
 

packages/dsw-document-worker/dsw/document_worker/context.py

@@ -1,9 +1,10 @@
 import dataclasses
 import pathlib
 
+import pluggy
+
 from dsw.database import Database
 from dsw.storage import S3Storage
-
 from .config import DocumentWorkerConfig
 
 
@@ -15,6 +16,7 @@ def __init__(self):
 
 @dataclasses.dataclass
 class AppContext:
+    pm: pluggy.PluginManager
     db: Database
     s3: S3Storage
     cfg: DocumentWorkerConfig
@@ -56,8 +58,11 @@ def get(cls) -> _Context:
 
     @classmethod
     def initialize(cls, db, s3, config, workdir):
+        # pylint: disable-next=import-outside-toplevel
+        from .plugins.manager import create_manager
         cls._instance = _Context(
             app=AppContext(
+                pm=create_manager(),
                 db=db,
                 s3=s3,
                 cfg=config,

packages/dsw-document-worker/dsw/document_worker/plugins/__init__.py

@@ -0,0 +1,3 @@
+from .specs import hookspec, hookimpl
+
+__all__ = ['hookspec', 'hookimpl']

packages/dsw-document-worker/dsw/document_worker/plugins/manager.py

@@ -0,0 +1,13 @@
+import pluggy
+
+from ..consts import PACKAGE_NAME, PLUGINS_ENTRYPOINT
+
+
+def create_manager():
+    # pylint: disable-next=import-outside-toplevel
+    import dsw.document_worker.plugins.specs as hookspecs
+
+    pm = pluggy.PluginManager(PACKAGE_NAME)
+    pm.load_setuptools_entrypoints(PLUGINS_ENTRYPOINT)
+    pm.add_hookspecs(hookspecs)
+    return pm

packages/dsw-document-worker/dsw/document_worker/plugins/specs.py

@@ -0,0 +1,56 @@
+# pylint: disable=unused-argument
+import pluggy
+import jinja2
+
+from ..consts import PACKAGE_NAME
+from ..templates.steps import Step
+
+hookspec = pluggy.HookspecMarker(PACKAGE_NAME)
+hookimpl = pluggy.HookimplMarker(PACKAGE_NAME)
+
+
+@hookspec
+def provide_steps() -> dict[str, type[Step]]:
+    """
+    Provide a dictionary of steps that the plugin can execute.
+
+    Steps are used to process the document generation in a specific way. Each step
+    must comply with the interface of the provided `Step` class. The plugin must make
+    sure to return a list of steps that it can execute and in compliance with the
+    `Step` class in the current implementation (use correct dsw-document-worker as
+    a dependency).
+
+    :return: a dictionary of steps that the plugin can execute
+    """
+    return {}
+
+
+@hookspec
+def enrich_document_context(context: dict) -> None:
+    """
+    Enrich the document context with custom data.
+
+    The plugin can enrich the provided document context with custom data that can be
+    used in the document generation process. The context is a dictionary that is passed
+    to each step during the document generation process. The plugin can manipulate the
+    context in any way it sees fit.
+
+    :param context: the document context to enrich
+    """
+
+
+@hookspec
+def enrich_jinja_env(jinja_env: jinja2.Environment, options: dict[str, str]) -> None:
+    """
+    Enrich the Jinja environment with custom filters, tests, policies, etc.
+
+    The plugin can enrich the provided Jinja environment with custom features like
+    filters, tests, policies, global variables and other. It can manipulate the Jinja
+    environment in any way it sees fit but should ensure that the environment is safe
+    to use in the document generation process. This Jinja environment is used within
+    all steps that are subclass of the `JinjaPoweredStep` (mainly the `jinja` step
+    implemented in class `Jinja2Step`).
+
+    :param jinja_env: the Jinja environment to enrich
+    :param options: the options provided to the step
+    """

packages/dsw-document-worker/dsw/document_worker/templates/steps/template.py

@@ -70,6 +70,11 @@ def __init__(self, template, options):
                 self.j2_env.add_extension('jinja2.ext.debug')
             self._apply_policies(options)
             self._add_j2_enhancements()
+
+            Context.get().app.pm.hook.enrich_jinja_env(
+                jinja_env=self.j2_env,
+                options=options,
+            )
         except jinja2.exceptions.TemplateSyntaxError as e:
             self.raise_exc(self._jinja_exception_msg(e))
         except Exception as e:

packages/dsw-document-worker/dsw/document_worker/templates/templates.py

@@ -11,7 +11,8 @@
 from ..consts import FormatField
 from ..context import Context
 from ..documents import DocumentFile
-from ..templates.formats import Format
+from .formats import Format
+from .steps.base import register_step, Step
 
 
 LOG = logging.getLogger(__name__)
@@ -217,6 +218,8 @@ def __getitem__(self, format_uuid: str) -> Format:
         return self.formats[format_uuid]
 
     def render(self, format_uuid: str, context: dict) -> DocumentFile:
+        Context.get().app.pm.hook.enrich_document_context(context=context)
+
         self.last_used = datetime.datetime.now(tz=datetime.UTC)
         return self[format_uuid].execute(context)
 
@@ -233,6 +236,14 @@ def get(cls) -> 'TemplateRegistry':
 
     def __init__(self):
         self._templates: dict[str, dict[str, Template]] = {}
+        self._load_plugin_steps()
+
+    def _load_plugin_steps(self):
+        for steps_dict in Context.get().app.pm.hook.provide_steps():
+            for name, step_class in steps_dict.items():
+                if not issubclass(step_class, Step):
+                    raise RuntimeError(f'Provided class "{step_class}" is not a subclass of Step')
+                register_step(name, step_class)
 
     def has_template(self, tenant_uuid: str, template_id: str) -> bool:
         return tenant_uuid in self._templates and \

packages/dsw-document-worker/pyproject.toml

@@ -28,6 +28,7 @@ dependencies = [
     'Markdown',
     'MarkupSafe',
     'pathvalidate',
+    'pluggy',
     'python-dateutil',
     'python-slugify',
     'rdflib',

packages/dsw-document-worker/requirements.txt

@@ -17,6 +17,7 @@ minio==7.2.13
 panflute==2.3.1
 pathvalidate==3.2.1
 pillow==11.0.0
+pluggy==1.5.0
 psycopg==3.2.3
 psycopg-binary==3.2.3
 pycparser==2.22

requirements.txt

@@ -32,6 +32,7 @@ panflute==2.3.1
 pathspec==0.12.1
 pathvalidate==3.2.1
 pillow==11.0.0
+pluggy==1.5.0
 propcache==0.2.1
 psycopg==3.2.3
 psycopg-binary==3.2.3