global: support Jinja templating for job args

* Closes #36.

Author: slint

invenio_jobs/models.py

@@ -23,6 +23,8 @@
 from sqlalchemy_utils.types import ChoiceType, JSONType, UUIDType
 from werkzeug.utils import cached_property
 
+from .utils import eval_cfg_value, walk_dict_values
+
 JSON = (
     db.JSON()
     .with_variant(postgresql.JSONB(none_as_null=True), "postgresql")
@@ -44,12 +46,17 @@ class Job(db.Model, Timestamp):
     default_args = db.Column(JSON, default=lambda: dict(), nullable=True)
     schedule = db.Column(JSON, nullable=True)
 
-    # TODO: See if we move this to an API class
     @property
     def last_run(self):
         """Last run of the job."""
         return self.runs.order_by(Run.created.desc()).first()
 
+    def last_run_with_status(self, status):
+        """Last run of the job with a specific status."""
+        return (
+            self.runs.filter(Run.status == status).order_by(Run.created.desc()).first()
+        )
+
     @property
     def parsed_schedule(self):
         """Return schedule parsed as crontab or timedelta."""
@@ -63,6 +70,13 @@ def parsed_schedule(self):
         elif stype == "interval":
             return timedelta(**schedule)
 
+    def render_args(self, **ctx):
+        """Render job arguments."""
+        args = deepcopy(self.default_args)
+        ctx.setdefault("job", self)
+        walk_dict_values(args, lambda val: eval_cfg_value(val, ctx))
+        return args
+
 
 class RunStatusEnum(enum.Enum):
     """Enumeration of a run's possible states."""

invenio_jobs/services/scheduler.py

@@ -9,13 +9,11 @@
 
 import traceback
 import uuid
-from typing import Any
 
 from celery.beat import ScheduleEntry, Scheduler, logger
 from invenio_db import db
-from sqlalchemy import and_
 
-from invenio_jobs.models import Job, Run, Task
+from invenio_jobs.models import Job, Run
 from invenio_jobs.tasks import execute_run
 
 
@@ -49,27 +47,23 @@ class RunScheduler(Scheduler):
     Entry = JobEntry
     entries = {}
 
-    def __init__(self, *args: Any, **kwargs: Any) -> None:
-        """Initialize the database scheduler."""
-        super().__init__(*args, **kwargs)
-
     @property
     def schedule(self):
         """Get currently scheduled entries."""
         return self.entries
 
-    # Celery internal override
+    #
+    # Celery overrides
+    #
     def setup_schedule(self):
         """Setup schedule."""
         self.sync()
 
-    # Celery internal override
     def reserve(self, entry):
         """Update entry to next run execution time."""
         new_entry = self.schedule[entry.job.id] = next(entry)
         return new_entry
 
-    # Celery internal override
     def apply_entry(self, entry, producer=None):
         """Create and apply a JobEntry."""
         with self.app.flask_app.app_context():
@@ -93,24 +87,27 @@ def apply_entry(self, entry, producer=None):
                 else:
                     logger.debug("%s sent.", entry.task)
 
-    # Celery internal override
     def sync(self):
         """Sync Jobs from db to the scheduler."""
         # TODO Should we also have a cleaup task for runs? "stale" run (status running, starttime > hour, Run pending for > 1 hr)
         with self.app.flask_app.app_context():
             jobs = Job.query.filter(
-                and_(Job.active == True, Job.schedule != None)
-            ).all()
+                Job.active.is_(True),
+                Job.schedule.isnot(None),
+            )
             self.entries = {}  # because some jobs might be deactivated
             for job in jobs:
                 self.entries[job.id] = JobEntry.from_job(job)
 
+    #
+    # Helpers
+    #
     def create_run(self, entry):
         """Create run from a JobEntry."""
-        job = Job.query.filter_by(id=entry.job.id).one()
+        job = Job.query.get(entry.job.id)
         run = Run(
             job=job,
-            args=job.default_args,
+            args=job.render_args(),
             queue=job.default_queue,
             task_id=uuid.uuid4(),
         )

invenio_jobs/utils.py

@@ -0,0 +1,46 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright (C) 2024 CERN.
+#
+# Invenio-Jobs is free software; you can redistribute it and/or modify it
+# under the terms of the MIT License; see LICENSE file for more details.
+
+"""Utilities."""
+
+import ast
+
+from jinja2.sandbox import SandboxedEnvironment
+
+jinja_env = SandboxedEnvironment()
+
+
+def eval_cfg_value(val, ctx):
+    tpl = jinja_env.from_string(val)
+    res = tpl.render(**ctx)
+
+    try:
+        res = ast.literal_eval(res)
+    except Exception:
+        pass
+
+    return res
+
+
+def walk_dict_values(obj, transform_fn):
+    for key, val in obj.items():
+        if isinstance(val, str):
+            try:
+                obj[key] = transform_fn(val)
+            except Exception:
+                pass
+        elif isinstance(val, dict):
+            walk_dict_values(val, transform_fn)
+        elif isinstance(val, list):
+            for i, v in enumerate(val):
+                if isinstance(v, dict):
+                    walk_dict_values(v, transform_fn)
+                elif isinstance(v, str):
+                    try:
+                        val[i] = transform_fn(v)
+                    except Exception:
+                        pass