Optimize memory usage (#86)

* Return generator of hosts

* Don't call keys() on zabbix_templates dict

* Use hosts generator wherever possible

* Limit API output

* Fetch status, proxy and inventory mode by default

* Remove redundant `dict.keys()` calls

zabbix_auto_config/processing.py

@@ -789,7 +789,7 @@ def do_update(self) -> None:
             return
         # Get all disabled hosts
         disabled_hosts = self.api.get_hosts(status=MonitoringStatus.OFF)
-        self.cleanup_maintenances(disabled_hosts)
+        self.cleanup_maintenances(list(disabled_hosts))
 
 
 class ZabbixHostUpdater(ZabbixUpdater):
@@ -917,7 +917,7 @@ def enable_host(self, db_host: models.Host) -> None:
             return
 
         try:
-            hosts = self.api.get_hosts(hostname, search=False)
+            hosts = list(self.api.get_hosts(hostname, search=False))
 
             if hosts:
                 host = hosts[0]
@@ -1185,8 +1185,8 @@ def do_update(self) -> None:
             else:
                 zabbix_managed_hosts.append(host)
 
-        db_hostnames = set(db_hosts.keys())
-        zabbix_hostnames = set(zabbix_hosts.keys())
+        db_hostnames = set(db_hosts)
+        zabbix_hostnames = set(zabbix_hosts)
         zabbix_managed_hostnames = {host.host for host in zabbix_managed_hosts}
         zabbix_manual_hostnames = {host.host for host in zabbix_manual_hosts}
 
@@ -1486,22 +1486,21 @@ def do_update(self) -> None:
             zabbix_templates[zabbix_template.host] = zabbix_template
 
         managed_template_names = managed_template_names.intersection(
-            set(zabbix_templates.keys())
+            set(zabbix_templates)
         )  # If the template isn't in zabbix we can't manage it
 
         # Get hosts from DB
         db_hosts = self.get_db_hosts()
 
         # Get hosts from Zabbix
-        _hosts = self.api.get_hosts(
+        zabbix_hosts = self.api.get_hosts(
             status=MonitoringStatus.ON,
             flags=0,
             select_groups=True,
             select_templates=True,
         )
-        zabbix_hosts = {host.host: host for host in _hosts}
 
-        for zabbix_hostname, zabbix_host in zabbix_hosts.items():
+        for zabbix_host in zabbix_hosts:
             if self.stop_event.is_set():
                 logging.debug("Told to stop. Breaking")
                 break
@@ -1514,21 +1513,19 @@ def do_update(self) -> None:
                 continue
 
             # Disabled hosts are not managed
-            if zabbix_hostname not in db_hosts:
+            if not (db_host := db_hosts.get(zabbix_host.host)):
                 logging.debug(
                     "Skipping host (It is not enabled in the database): %s", zabbix_host
                 )
                 continue
 
-            db_host = db_hosts[zabbix_hostname]
-
             # Determine managed templates
             synced_template_names: Set[str] = set()
             for prop in db_host.properties:
                 if template_names := self.property_template_map.get(prop):
                     synced_template_names.update(template_names)
             synced_template_names = synced_template_names.intersection(
-                set(zabbix_templates.keys())
+                set(zabbix_templates)  # list of dict keys
             )  # If the template isn't in zabbix we can't manage it
 
             host_templates: Dict[str, Template] = {}
@@ -1539,32 +1536,32 @@ def do_update(self) -> None:
             host_templates_to_remove: Dict[str, Template] = {}
 
             # Update templates on host
-            for template_name in list(host_templates.keys()):
+            for template_name in list(host_templates):
                 if (
                     template_name in managed_template_names
                     and template_name not in synced_template_names
                 ):
                     logging.debug(
                         "Going to remove template '%s' from host '%s'.",
                         template_name,
-                        zabbix_hostname,
+                        zabbix_host.host,
                     )
                     host_templates_to_remove[template_name] = host_templates[
                         template_name
                     ]
                     del host_templates[template_name]
             for template_name in synced_template_names:
-                if template_name not in host_templates.keys():
+                if template_name not in host_templates:
                     logging.debug(
                         "Going to add template '%s' to host '%s'.",
                         template_name,
-                        zabbix_hostname,
+                        zabbix_host.host,
                     )
                     host_templates[template_name] = zabbix_templates[template_name]
             if host_templates != old_host_templates:
                 logging.info(
                     "Updating templates on host '%s'. Old: %s. New: %s",
-                    zabbix_hostname,
+                    zabbix_host.host,
                     ", ".join(old_host_templates.keys()),
                     ", ".join(host_templates.keys()),
                 )
@@ -1744,16 +1741,14 @@ def do_update(self) -> None:
         db_hosts = self.get_db_hosts()
 
         # Get hosts from Zabbix
-        _hosts = self.api.get_hosts(
+        zabbix_hosts = self.api.get_hosts(
             status=MonitoringStatus.ON,
             flags=0,
             select_groups=True,
             select_templates=True,
         )
-        zabbix_hosts = {host.host: host for host in _hosts}
-
         # Iterate over hosts in Zabbix and update synced hosts
-        for zabbix_hostname, zabbix_host in zabbix_hosts.items():
+        for zabbix_host in zabbix_hosts:
             if self.stop_event.is_set():
                 logging.debug("Told to stop. Breaking")
                 break
@@ -1766,13 +1761,13 @@ def do_update(self) -> None:
                 continue
 
             # Disabled hosts are not managed
-            if zabbix_hostname not in db_hosts:
+            if zabbix_host.host not in db_hosts:
                 logging.debug(
                     "Skipping host (It is not enabled in the database): %s", zabbix_host
                 )
                 continue
 
-            db_host = db_hosts[zabbix_hostname]
+            db_host = db_hosts[zabbix_host.host]
 
             # Determine host groups to sync for host
             # Sync host groups derived from its properties, siteadmins, sources, etc.
@@ -1803,7 +1798,7 @@ def do_update(self) -> None:
                 host_hostgroups[zabbix_hostgroup.name] = zabbix_hostgroup
             old_host_hostgroups = host_hostgroups.copy()
 
-            for hostgroup_name in list(host_hostgroups.keys()):
+            for hostgroup_name in list(host_hostgroups):
                 # TODO: Here lies a bug due to managed_hostgroup_names not being properly updated above?
                 # NOTE (pederhan): Not sure what this refers to?
                 if (
@@ -1842,7 +1837,7 @@ def do_update(self) -> None:
             if sorted(host_hostgroups) != sorted(old_host_hostgroups):
                 logging.info(
                     "Updating host groups on host '%s'. Old: %s. New: %s",
-                    zabbix_hostname,
+                    zabbix_host.host,
                     # Just re-compute here (it's cheap enough)
                     ", ".join(sorted(old_host_hostgroups)),
                     ", ".join(sorted(host_hostgroups)),

zabbix_auto_config/pyzabbix/client.py

@@ -19,6 +19,7 @@
 from typing import TYPE_CHECKING
 from typing import Any
 from typing import Dict
+from typing import Iterator
 from typing import List
 from typing import Literal
 from typing import MutableMapping
@@ -618,6 +619,7 @@ def get_host(
             status=status,
             agent_status=agent_status,
         )
+        hosts = list(hosts)  # consume the iterator
         if not hosts:
             raise ZabbixNotFoundError(
                 f"Host {name_or_id!r} not found. Check your search pattern and filters."
@@ -645,8 +647,7 @@ def get_hosts(
         search: Optional[
             bool
         ] = True,  # we generally always want to search when multiple hosts are requested
-        # **filter_kwargs,
-    ) -> List[Host]:
+    ) -> Iterator[Host]:
         """Fetch all hosts matching the given criteria(s).
 
         Hosts can be filtered by name or ID. Names and IDs cannot be mixed.
@@ -680,7 +681,9 @@ def get_hosts(
         Returns:
             List[Host]: _description_
         """
-        params: ParamsType = {"output": "extend"}
+        params: ParamsType = {
+            "output": ["hostid", "host", "proxyid", "status", "inventory_mode"]
+        }
         filter_params: ParamsType = {}
         search_params: ParamsType = {}
 
@@ -733,9 +736,9 @@ def get_hosts(
             # still returns the result under the "groups" property
             # even if we use the new 6.2 selectHostGroups param
             param = compat.param_host_get_groups(self.version)
-            params[param] = "extend"
+            params[param] = ["groupid", "name"]
         if select_templates:
-            params["selectParentTemplates"] = "extend"
+            params["selectParentTemplates"] = ["templateid", "host"]
         if select_inventory:
             params["selectInventory"] = "extend"
         if select_macros:
@@ -751,7 +754,8 @@ def get_hosts(
 
         resp: List[Any] = self.host.get(**params) or []
         # TODO add result to cache
-        return [Host(**resp) for resp in resp]
+        for r in resp:
+            yield Host.model_validate(r)
 
     def create_host(
         self,
@@ -1454,7 +1458,7 @@ def get_templates(
         select_parent_templates: bool = False,
     ) -> List[Template]:
         """Fetch one or more templates given a name or ID."""
-        params: ParamsType = {"output": "extend"}
+        params: ParamsType = {"output": ["templateid", "host"]}
         search_params: ParamsType = {}
 
         # TODO: refactor this along with other methods that take names or ids (or wildcards)
@@ -1474,11 +1478,11 @@ def get_templates(
         if search_params:
             params["search"] = search_params
         if select_hosts:
-            params["selectHosts"] = "extend"
+            params["selectHosts"] = ["hostid", "host"]
         if select_templates:
-            params["selectTemplates"] = "extend"
+            params["selectTemplates"] = ["templateid", "host"]
         if select_parent_templates:
-            params["selectParentTemplates"] = "extend"
+            params["selectParentTemplates"] = ["templateid", "host"]
 
         try:
             templates = self.template.get(**params)

zabbix_auto_config/pyzabbix/types.py

@@ -224,7 +224,7 @@ class Template(ZabbixAPIBaseModel):
 class TemplateGroup(ZabbixAPIBaseModel):
     groupid: str
     name: str
-    uuid: str
+    uuid: str = ""
     templates: List[Template] = []