Re-login on expired session (#89)

* Re-login on expired session * Change log level for re-login msg
unioslo · Oct 17, 2024 · a42e0d1 · a42e0d1
1 parent 16ed546
commit a42e0d1
Show file tree

Hide file tree

Showing 4 changed files with 41 additions and 18 deletions.
diff --git a/zabbix_auto_config/exceptions.py b/zabbix_auto_config/exceptions.py
@@ -52,6 +52,10 @@ def reason(self) -> str:
         return reason
 
 
+class ZabbixAPISessionExpired(ZabbixAPIRequestError):
+    """Zabbix API session expired."""
+
+
 class ZabbixAPIResponseParsingError(ZabbixAPIRequestError):
     """Zabbix API request error."""
 

diff --git a/zabbix_auto_config/processing.py b/zabbix_auto_config/processing.py
@@ -41,6 +41,7 @@
 from zabbix_auto_config.exceptions import SourceCollectorError
 from zabbix_auto_config.exceptions import SourceCollectorTypeError
 from zabbix_auto_config.exceptions import ZabbixAPIException
+from zabbix_auto_config.exceptions import ZabbixAPISessionExpired
 from zabbix_auto_config.exceptions import ZabbixNotFoundError
 from zabbix_auto_config.exceptions import ZACException
 from zabbix_auto_config.failsafe import check_failsafe
@@ -103,6 +104,13 @@ def run(self) -> None:
                         logging.error("Timeout exception: %s", str(e))
                     elif isinstance(e, ZACException):
                         logging.error("Work exception: %s", str(e))
+                    elif isinstance(e, ZabbixAPISessionExpired):
+                        logging.error("API Session expired: %s", str(e))
+                        if isinstance(self, ZabbixUpdater):
+                            logging.info(
+                                "Reconnecting to Zabbix API and retrying update"
+                            )
+                            self.login()
                     elif isinstance(e, ZabbixAPIException):
                         logging.error("API exception: %s", str(e))
                     else:
@@ -631,6 +639,16 @@ def __init__(
 
         self.update_interval = 60  # default. Overriden in subclasses
 
+        self.property_template_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "property_template_map.txt")
+        )
+        self.property_hostgroup_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "property_hostgroup_map.txt")
+        )
+        self.siteadmin_hostgroup_map = utils.read_map_file(
+            os.path.join(self.config.map_dir, "siteadmin_hostgroup_map.txt")
+        )
+
         pyzabbix_logger = logging.getLogger("pyzabbix")
         pyzabbix_logger.setLevel(logging.ERROR)
 
@@ -639,6 +657,13 @@ def __init__(
             timeout=self.config.timeout,  # timeout for connect AND read
             read_only=self.config.dryrun,  # prevent accidental changes
         )
+
+        self.login()
+        ver = self.api.apiinfo.version()
+        self.zabbix_version = Version(ver)
+        logging.info("Connected to Zabbix API version: %s", ver)
+
+    def login(self) -> None:
         try:
             self.api.login(self.config.username, self.config.password)
         except httpx.ConnectError as e:
@@ -653,20 +678,6 @@ def __init__(
             logging.error("Unable to login to Zabbix API: %s", str(e))
             raise ZACException(*e.args)
 
-        self.property_template_map = utils.read_map_file(
-            os.path.join(self.config.map_dir, "property_template_map.txt")
-        )
-        self.property_hostgroup_map = utils.read_map_file(
-            os.path.join(self.config.map_dir, "property_hostgroup_map.txt")
-        )
-        self.siteadmin_hostgroup_map = utils.read_map_file(
-            os.path.join(self.config.map_dir, "siteadmin_hostgroup_map.txt")
-        )
-
-        ver = self.api.apiinfo.version()
-        self.zabbix_version = Version(ver)
-        logging.info("Connected to Zabbix API version: %s", ver)
-
     def work(self) -> None:
         start_time = time.time()
         logging.info("Zabbix update starting")

diff --git a/zabbix_auto_config/pyzabbix/client.py b/zabbix_auto_config/pyzabbix/client.py
@@ -34,6 +34,7 @@
 from zabbix_auto_config.exceptions import ZabbixAPIException
 from zabbix_auto_config.exceptions import ZabbixAPIRequestError
 from zabbix_auto_config.exceptions import ZabbixAPIResponseParsingError
+from zabbix_auto_config.exceptions import ZabbixAPISessionExpired
 from zabbix_auto_config.exceptions import ZabbixNotFoundError
 from zabbix_auto_config.pyzabbix import compat
 from zabbix_auto_config.pyzabbix.enums import AgentAvailable
@@ -296,8 +297,13 @@ def do_request(
             # some errors don't contain 'data': workaround for ZBX-9340
             if not resp.error.data:
                 resp.error.data = "No data"
-            raise ZabbixAPIRequestError(
-                f"Error: {resp.error.message} {resp.error.data}",
+            msg = f"Error: {resp.error.message} {resp.error.data}"
+            if "re-login" in msg:
+                cls = ZabbixAPISessionExpired
+            else:
+                cls = ZabbixAPIRequestError
+            raise cls(
+                msg,
                 api_response=resp,
                 response=response,
             )
@@ -753,7 +759,9 @@ def get_hosts(
             params["sortorder"] = sort_order
 
         resp: List[Any] = self.host.get(**params) or []
-        # TODO add result to cache
+
+        # Instantiate one at the time when iterating
+        # which should avoid some memory pressure when there are many hosts
         for r in resp:
             yield Host.model_validate(r)
 

diff --git a/zabbix_auto_config/pyzabbix/types.py b/zabbix_auto_config/pyzabbix/types.py
@@ -147,7 +147,7 @@ class ZabbixAPIResponse(BaseModel):
 
     jsonrpc: str
     id: int
-    result: Any = None  # can subclass this and specify types (ie. ZabbixAPIListResponse, ZabbixAPIStrResponse, etc.)
+    result: Any = None
     """Result of API call, if request succeeded."""
     error: Optional[ZabbixAPIError] = None
     """Error info, if request failed."""