From 7cdf887de2bba45bf3d4e6ccf57c798ee0a335a9 Mon Sep 17 00:00:00 2001
From: Marek Zbroch <marek.zbroch@networktocode.com>
Date: Tue, 23 Nov 2021 13:07:40 +0100
Subject: [PATCH] Added HTTP retries for Nautobot API

---
 .../adapters/nautobot_api/adapter.py            | 14 ++++++++++++--
 .../adapters/nautobot_api/inventory.py          | 17 ++++++++++++++---
 .../adapters/nautobot_api/settings.py           |  2 ++
 network_importer/adapters/nautobot_api/tasks.py | 10 ++++++++++
 4 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/network_importer/adapters/nautobot_api/adapter.py b/network_importer/adapters/nautobot_api/adapter.py
index a13e08b3..a2f27bea 100644
--- a/network_importer/adapters/nautobot_api/adapter.py
+++ b/network_importer/adapters/nautobot_api/adapter.py
@@ -3,9 +3,11 @@
 import warnings
 
 import pynautobot
+from diffsync.exceptions import ObjectAlreadyExists, ObjectNotFound
 from packaging.version import Version, InvalidVersion
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
-from diffsync.exceptions import ObjectAlreadyExists, ObjectNotFound
 import network_importer.config as config  # pylint: disable=import-error
 from network_importer.adapters.base import BaseAdapter  # pylint: disable=import-error
 from network_importer.adapters.nautobot_api.models import (  # pylint: disable=import-error
@@ -17,8 +19,8 @@
     NautobotPrefix,
     NautobotVlan,
 )
-from network_importer.adapters.nautobot_api.tasks import query_device_info_from_nautobot
 from network_importer.adapters.nautobot_api.settings import InventorySettings, AdapterSettings
+from network_importer.adapters.nautobot_api.tasks import query_device_info_from_nautobot
 
 warnings.filterwarnings("ignore", category=DeprecationWarning)
 
@@ -104,6 +106,14 @@ def load(self):
         else:
             self.nautobot.http_session.verify = True
 
+        if inventory_settings.http_retries:
+            retries = Retry(total=inventory_settings.http_retries,
+                            backoff_factor=0.5,
+                            status_forcelist=[429, 500, 502, 503, 504, ],
+                            allowed_methods=False
+                            )
+            self.nautobot.http_session.mount(self.nautobot.base_url, HTTPAdapter(max_retries=retries))
+
         self._check_nautobot_version()
 
         sites = {}
diff --git a/network_importer/adapters/nautobot_api/inventory.py b/network_importer/adapters/nautobot_api/inventory.py
index 4533b7f8..86025ee9 100644
--- a/network_importer/adapters/nautobot_api/inventory.py
+++ b/network_importer/adapters/nautobot_api/inventory.py
@@ -4,14 +4,17 @@
 
 import sys
 from typing import Any, List
-import pynautobot
-from pydantic import ValidationError
 
+import pynautobot
 from nornir.core.inventory import Defaults, Groups, Hosts, Inventory, ParentGroups, ConnectionOptions
 from nornir.core.plugins.inventory import InventoryPluginRegister
+from pydantic import ValidationError
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
+
+from network_importer.adapters.nautobot_api.settings import InventorySettings
 from network_importer.inventory import NetworkImporterInventory, NetworkImporterHost
 from network_importer.utils import build_filter_params
-from network_importer.adapters.nautobot_api.settings import InventorySettings
 
 
 class NautobotAPIInventory(NetworkImporterInventory):
@@ -51,6 +54,14 @@ def __init__(self, *args, **kwargs: Any,) -> None:
         if not self.settings.verify_ssl:
             self.session.http_session.verify = False
 
+        if self.settings.http_retries:
+            retries = Retry(total=self.settings.http_retries,
+                            backoff_factor=0.5,
+                            status_forcelist=[429, 500, 502, 503, 504, ],
+                            allowed_methods=False
+                            )
+            self.session.http_session.mount(self.session.base_url, HTTPAdapter(max_retries=retries))
+
     def load(self):
         """Load inventory by fetching devices from nautobot."""
         if self.filter_parameters:
diff --git a/network_importer/adapters/nautobot_api/settings.py b/network_importer/adapters/nautobot_api/settings.py
index 1734ac3c..adfc6d11 100644
--- a/network_importer/adapters/nautobot_api/settings.py
+++ b/network_importer/adapters/nautobot_api/settings.py
@@ -18,6 +18,7 @@ class InventorySettings(BaseSettings):
     address: Optional[str] = "http://localhost"
     token: Optional[str] = None
     verify_ssl: Union[bool, str] = True
+    http_retries: Optional[int] = 0
 
     use_primary_ip: Optional[bool] = True
     fqdn: Optional[str] = None
@@ -30,4 +31,5 @@ class Config:
             "address": {"env": "NAUTOBOT_ADDRESS"},
             "token": {"env": "NAUTOBOT_TOKEN"},
             "verify_ssl": {"env": "NAUTOBOT_VERIFY_SSL"},
+            "http_retries": {"env": "NAUTOBOT_HTTP_RETRY"},
         }
diff --git a/network_importer/adapters/nautobot_api/tasks.py b/network_importer/adapters/nautobot_api/tasks.py
index cf4fbd62..f1324a8a 100644
--- a/network_importer/adapters/nautobot_api/tasks.py
+++ b/network_importer/adapters/nautobot_api/tasks.py
@@ -3,6 +3,8 @@
 
 import pynautobot
 from nornir.core.task import Result, Task
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 
 import network_importer.config as config  # pylint: disable=import-error
 from network_importer.adapters.nautobot_api.settings import InventorySettings
@@ -34,6 +36,14 @@ def query_device_info_from_nautobot(task: Task) -> Result:
     else:
         nautobot.http_session.verify = True
 
+    if inventory_settings.http_retries:
+        retries = Retry(total=inventory_settings.http_retries,
+                        backoff_factor=0.5,
+                        status_forcelist=[429, 500, 502, 503, 504, ],
+                        allowed_methods=False
+                        )
+        nautobot.http_session.mount(nautobot.base_url, HTTPAdapter(max_retries=retries))
+
     # Set a Results dictionary
     results = {
         "device": None,