From f773898b71b192753f89f8a5f8405b78e3011a39 Mon Sep 17 00:00:00 2001
From: MOIS3Y <s.zhukovskii@ispsystem.com>
Date: Thu, 6 Jun 2024 20:59:03 +0900
Subject: [PATCH] Update: BaseAPI class now has error handler

---
 mgrctl/api/base.py     | 67 ++++++++++++++++++++++++++++--------------
 mgrctl/settings/api.py |  4 +++
 2 files changed, 49 insertions(+), 22 deletions(-)

diff --git a/mgrctl/api/base.py b/mgrctl/api/base.py
index 6a0fd04..a7fe3c8 100644
--- a/mgrctl/api/base.py
+++ b/mgrctl/api/base.py
@@ -1,14 +1,18 @@
 import sys
 import json
+import click
 import urllib
 import requests
 
+from time import sleep
+
 from mgrctl.settings.api import (
     API_URL,
     API_HEADERS,
     API_EMAIL,
     API_PASSWORD,
-    API_VERIFY_SSL
+    API_VERIFY_SSL,
+    API_COUNT_TRY_CONNECTIONS
 )
 
 
@@ -27,47 +31,66 @@ class BaseAPI(object):
         return f'{self.API_URL}/{self.API_DEFINITION}/{self.API_VERSION}{url}'
 
     def call_api(self, url, method='GET', headers={}, data={}):
-        # Open session
-        with requests.Session() as session:
+        attempt = API_COUNT_TRY_CONNECTIONS
+        while attempt:
             try:
-                url = self._gen_request_url(url)
+                uri = self._gen_request_url(url)
                 headers = self.API_HEADERS if not headers else headers
                 params_str = urllib.parse.urlencode(data, safe="+'()")
                 if method == 'POST':
-                    api_request = session.post(
-                        url=url,
+                    api_request = requests.post(
+                        url=uri,
                         json=data,
                         headers=headers,
                         verify=self.API_VERIFY_SSL
                     )
                 if method == 'GET':
-                    url = f'{url}?{params_str}' if params_str else url
-                    api_request = session.get(
-                        url=url,
+                    uri = f'{uri}?{params_str}' if params_str else uri
+                    api_request = requests.get(
+                        url=uri,
                         headers=headers,
                         verify=self.API_VERIFY_SSL
                     )
             except Exception as error:
-                api_request = {
-                    'result': False,
-                    'error': type(error).__name__
-                }
-                return api_request
-            finally:
-                session.close()
+                click.echo(f'Error: {type(error).__name__}')
+                sys.exit()
 
-        # Get response
+            # Get response:
+            response = self._parse_response(api_request)
+
+            # Validate response:
+            if self._error_handler(response):
+                attempt -= 1
+                sleep(2)  # wait 2 second timeout
+                continue  # new attempt connection
+
+            return response
+
+    def _parse_response(self, api_request):
         try:
             response = json.loads(api_request.text)
-            if 'error' in response and response['error']:
-                print(response['error'])
-                raise sys.exit()
             return response
         except json.decoder.JSONDecodeError:
-            response = {'error': 'Can not parse response'}
-            print(response)
+            click.echo('Error: Invalid API response')
             raise sys.exit()
 
+    def _is_error(self, response):
+        if response.get('error'):
+            return True
+
+    def _is_error_3004(self, error):
+        if error.get('code') == 3004:
+            return True
+
+    def _error_handler(self, response):
+        if self._is_error(response):
+            error = response.get('error')
+            if self._is_error_3004(error):
+                return True
+            else:
+                click.echo(f'Error: API return {response}')
+                raise sys.exit()
+
 
 class BaseAuthAPI(BaseAPI):
     def __init__(self):
diff --git a/mgrctl/settings/api.py b/mgrctl/settings/api.py
index 28610c8..b3addd3 100644
--- a/mgrctl/settings/api.py
+++ b/mgrctl/settings/api.py
@@ -1,5 +1,6 @@
 from requests.packages import urllib3
 
+from mgrctl.settings.environment import env
 from mgrctl.settings.platform import (
     PLATFORM_TYPE,
     PLATFORM_VERIFY_SSL,
@@ -26,6 +27,9 @@ API_HEADERS = {"Internal-Auth": "on", "Accept": "application/json"}
 # Alias for import:
 API_VERIFY_SSL = PLATFORM_VERIFY_SSL
 
+# API 3004 Unavailable error handler:
+API_COUNT_TRY_CONNECTIONS = env.int('API_COUNT_TRY_CONNECTIONS', 3)
+
 # Suppress warning from urllib3:
 if not PLATFORM_VERIFY_SSL_WARNING:
     # ! This is not recommended,