From ac5a61aad41b8230594e38c080f4689964e3accf Mon Sep 17 00:00:00 2001
From: Scott Wallace <scott@wallace.sh>
Date: Sat, 10 Oct 2020 16:24:44 +0100
Subject: [PATCH] Move to use `prometheus_client`

---
 .gitignore       |  1 +
 Dockerfile       |  2 +-
 dockstat.py      | 62 +++++++++++++++++++++++++++---------------------
 requirements.txt |  1 +
 4 files changed, 38 insertions(+), 28 deletions(-)
 create mode 100644 .gitignore

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..27573b8
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+.pyenv/
diff --git a/Dockerfile b/Dockerfile
index e2dd0ef..63935e8 100644
--- a/Dockerfile
+++ b/Dockerfile
@@ -12,7 +12,7 @@ ADD requirements.txt .
 RUN python -m pip install -r requirements.txt
 
 WORKDIR /app
-ADD . /app
+ADD dockstat.py /app
 
 # Switching to a non-root user, please refer to https://aka.ms/vscode-docker-python-user-rights
 # RUN useradd appuser && chown -R appuser /app
diff --git a/dockstat.py b/dockstat.py
index 2794052..7d76a9e 100644
--- a/dockstat.py
+++ b/dockstat.py
@@ -5,23 +5,34 @@ Module to act as a Prometheus Exporter for Docker containers with a
 """
 
 import argparse
-import http.server
 import os.path
-import socketserver
 import sys
-import time
+from http.server import HTTPServer
 
 import docker
+from prometheus_client import (
+    CollectorRegistry,
+    Gauge,
+    generate_latest,
+    MetricsHandler,
+)
 
 LISTEN_PORT = 8080
 HEALTHY_STR = 'healthy'
 
 
-class SimpleHTTPRequestHandler(http.server.BaseHTTPRequestHandler):
+class HTTPHandler(MetricsHandler):
     """
     Class to encompass the requirements of a Prometheus Exporter
         for Docker containers with a healthcheck configured
     """
+
+    def __init__(self, *args, **kwargs):
+        self.docker_api = docker.APIClient()
+        self.docker_client = docker.from_env()
+        super().__init__(*args, **kwargs)
+
+    # Override built-in method
     # pylint: disable=invalid-name
     def do_GET(self):
         """
@@ -33,7 +44,6 @@ class SimpleHTTPRequestHandler(http.server.BaseHTTPRequestHandler):
         if self.path == '/healthcheck':
             self._healthcheck()
 
-
     def _healthcheck(self, message=True):
         """
         Method to return 200 or 500 response and an optional message
@@ -50,35 +60,37 @@ class SimpleHTTPRequestHandler(http.server.BaseHTTPRequestHandler):
         if message:
             self.wfile.write(b'OK')
 
-
     def _metrics(self):
         """
         Method to handle the request for metrics
         """
         self._healthcheck(message=False)
 
-        api = docker.APIClient()
+        registry = CollectorRegistry()
 
-        client = docker.from_env()
-        for container in client.containers.list():
-            now = int(round(time.time() * 1000))
-            data = api.inspect_container(container.id)
+        gauge = Gauge(
+            'container_inspect_state_health_status',
+            "Container's healthcheck value (binary)",
+            labelnames=['id', 'name', 'value'],
+            registry=registry
+        )
+
+        for container in self.docker_client.containers.list():
+            data = self.docker_api.inspect_container(container.id)
 
             try:
                 health_str = data["State"]["Health"]["Status"]
+                label_values = [
+                    container.id,
+                    container.name,
+                    health_str,
+                ]
             except KeyError:
                 pass
             else:
-                self.wfile.write(
-                    bytes(
-                        f'container_inspect_state_health_status{{'
-                        f'id="{container.id}",'
-                        f'name="{container.name}",'
-                        f'value="{health_str}"'
-                        f'}} '
-                        f'{int(health_str == HEALTHY_STR)} {now}\n'.encode()
-                    )
-                )
+                gauge.labels(*label_values).set(int(health_str == HEALTHY_STR))
+
+        self.wfile.write(generate_latest(registry))
 
 
 def healthy():
@@ -105,7 +117,6 @@ if __name__ == '__main__':
 
         return parser.parse_args()
 
-
     def main():
         """
         main()
@@ -116,11 +127,8 @@ if __name__ == '__main__':
             # Invert the sense of 'healthy' for Unix CLI usage
             return not healthy()
 
-        Handler = SimpleHTTPRequestHandler
+        HTTPServer(('', LISTEN_PORT), HTTPHandler).serve_forever()
 
-        with socketserver.TCPServer(('', LISTEN_PORT), Handler) as httpd:
-            httpd.serve_forever()
-
-        return True
+        return 0
 
     sys.exit(main())
diff --git a/requirements.txt b/requirements.txt
index 1a6c338..90b4360 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,2 +1,3 @@
 # To ensure app dependencies are ported from your virtual environment/host machine into your container, run 'pip freeze > requirements.txt' in the terminal to overwrite this file
 docker
+prometheus_client