fix/task-conn-error-handling (#168)

* draft: retry on connection error when polling task status

* change exception type

* conn drop condition is now parametrized, add condition to template api methods when applicable

* bump SDK version

Author: sbasan

.gitignore

@@ -132,4 +132,5 @@ dmypy.json
 poetry.lock
 
 # Misc
-**/.DS_Store
+**/.DS_Store
+.vscode

catalystwan/api/task_status_api.py

@@ -5,9 +5,10 @@
 import logging
 from typing import TYPE_CHECKING, List, cast
 
-from tenacity import retry, retry_if_result, stop_after_attempt, wait_fixed  # type: ignore
+from tenacity import retry, retry_if_exception_type, retry_if_result, stop_after_attempt, wait_fixed  # type: ignore
+from tenacity.retry import retry_base  # type: ignore
 
-from catalystwan.exceptions import TaskValidationError
+from catalystwan.exceptions import ManagerRequestException, TaskValidationError
 
 if TYPE_CHECKING:
     from catalystwan.session import ManagerSession
@@ -59,6 +60,7 @@ def wait_for_completed(
         failure_statuses_ids: List[OperationStatusId] = [
             OperationStatusId.FAILURE,
         ],
+        expect_conn_drop: bool = False,
     ) -> TaskResult:
         """
         Method to check subtasks statuses of the task
@@ -91,6 +93,8 @@ def wait_for_completed(
             success_statuses_ids (Union[List[OperationStatus], str]): list of positive sub-tasks statuses id's
             fails_statuses_id (Union[List[OperationStatusId], str]): list of negative sub-tasks statuses
             fails_statuses_ids (Union[List[OperationStatusId], str]): list of negative sub-tasks statuses id's
+            expect_conn_drop (bool): set to true when performing action which can result in connection drop like:
+                                     server reboot or server configuration change
 
         Returns:
             TaskResult(): result attr is True if all subtasks are success
@@ -136,10 +140,14 @@ def check_status(task_data: List[SubTaskData]) -> bool:
         def log_exception(self) -> None:
             logger.error("Operation status not achieved in given time")
 
+        retry_condition: retry_base = retry_if_result(check_status)
+        if expect_conn_drop:
+            retry_condition |= retry_if_exception_type(ManagerRequestException)
+
         @retry(
             wait=wait_fixed(interval_seconds),
             stop=stop_after_attempt(int(timeout_seconds / interval_seconds)),
-            retry=retry_if_result(check_status),
+            retry=retry_condition,
             retry_error_callback=log_exception,
         )
         def wait_for_action_finish() -> List[SubTaskData]:

catalystwan/api/template_api.py

@@ -44,6 +44,7 @@
 from catalystwan.typed_list import DataSequence
 from catalystwan.utils.device_model import DeviceModel
 from catalystwan.utils.dict import merge
+from catalystwan.utils.personality import Personality
 from catalystwan.utils.pydantic_field import get_extra_field
 from catalystwan.utils.template_type import TemplateType
 
@@ -190,7 +191,9 @@ def get_device_specific_variables(name: str):
         endpoint = "/dataservice/template/device/config/attachfeature"
         logger.info(f"Attaching a template: {name} to the device: {device.hostname}.")
         response = self.session.post(url=endpoint, json=payload).json()
-        task = Task(session=self.session, task_id=response["id"]).wait_for_completed(timeout_seconds=timeout_seconds)
+        task = Task(session=self.session, task_id=response["id"]).wait_for_completed(
+            timeout_seconds=timeout_seconds, expect_conn_drop=device.personality is Personality.VMANAGE
+        )
         if task.result:
             return True
         logger.warning(f"Failed to attach template: {name} to the device: {device.hostname}.")
@@ -235,7 +238,9 @@ def _attach_cli(self, name: str, device: Device, is_edited: bool = False, timeou
         endpoint = "/dataservice/template/device/config/attachcli"
         logger.info(f"Attaching a template: {name} to the device: {device.hostname}.")
         response = self.session.post(url=endpoint, json=payload).json()
-        task = Task(session=self.session, task_id=response["id"]).wait_for_completed(timeout_seconds=timeout_seconds)
+        task = Task(session=self.session, task_id=response["id"]).wait_for_completed(
+            timeout_seconds=timeout_seconds, expect_conn_drop=device.personality is Personality.VMANAGE
+        )
         if task.result:
             return True
         logger.warning(f"Failed to attach tempate: {name} to the device: {device.hostname}.")
@@ -259,7 +264,9 @@ def deatach(self, device: Device) -> bool:
         endpoint = "/dataservice/template/config/device/mode/cli"
         logger.info(f"Changing mode to cli mode for {device.hostname}.")
         response = self.session.post(url=endpoint, json=payload).json()
-        task = Task(session=self.session, task_id=response["id"]).wait_for_completed()
+        task = Task(session=self.session, task_id=response["id"]).wait_for_completed(
+            expect_conn_drop=device.personality is Personality.VMANAGE
+        )
         if task.result:
             return True
         logger.warning(f"Failed to change to cli mode for device: {device.hostname}.")
@@ -353,16 +360,13 @@ def _delete_cli_template(self, name: str) -> bool:
         return True
 
     @overload
-    def edit(self, template: FeatureTemplate) -> Any:
-        ...
+    def edit(self, template: FeatureTemplate) -> Any: ...
 
     @overload
-    def edit(self, template: CLITemplate) -> Any:
-        ...
+    def edit(self, template: CLITemplate) -> Any: ...
 
     @overload
-    def edit(self, template: DeviceTemplate) -> Any:
-        ...
+    def edit(self, template: DeviceTemplate) -> Any: ...
 
     def edit(self, template):
         template_info = self.get(template).filter(name=template.template_name).single_or_default()
@@ -394,16 +398,13 @@ def _edit_feature_template(self, template: FeatureTemplate, data: FeatureTemplat
         return response
 
     @overload
-    def create(self, template: FeatureTemplate, debug=False) -> str:
-        ...
+    def create(self, template: FeatureTemplate, debug=False) -> str: ...
 
     @overload
-    def create(self, template: DeviceTemplate) -> str:
-        ...
+    def create(self, template: DeviceTemplate) -> str: ...
 
     @overload
-    def create(self, template: CLITemplate) -> str:
-        ...
+    def create(self, template: CLITemplate) -> str: ...
 
     def create(self, template, debug: bool = False):
         if isinstance(template, list):

pyproject.toml

@@ -1,8 +1,8 @@
 [tool.poetry]
 name = "catalystwan"
-version = "0.41.3"
+version = "0.41.4"
 description = "Cisco Catalyst WAN SDK for Python"
-authors = ["kagorski <kagorski@cisco.com>"]
+authors = ["sbasan <sbasan@cisco.com>"]
 readme = "README.md"
 repository = "https://github.com/cisco-en-programmability/catalystwan-sdk"