Write failsafe hosts to file (#74)

* Write failsafe hosts to file

* Fix tests, add `config` fixture

config.sample.toml

@@ -2,8 +2,9 @@
 source_collector_dir = "path/to/source_collector_dir/"
 host_modifier_dir = "path/to/host_modifier_dir/"
 db_uri = "dbname='zac' user='zabbix' host='localhost' password='secret' port=5432 connect_timeout=2"
-health_file = "/tmp/zac_health.json"
 log_level = "DEBUG"
+health_file = "/tmp/zac_health.json"
+failsafe_file = "/tmp/zac_failsafe.json"
 
 [zabbix]
 map_dir = "path/to/map_dir/"

pyproject.toml

@@ -27,7 +27,7 @@ classifiers = [
 dependencies = [
     "multiprocessing-logging==0.3.1",
     "psycopg2>=2.9.5",
-    "pydantic>=2.0.0",
+    "pydantic>=2.6.0",
     "pyzabbix>=1.3.0",
     "requests>=1.0.0",
     "tomli>=2.0.0",

tests/conftest.py

@@ -6,6 +6,10 @@
 import pytest
 from unittest import mock
 
+import tomli
+
+from zabbix_auto_config import models
+
 
 @pytest.fixture(scope="function")
 def minimal_hosts():
@@ -102,6 +106,10 @@ def sample_config():
         yield config.read()
 
 
+@pytest.fixture(name="config")
+def config(sample_config: str) -> Iterable[models.Settings]:
+    yield models.Settings(**tomli.loads(sample_config))
+
 @pytest.fixture
 def hostgroup_map_file(tmp_path: Path) -> Iterable[Path]:
     contents = """

tests/test_processing/test_zabbixupdater.py

@@ -8,6 +8,7 @@
 from zabbix_auto_config import exceptions
 
 from zabbix_auto_config.models import ZabbixSettings
+from zabbix_auto_config.models import Settings
 from zabbix_auto_config.processing import ZabbixUpdater
 from zabbix_auto_config.state import get_manager
 
@@ -27,20 +28,21 @@ def __init__(self, *args, **kwargs) -> None:
 
 @pytest.mark.timeout(10)
 @patch("pyzabbix.ZabbixAPI", TimeoutAPI())  # mock with timeout on login
-def test_zabbixupdater_connect_timeout(mock_psycopg2_connect):
-    with pytest.raises(exceptions.ZACException) as exc_info:
-        ZabbixUpdater(
-            name="connect-timeout",
-            db_uri="",
-            state=get_manager().State(),
-            zabbix_config=ZabbixSettings(
+def test_zabbixupdater_connect_timeout(mock_psycopg2_connect, config: Settings):
+    config.zabbix = ZabbixSettings(
                 map_dir="",
                 url="",
                 username="",
                 password="",
                 dryrun=False,
                 timeout=1,
-            ),
+    )
+    with pytest.raises(exceptions.ZACException) as exc_info:
+        ZabbixUpdater(
+            name="connect-timeout",
+            db_uri="",
+            state=get_manager().State(),
+            settings=config,
         )
     assert "connect timeout" in exc_info.exconly()
 
@@ -51,26 +53,29 @@ def do_update(self):
 
 
 @pytest.mark.timeout(5)
-def test_zabbixupdater_read_timeout(tmp_path: Path, mock_psycopg2_connect):
+def test_zabbixupdater_read_timeout(
+    tmp_path: Path, mock_psycopg2_connect, config: Settings
+):
     # TODO: use mapping file fixtures from #67
     map_dir = tmp_path / "maps"
     map_dir.mkdir()
     (map_dir / "property_template_map.txt").touch()
     (map_dir / "property_hostgroup_map.txt").touch()
     (map_dir / "siteadmin_hostgroup_map.txt").touch()
 
-    process = TimeoutUpdater(
-        name="read-timeout",
-        db_uri="",
-        state=get_manager().State(),
-        zabbix_config=ZabbixSettings(
+    config.zabbix = ZabbixSettings(
             map_dir=str(map_dir),
             url="",
             username="",
             password="",
             dryrun=False,
             timeout=1,
-        ),
+    )
+    process = TimeoutUpdater(
+        name="read-timeout",
+        db_uri="",
+        state=get_manager().State(),
+        settings=config,
     )
 
     # Start the process and wait for it to be marked as unhealthy

zabbix_auto_config/__init__.py

@@ -164,23 +164,23 @@ def main():
             "zabbix-host-updater",
             state_manager.State(),
             config.zac.db_uri,
-            config.zabbix,
+            config,
         )
         processes.append(process)
 
         process = processing.ZabbixHostgroupUpdater(
             "zabbix-hostgroup-updater",
             state_manager.State(),
             config.zac.db_uri,
-            config.zabbix,
+            config,
         )
         processes.append(process)
 
         process = processing.ZabbixTemplateUpdater(
             "zabbix-template-updater",
             state_manager.State(),
             config.zac.db_uri,
-            config.zabbix,
+            config,
         )
         processes.append(process)
     except exceptions.ZACException as e:

zabbix_auto_config/models.py

@@ -2,7 +2,7 @@
 from pathlib import Path
 from typing import Any, Dict, List, Optional, Set, Tuple, Union
 
-from pydantic import BaseModel
+from pydantic import BaseModel, ValidationInfo
 from pydantic import BaseModel as PydanticBaseModel
 from pydantic import (
     ConfigDict,
@@ -84,6 +84,18 @@ class ZacSettings(ConfigBaseModel):
     db_uri: str
     health_file: Optional[Path] = None
     log_level: int = Field(logging.DEBUG, description="The log level to use.")
+    failsafe_file: Optional[Path] = None
+
+    @field_validator("failsafe_file", "health_file", mode="after")
+    @classmethod
+    def _validate_file_path(cls, v: Optional[Path], info: ValidationInfo) -> Optional[Path]:
+        if v is None:
+            return v
+        if v.exists() and v.is_dir():
+            raise ValueError(f"'{info.field_name}' cannot be a directory")
+        if not v.exists():
+            utils.make_parent_dirs(v)
+        return v
 
     @field_serializer("log_level")
     def _serialize_log_level(self, v: str) -> str:
@@ -263,3 +275,12 @@ def merge(self, other: "Host") -> None:
             self.proxy_pattern = sorted(list(proxy_patterns))[0]
         elif len(proxy_patterns) == 1:
             self.proxy_pattern = proxy_patterns.pop()
+
+
+class HostActions(BaseModel):
+    add: List[str] = []
+    remove: List[str] = []
+
+    def write_json(self, path: Path) -> None:
+        """Writes a JSON serialized representation of self to a file."""
+        utils.write_file(path, self.model_dump_json(indent=2))

zabbix_auto_config/processing.py

@@ -105,7 +105,7 @@ class SourceCollectorProcess(BaseProcess):
     def __init__(
         self,
         name: str,
-        state: dict,
+        state: State,
         module: SourceCollectorModule,
         config: models.SourceCollectorSettings,
         source_hosts_queue: multiprocessing.Queue,
@@ -555,7 +555,7 @@ def merge_sources(self):
 
 
 class ZabbixUpdater(BaseProcess):
-    def __init__(self, name, state, db_uri, zabbix_config: models.ZabbixSettings):
+    def __init__(self, name, state, db_uri, settings: models.Settings):
         super().__init__(name, state)
 
         self.db_uri = db_uri
@@ -568,7 +568,8 @@ def __init__(self, name, state, db_uri, zabbix_config: models.ZabbixSettings):
             logging.error("Unable to connect to database. Process exiting with error")
             raise exceptions.ZACException(*e.args)
 
-        self.config = zabbix_config
+        self.config = settings.zabbix
+        self.settings = settings
 
         self.update_interval = 60
 
@@ -732,6 +733,20 @@ def set_tags(self, zabbix_host, tags):
         else:
             logging.info("DRYRUN: Setting tags (%s) on host: '%s' (%s)", tags, zabbix_host["host"], zabbix_host["hostid"])
 
+
+    def write_failsafe_hosts(self, to_add: List[str], to_remove: List[str]) -> None:
+        if not self.settings.zac.failsafe_file:
+            logging.info(
+                "Unable to write failsafe hosts. No diagnostics directory configured."
+            )
+            return
+        h = models.HostActions(add=to_add, remove=to_remove)
+        h.write_json(self.settings.zac.failsafe_file)
+        logging.info(
+            "Wrote list of hosts to add and remove to %s",
+            self.settings.zac.failsafe_file,
+        )
+
     def do_update(self):
         with self.db_connection, self.db_connection.cursor() as db_cursor:
             db_cursor.execute(f"SELECT data FROM {self.db_hosts_table} WHERE data->>'enabled' = 'true'")
@@ -782,6 +797,7 @@ def do_update(self):
 
         if len(hostnames_to_remove) > self.config.failsafe or len(hostnames_to_add) > self.config.failsafe:
             logging.warning("Too many hosts to change (failsafe=%d). Remove: %d, Add: %d. Aborting", self.config.failsafe, len(hostnames_to_remove), len(hostnames_to_add))
+            self.write_failsafe_hosts(hostnames_to_add, hostnames_to_remove)
             raise exceptions.ZACException("Failsafe triggered")
 
         for hostname in hostnames_to_remove:

zabbix_auto_config/utils.py

@@ -8,6 +8,7 @@
 import re
 from typing import Dict, Iterable, List, Mapping, MutableMapping, Set, Tuple, Union
 
+
 def is_valid_regexp(pattern: str):
     try:
         re.compile(pattern)
@@ -166,3 +167,33 @@ def drain_queue(q: multiprocessing.Queue) -> None:
 def timedelta_to_str(td: datetime.timedelta) -> str:
     """Converts a timedelta to a string of the form HH:MM:SS."""
     return str(td).partition(".")[0]
+
+
+def write_file(path: Union[str, Path], content: str) -> None:
+    """Writes `content` to `path`. Ensures content ends with a newline."""
+    path = Path(path)
+    # Ensure parent dirs exist
+    make_parent_dirs(path)
+
+    try:
+        with open(path, "w") as f:
+            if not content.endswith("\n"):
+                content += "\n"
+            f.write(content)
+    except OSError as e:
+        logging.error("Failed to write to file '%s': %s", path, e)
+        raise
+
+
+def make_parent_dirs(path: Union[str, Path]) -> Path:
+    """Attempts to create all parent directories given a path.
+
+    NOTE: Intended for usage with Pydantic models, and as such it will raise
+    a ValueError instead of OSError if the directory cannot be created."""
+    path = Path(path)
+
+    try:
+        path.parent.mkdir(parents=True, exist_ok=True)
+    except OSError as e:
+        raise ValueError(f"Failed to create parent directories for {path}: {e}") from e
+    return path