diff --git a/packages/dsw-document-worker/Dockerfile b/packages/dsw-document-worker/Dockerfile index e681ab4..a62522c 100644 --- a/packages/dsw-document-worker/Dockerfile +++ b/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"] diff --git a/packages/dsw-document-worker/dsw/document_worker/cli.py b/packages/dsw-document-worker/dsw/document_worker/cli.py index 8a8a5c1..d07b0e4 100644 --- a/packages/dsw-document-worker/dsw/document_worker/cli.py +++ b/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__}') diff --git a/packages/dsw-document-worker/dsw/document_worker/consts.py b/packages/dsw-document-worker/dsw/document_worker/consts.py index 4b0024f..0ab6bf9 100644 --- a/packages/dsw-document-worker/dsw/document_worker/consts.py +++ b/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' diff --git a/packages/dsw-document-worker/dsw/document_worker/context.py b/packages/dsw-document-worker/dsw/document_worker/context.py index cc67d45..c3e4bff 100644 --- a/packages/dsw-document-worker/dsw/document_worker/context.py +++ b/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, diff --git a/packages/dsw-document-worker/dsw/document_worker/plugins/__init__.py b/packages/dsw-document-worker/dsw/document_worker/plugins/__init__.py new file mode 100644 index 0000000..93c1b95 --- /dev/null +++ b/packages/dsw-document-worker/dsw/document_worker/plugins/__init__.py @@ -0,0 +1,3 @@ +from .specs import hookspec, hookimpl + +__all__ = ['hookspec', 'hookimpl'] diff --git a/packages/dsw-document-worker/dsw/document_worker/plugins/manager.py b/packages/dsw-document-worker/dsw/document_worker/plugins/manager.py new file mode 100644 index 0000000..5d0f6d4 --- /dev/null +++ b/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 diff --git a/packages/dsw-document-worker/dsw/document_worker/plugins/specs.py b/packages/dsw-document-worker/dsw/document_worker/plugins/specs.py new file mode 100644 index 0000000..38af73d --- /dev/null +++ b/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 + """ diff --git a/packages/dsw-document-worker/dsw/document_worker/templates/steps/template.py b/packages/dsw-document-worker/dsw/document_worker/templates/steps/template.py index b012d54..5b50447 100644 --- a/packages/dsw-document-worker/dsw/document_worker/templates/steps/template.py +++ b/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: diff --git a/packages/dsw-document-worker/dsw/document_worker/templates/templates.py b/packages/dsw-document-worker/dsw/document_worker/templates/templates.py index 6e73a71..e6ab599 100644 --- a/packages/dsw-document-worker/dsw/document_worker/templates/templates.py +++ b/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 \ diff --git a/packages/dsw-document-worker/pyproject.toml b/packages/dsw-document-worker/pyproject.toml index 67dfdbf..cf01484 100644 --- a/packages/dsw-document-worker/pyproject.toml +++ b/packages/dsw-document-worker/pyproject.toml @@ -28,6 +28,7 @@ dependencies = [ 'Markdown', 'MarkupSafe', 'pathvalidate', + 'pluggy', 'python-dateutil', 'python-slugify', 'rdflib', diff --git a/packages/dsw-document-worker/requirements.txt b/packages/dsw-document-worker/requirements.txt index caeb60f..816902f 100644 --- a/packages/dsw-document-worker/requirements.txt +++ b/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 diff --git a/requirements.txt b/requirements.txt index f2c1527..015b4e9 100644 --- a/requirements.txt +++ b/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