From cdba6aa24f66c3ad4e71ebfbe7303ddebe4bbd4a Mon Sep 17 00:00:00 2001
From: Ian Fijolek <ian@iamthefij.com>
Date: Mon, 28 Aug 2023 13:54:24 -0700
Subject: [PATCH] Add script to check for missing services

---
 scripts/nomad_missing_services.py | 70 +++++++++++++++++++++++++++++++
 1 file changed, 70 insertions(+)
 create mode 100755 scripts/nomad_missing_services.py

diff --git a/scripts/nomad_missing_services.py b/scripts/nomad_missing_services.py
new file mode 100755
index 0000000..79b0aa9
--- /dev/null
+++ b/scripts/nomad_missing_services.py
@@ -0,0 +1,70 @@
+#! /usr/bin/env python3
+from os import environ
+from typing import Any
+from typing import cast
+
+import requests
+
+
+NOMAD_ADDR = environ.get("NOMAD_ADDR", "http://127.0.0.1:4646")
+NOMAD_TOKEN = environ.get("NOMAD_TOKEN")
+
+
+def nomad_req(
+    *path: str, params: dict[str, Any] | None = None, method="GET"
+) -> list[dict[str, Any]] | dict[str, Any] | str:
+    headers = {}
+    if NOMAD_TOKEN:
+        headers["X-Nomad-Token"] = NOMAD_TOKEN
+
+    response = requests.request(
+        method,
+        f"{NOMAD_ADDR}/v1/{'/'.join(path)}",
+        params=params,
+        headers=headers,
+    )
+    response.raise_for_status()
+
+    try:
+        return response.json()
+    except requests.exceptions.JSONDecodeError:
+        return response.text
+
+
+def extract_job_services(job: dict[str, Any]) -> set[str]:
+    services: set[str] = set()
+    for group in job["TaskGroups"]:
+        for service in group.get("Services") or []:
+            services.add(service["Name"])
+        for task in group["Tasks"]:
+            for service in task.get("Services") or []:
+                services.add(service["Name"])
+
+    return services
+
+exit_code = 0
+
+for job in nomad_req("jobs"):
+    job = cast(dict[str, Any], job)
+
+    if job["Type"] in ("batch", "sysbatch"):
+        continue
+
+    job_detail = nomad_req("job", job["ID"])
+    job_detail = cast(dict[str, Any], job_detail)
+
+    expected_services = extract_job_services(job_detail)
+
+    found_services: set[str] = set()
+    for service in nomad_req("job", job_detail["ID"], "services"):
+        service = cast(dict[str, Any], service)
+        found_services.add(service["ServiceName"])
+
+    missing_services = expected_services - found_services
+    for missing_service in missing_services:
+        print(f"ERROR: Missing service {missing_service} for job {job_detail['Name']}")
+        print(job)
+        exit_code = 1
+
+
+exit(exit_code)