Merge pull request #181 from tbotnz/per-job-ttl

bugfix on user script models and TTL support per task

netpalm/backend/core/manager/netpalm_manager.py

@@ -31,6 +31,7 @@
 from netpalm.backend.core.models.task import Response
 
 from netpalm.backend.plugins.utilities.webhook.webhook import exec_webhook_func
+from netpalm.backend.plugins.calls.scriptrunner.script import script_model_finder
 
 log = logging.getLogger(__name__)
 
@@ -124,13 +125,10 @@ def set_config_restconf(self, setcfg: Restconf):
         """ executes the restconf setconfig method async and returns the response obj """
         return self._set_config(setcfg, library="restconf")
 
-    def execute_script(self, script: Script):
+    def execute_script(self, **kwargs):
         """ executes the netpalm script method async and returns the response obj """
-        if isinstance(script, dict):
-            req_data = script
-        else:
-            req_data = script.dict(exclude_none=True)
-
+        log.debug(f"execute_script: called with {kwargs}")
+        req_data = kwargs
         # check if pinned required
         if req_data.get("queue_strategy") == "pinned":
             if isinstance(req_data.get("connection_args"), dict):

netpalm/backend/core/models/models.py

@@ -64,6 +64,7 @@ class SetConfig(BaseModel):
     pre_checks: Optional[List[GenericPrePostCheck]] = None
     post_checks: Optional[List[GenericPrePostCheck]] = None
     enable_mode: bool = False
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {
@@ -122,6 +123,7 @@ class Script(BaseModel):
     webhook: Optional[Webhook] = None
     queue_strategy: Optional[QueueStrategy] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {
@@ -134,6 +136,20 @@ class Config:
             }
         }
 
+class ScriptCustom(BaseModel):
+    script: str
+    webhook: Optional[Webhook] = None
+    queue_strategy: Optional[QueueStrategy] = None
+    cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "script": "hello_world",
+                "queue_strategy": "fifo"
+            }
+        }
 
 class GetConfig(BaseModel):
     library: LibraryName
@@ -144,6 +160,7 @@ class GetConfig(BaseModel):
     queue_strategy: Optional[QueueStrategy] = None
     post_checks: Optional[List[GenericPrePostCheck]] = []
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/ncclient.py

@@ -71,6 +71,7 @@ class NcclientSetConfig(BaseModel):
     j2config: Optional[J2Config] = None
     webhook: Optional[Webhook] = None
     queue_strategy: Optional[QueueStrategy] = None
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {
@@ -100,6 +101,7 @@ class NcclientGetConfig(BaseModel):
     webhook: Optional[Webhook] = None
     queue_strategy: Optional[QueueStrategy] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {
@@ -134,6 +136,7 @@ class NcclientGet(BaseModel):
     args: NcclientGetArgs
     queue_strategy: Optional[QueueStrategy] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/netmiko.py

@@ -74,6 +74,7 @@ class NetmikoGetConfig(BaseModel):
     queue_strategy: Optional[QueueStrategy] = None
     post_checks: Optional[List[GenericPrePostCheck]] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {
@@ -109,6 +110,7 @@ class NetmikoSetConfig(BaseModel):
     pre_checks: Optional[List[GenericPrePostCheck]] = None
     post_checks: Optional[List[GenericPrePostCheck]] = None
     enable_mode: bool = False
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/puresnmp.py

@@ -34,6 +34,7 @@ class PureSNMPGetConfig(BaseModel):
     webhook: Optional[Webhook] = None
     queue_strategy: Optional[QueueStrategy] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/restconf.py

@@ -38,6 +38,7 @@ class Restconf(BaseModel):
     webhook: Optional[Webhook] = None
     queue_strategy: Optional[QueueStrategy] = None
     cache: Optional[CacheConfig] = {}
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/service.py

@@ -39,6 +39,7 @@ class ServiceModel(BaseModel):
     operation: ServiceLifecycle
     args: dict
     queue_strategy: Optional[QueueStrategy] = None
+    ttl: Optional[int] = None
 
     class Config:
         schema_extra = {

netpalm/backend/core/models/task.py

@@ -26,7 +26,6 @@ class TaskMetaData(BaseModel):
     total_elapsed_seconds: Optional[str]
     assigned_worker: Optional[str]
 
-
 class TaskError(BaseModel):
     exception_class: str
     exception_args: List[str]

netpalm/backend/core/redis/rediz.py

@@ -360,10 +360,18 @@ def render_task_response(self, task_job):
         return resultdata
 
     def sendtask(self, q, exe, **kwargs):
+
+        log.debug(f'sendtask: {kwargs["kwargs"]}')
+        ttl = kwargs["kwargs"].get("ttl")
         meta_template = self.get_redis_meta_template()
-        task = self.local_queuedb[q]["queue"].enqueue_call(func=self.routes[exe], description=q, ttl=self.ttl,
-                                                           result_ttl=self.task_result_ttl, kwargs=kwargs["kwargs"],
-                                                           meta=meta_template, timeout=self.timeout)
+        if not ttl:
+            task = self.local_queuedb[q]["queue"].enqueue_call(func=self.routes[exe], description=q, ttl=self.ttl,
+                                                            result_ttl=self.task_result_ttl, kwargs=kwargs["kwargs"],
+                                                            meta=meta_template, timeout=self.timeout)
+        else:
+            task = self.local_queuedb[q]["queue"].enqueue_call(func=self.routes[exe], description=q, ttl=ttl,
+                                                            result_ttl=ttl, kwargs=kwargs["kwargs"],
+                                                            meta=meta_template, timeout=ttl)
         resultdata = self.render_task_response(task)
         return resultdata
 

netpalm/backend/core/routes/routes.py

@@ -2,7 +2,7 @@
 from netpalm.backend.plugins.calls.dryrun.dryrun import dryrun
 from netpalm.backend.plugins.calls.getconfig.exec_command import exec_command
 from netpalm.backend.plugins.calls.getconfig.ncclient_get import ncclient_get
-from netpalm.backend.plugins.calls.scriptrunner.script import script_exec
+from netpalm.backend.plugins.calls.scriptrunner.script import script_kiddy
 from netpalm.backend.plugins.calls.service.procedures import (
     create,
     update,
@@ -27,7 +27,7 @@
     "pushtemplate": pushtemplate,
     "removetemplate": removetemplate,
     "ls": list_files,
-    "script": script_exec,
+    "script": script_kiddy,
     "j2gettemplate": j2gettemplate,
     "render_j2template": render_j2template,
     "dryrun": dryrun,

netpalm/backend/core/utilities/rediz_meta.py

@@ -62,7 +62,6 @@ def write_mandatory_meta():
     job = get_current_job()
     if job is None:  # it will be None in many/all unit tests
         return
-
     job.meta["assigned_worker"] = config.worker_name
     job.save_meta()
 

netpalm/backend/plugins/calls/scriptrunner/script.py

@@ -1,39 +1,102 @@
 import importlib
+import inspect
+
+import logging
 
 from netpalm.backend.core.confload.confload import config
 from netpalm.backend.core.utilities.rediz_meta import render_netpalm_payload, write_mandatory_meta
 from netpalm.backend.core.utilities.rediz_meta import write_meta_error
 from netpalm.backend.plugins.utilities.webhook.webhook import exec_webhook_func
 
+from netpalm.backend.core.models.models import Script, ScriptCustom
+
+log = logging.getLogger(__name__)
+
 
-class script_kiddy:
-    def __init__(self, **kwargs):
-        self.scrp_path = config.custom_scripts
-        self.kwarg = kwargs.get('kwargs', False)
-        self.arg = self.kwarg.get('args', False)
-        self.script = self.kwarg.get('script', False)
-        self.script_name = self.scrp_path.replace('/', '.') + self.script
+def script_model_finder(script_name: str):
+    log.debug(f"script_model_finder: locating model for  {script_name}")
+    model = Script
+    model_defined = False
+    model_mode = None
+    # first check whether there is the legacy _model.py file against the script name
+    try:
+        model_name = f"{script_name}_model"
+        template_model_path_raw = config.custom_scripts
+        template_model_path = template_model_path_raw.replace('/', '.') + model_name
+        module = importlib.import_module(template_model_path)
+        model = getattr(module, model_name)
+        model_defined = True
+        model_mode = "legacy"
+    except Exception as e:
+        log.debug(f"script_model_finder: no legacy model found for {script_name} import error {e} attempting with newer model in file")
+        model = Script
+        pass
 
-    def s_exec(self):
-        module = importlib.import_module(self.script_name)
+    # if this does not exist, check within the file to see if a model exists 
+    # clean this up at some point -_-'
+
+    try:
+        model_name = f"{script_name}"
+        template_model_path_raw = config.custom_scripts
+        template_model_path = template_model_path_raw.replace('/', '.') + model_name
+        module = importlib.import_module(template_model_path)
         runscrp = getattr(module, "run")
-        res = runscrp(kwargs=self.arg)
-        return res
+        for item in inspect.getfullargspec(runscrp):
+            if type(item) is dict:
+                for key, value in item.items():
+                    if issubclass(value, ScriptCustom):
+                        model = value
+                        model_defined = True
+                        model_mode = "new"
+    except Exception as e:
+        pass
+    log.debug(f"script_model_finder: returning  {model}")
+    return model, model_defined, model_mode
 
 
-def script_exec(**kwargs):
+def script_kiddy(**kwargs):
     webhook = kwargs.get("webhook", False)
     result = False
 
+    log.debug(f'script_kiddy: locating model for script {kwargs["script"]}')
+    model = script_model_finder(kwargs["script"])
+    model_to_validate = model[0]
+    model_is_defined = model[1]
+    model_mode = model[2]
+    log.debug(f"script_kiddy: model located is {model} and a user model was found is {model_is_defined}")
+
     try:
         write_mandatory_meta()
-        scrip = script_kiddy(kwargs=kwargs)
-        result = scrip.s_exec()
 
+        # execute the script
+        scrp_path = config.custom_scripts
+        kwarg = kwargs
+        arg = kwarg.get('args', False)
+        script_name = kwarg.get('script', False)
+        script_path_full_name = scrp_path.replace('/', '.') + script_name
+        log.debug(f'script_kiddy: attempting to import script {script_path_full_name} for run')
+
+        module = importlib.import_module(script_path_full_name)
+        runscrp = getattr(module, "run")
+    except Exception as e:
+        log.error(f"script_kiddy: could not import {script_path_full_name} with {e}")
+        write_meta_error(e)
+
+    try:
+        log.debug(f'script_kiddy: attempting to run script {script_path_full_name}')
+        if not model_is_defined or model_mode is "legacy":
+            result = runscrp(kwargs=arg)
+        else:
+            data_to_send = model_to_validate(**kwarg)
+            result = runscrp(data_to_send)
+        # if webhook used do that too
         if webhook:
             current_jobdata = render_netpalm_payload(job_result=result)
             exec_webhook_func(jobdata=current_jobdata, webhook_payload=webhook)
     except Exception as e:
+        log.error(f"script_kiddy: could not run script {script_path_full_name} with {e}")
         write_meta_error(e)
 
     return result
+
+

netpalm/backend/plugins/extensibles/custom_scripts/hello_world_embedded_pydanticmodel.py

@@ -1,13 +1,14 @@
 from typing import Optional, Any, List
-from pydantic import BaseModel
+from netpalm.backend.core.models.models import ScriptCustom
 
-class CustomScriptModel(BaseModel):
-    args: Optional[str] = None
+class MyCustomScriptModel(ScriptCustom):
+    script: str
+    test: Optional[str] = None
 
-def run(payload: CustomScriptModel):
+def run(payload: MyCustomScriptModel):
     try:
         # mandatory get of kwargs - payload comes through as {"kwargs": {"hello": "world"}}
-        args = payload.args
+        args = payload.test
         # reutn "world"
         return args
     except Exception as e: