add try except finally around run method within the decorator function

nautobot · Mar 13, 2024 · aae5fa4 · aae5fa4
1 parent 3c12378
commit aae5fa4
Show file tree

Hide file tree

Showing 4 changed files with 96 additions and 114 deletions.
diff --git a/nautobot_golden_config/jobs.py b/nautobot_golden_config/jobs.py
@@ -118,21 +118,28 @@ def gc_repo_wrapper(self, *args, **kwargs):
         )
         current_repos = get_refreshed_repos(job_obj=self, repo_types=gitrepo_types, data=self.qs)
         # This is where the specific jobs run method runs via this decorator.
-        func(self, *args, **kwargs)
-        now = make_aware(datetime.now())
-        self.logger.debug(
-            f"Finished the {self.Meta.name} job execution.",
-            extra={"grouping": "GC After Run"},
-        )
-        if current_repos:
-            for _, repo in current_repos.items():
-                if repo["to_commit"]:
-                    self.logger.debug(
-                        f"Pushing {self.Meta.name} results to repo {repo['repo_obj'].base_url}.",
-                        extra={"grouping": "GC Repo Commit and Push"},
-                    )
-                    repo["repo_obj"].commit_with_added(f"{self.Meta.name.upper()} JOB {now}")
-                    repo["repo_obj"].push()
+        try:
+            func(self, *args, **kwargs)
+        except Exception as error:  # pylint disable=broad-exception-caught
+            error_msg = f"`E3001:` General Exception handler, original error message ```{error}```"
+            # Raise error only if the job kwarg (checkbox) is selected to do so on the job execution form.
+            if kwargs["fail_job_on_task_failure"]:
+                raise NornirNautobotException(error_msg) from error
+        finally:
+            now = make_aware(datetime.now())
+            self.logger.debug(
+                f"Finished the {self.Meta.name} job execution.",
+                extra={"grouping": "GC After Run"},
+            )
+            if current_repos:
+                for _, repo in current_repos.items():
+                    if repo["to_commit"]:
+                        self.logger.debug(
+                            f"Pushing {self.Meta.name} results to repo {repo['repo_obj'].base_url}.",
+                            extra={"grouping": "GC Repo Commit and Push"},
+                        )
+                        repo["repo_obj"].commit_with_added(f"{self.Meta.name.upper()} JOB {now}")
+                        repo["repo_obj"].push()
 
     return gc_repo_wrapper
 
@@ -161,6 +168,9 @@ class FormEntry:  # pylint disable=too-few-public-method
         label="Device Status",
     )
     debug = BooleanVar(description="Enable for more verbose debug logging")
+    fail_job_on_task_failure = BooleanVar(
+        description="If any device in the tasks list fails, fail the entire job result."
+    )
 
 
 class GoldenConfigJobMixin(Job):  # pylint: disable=abstract-method

diff --git a/nautobot_golden_config/nornir_plays/config_backup.py b/nautobot_golden_config/nornir_plays/config_backup.py
@@ -10,12 +10,11 @@
 from nornir import InitNornir
 from nornir.core.plugins.inventory import InventoryPluginRegister
 from nornir.core.task import Result, Task
-from nornir_nautobot.exceptions import NornirNautobotException
 from nornir_nautobot.plugins.tasks.dispatcher import dispatcher
 from nautobot_golden_config.models import ConfigRemove, ConfigReplace, GoldenConfig
 from nautobot_golden_config.nornir_plays.processor import ProcessGoldenConfig
 from nautobot_golden_config.utilities.db_management import close_threaded_db_connections
-from nautobot_golden_config.utilities.helper import (  # get_device_to_settings_map,; get_job_filter,
+from nautobot_golden_config.utilities.helper import (
     dispatch_params,
     render_jinja_template,
     verify_settings,
@@ -103,41 +102,29 @@ def config_backup(job_result, log_level, qs, device_to_settings_map):
         if not replace_regex_dict.get(regex.platform.network_driver):
             replace_regex_dict[regex.platform.network_driver] = []
         replace_regex_dict[regex.platform.network_driver].append({"replace": regex.replace, "regex": regex.regex})
-    try:
-        with InitNornir(
-            runner=NORNIR_SETTINGS.get("runner"),
-            logging={"enabled": False},
-            inventory={
-                "plugin": "nautobot-inventory",
-                "options": {
-                    "credentials_class": NORNIR_SETTINGS.get("credentials"),
-                    "params": NORNIR_SETTINGS.get("inventory_params"),
-                    "queryset": qs,
-                    "defaults": {"now": now},
-                },
+    with InitNornir(
+        runner=NORNIR_SETTINGS.get("runner"),
+        logging={"enabled": False},
+        inventory={
+            "plugin": "nautobot-inventory",
+            "options": {
+                "credentials_class": NORNIR_SETTINGS.get("credentials"),
+                "params": NORNIR_SETTINGS.get("inventory_params"),
+                "queryset": qs,
+                "defaults": {"now": now},
             },
-        ) as nornir_obj:
-            nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
-
-            logger.debug("Run nornir backup tasks.")
-            nr_with_processors.run(
-                task=run_backup,
-                name="BACKUP CONFIG",
-                logger=logger,
-                device_to_settings_map=device_to_settings_map,
-                remove_regex_dict=remove_regex_dict,
-                replace_regex_dict=replace_regex_dict,
-            )
-            logger.debug("Completed configuration from devices.")
-
-    except ValueError as error:
-        # This block means the nornir processor raised ValueError which is "somewhat" expected.
-        # It should not reraise and fail the entire job.
-        error_msg = f"`E3001:` General Exception handler, original error message ```{error}```"
-        logger.error(error_msg)
-    except Exception as error:
-        error_msg = f"`E3001:` Unknown Exception handler, original error message ```{error}```"
-        logger.error(error_msg)
-        raise NornirNautobotException(error_msg) from err
-
+        },
+    ) as nornir_obj:
+        nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
+
+        logger.debug("Run nornir backup tasks.")
+        nr_with_processors.run(
+            task=run_backup,
+            name="BACKUP CONFIG",
+            logger=logger,
+            device_to_settings_map=device_to_settings_map,
+            remove_regex_dict=remove_regex_dict,
+            replace_regex_dict=replace_regex_dict,
+        )
+        logger.debug("Completed configuration from devices.")
     logger.debug("Completed configuration backup job for devices.")
diff --git a/nautobot_golden_config/nornir_plays/config_compliance.py b/nautobot_golden_config/nornir_plays/config_compliance.py
@@ -177,35 +177,27 @@ def config_compliance(job_result, log_level, qs, device_to_settings_map):
 
     for settings in set(device_to_settings_map.values()):
         verify_settings(logger, settings, ["backup_path_template", "intended_path_template"])
-
-    try:
-        with InitNornir(
-            runner=NORNIR_SETTINGS.get("runner"),
-            logging={"enabled": False},
-            inventory={
-                "plugin": "nautobot-inventory",
-                "options": {
-                    "credentials_class": NORNIR_SETTINGS.get("credentials"),
-                    "params": NORNIR_SETTINGS.get("inventory_params"),
-                    "queryset": qs,
-                    "defaults": {"now": now},
-                },
+    with InitNornir(
+        runner=NORNIR_SETTINGS.get("runner"),
+        logging={"enabled": False},
+        inventory={
+            "plugin": "nautobot-inventory",
+            "options": {
+                "credentials_class": NORNIR_SETTINGS.get("credentials"),
+                "params": NORNIR_SETTINGS.get("inventory_params"),
+                "queryset": qs,
+                "defaults": {"now": now},
             },
-        ) as nornir_obj:
-            nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
-
-            logger.debug("Run nornir compliance tasks.")
-            nr_with_processors.run(
-                task=run_compliance,
-                name="RENDER COMPLIANCE TASK GROUP",
-                logger=logger,
-                device_to_settings_map=device_to_settings_map,
-                rules=rules,
-            )
-
-    except Exception as error:
-        error_msg = f"`E3001:` General Exception handler, original error message ```{error}```"
-        logger.error(error_msg)
-        raise NornirNautobotException(error_msg) from error
-
+        },
+    ) as nornir_obj:
+        nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
+
+        logger.debug("Run nornir compliance tasks.")
+        nr_with_processors.run(
+            task=run_compliance,
+            name="RENDER COMPLIANCE TASK GROUP",
+            logger=logger,
+            device_to_settings_map=device_to_settings_map,
+            rules=rules,
+        )
     logger.debug("Completed compliance job for devices.")
diff --git a/nautobot_golden_config/nornir_plays/config_intended.py b/nautobot_golden_config/nornir_plays/config_intended.py
@@ -110,35 +110,28 @@ def config_intended(job_result, log_level, job_class_instance, qs, device_to_set
 
     # Retrieve filters from the Django jinja template engine
     jinja_env = get_django_env()
-
-    try:
-        with InitNornir(
-            runner=NORNIR_SETTINGS.get("runner"),
-            logging={"enabled": False},
-            inventory={
-                "plugin": "nautobot-inventory",
-                "options": {
-                    "credentials_class": NORNIR_SETTINGS.get("credentials"),
-                    "params": NORNIR_SETTINGS.get("inventory_params"),
-                    "queryset": qs,
-                    "defaults": {"now": now},
-                },
+    with InitNornir(
+        runner=NORNIR_SETTINGS.get("runner"),
+        logging={"enabled": False},
+        inventory={
+            "plugin": "nautobot-inventory",
+            "options": {
+                "credentials_class": NORNIR_SETTINGS.get("credentials"),
+                "params": NORNIR_SETTINGS.get("inventory_params"),
+                "queryset": qs,
+                "defaults": {"now": now},
             },
-        ) as nornir_obj:
-            nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
-
-            logger.debug("Run nornir render config tasks.")
-            # Run the Nornir Tasks
-            nr_with_processors.run(
-                task=run_template,
-                name="RENDER CONFIG",
-                logger=logger,
-                device_to_settings_map=device_to_settings_map,
-                job_class_instance=job_class_instance,
-                jinja_env=jinja_env,
-            )
-
-    except Exception as error:
-        error_msg = f"`E3001:` General Exception handler, original error message ```{error}```"
-        logger.error(error_msg)
-        raise NornirNautobotException(error_msg) from error
+        },
+    ) as nornir_obj:
+        nr_with_processors = nornir_obj.with_processors([ProcessGoldenConfig(logger)])
+
+        logger.debug("Run nornir render config tasks.")
+        # Run the Nornir Tasks
+        nr_with_processors.run(
+            task=run_template,
+            name="RENDER CONFIG",
+            logger=logger,
+            device_to_settings_map=device_to_settings_map,
+            job_class_instance=job_class_instance,
+            jinja_env=jinja_env,
+        )