Generate graphs when data is updated.

4 files changed, 95 insertions, 4 deletions

diff --git a/kerbana/settings.py b/kerbana/settings.py
index 5d65b47..383d86f 100644
--- a/kerbana/settings.py
+++ b/kerbana/settings.py
@@ -154,4 +154,30 @@ RRA:MAX:0.5:288:365
 RRA:MIN:0.5:288:365
 """
 
+RRD_GRAPH_CONFIG = """
+--imgformat
+  PNG
+--width
+  800
+--height
+  200
+--start
+  -86400
+--end
+  -1
+--vertical-label
+  {ds_names[0]}
+--title
+  {title}
+--right-axis
+  1:0
+--alt-autoscale
+DEF:{ds_names[0]}={ds_paths[0]}:{ds_names[0]}:AVERAGE
+VDEF:max={ds_names[0]},MAXIMUM
+VDEF:min={ds_names[0]},MINIMUM
+LINE2:{ds_names[0]}#0000FF
+LINE1:max#FF0000
+LINE1:min#FF0000
+"""
+
 from .config import *  # noqa
diff --git a/rrd/models.py b/rrd/models.py
index 8d995c8..1d1e5e9 100644
--- a/rrd/models.py
+++ b/rrd/models.py
@@ -80,6 +80,9 @@ class DataSource(ModelWithPerms):
         except ValueError as e:
             log.warning("Could not update ds: %s", e)
 
+        for graph in self.graph_set.all():
+            graph.update()
+
 
 class Graph(ModelWithPerms):
     title = models.CharField(max_length=64)
@@ -89,11 +92,35 @@ class Graph(ModelWithPerms):
         recursive=True,
         max_length=512,
     )
-    rrd_config = models.TextField()
+    rrd_config = models.TextField(
+        default=settings.RRD_GRAPH_CONFIG
+    )
 
     def __str__(self):
         return self.title
 
+    def update(self):
+        graph_path = os.path.join(settings.RRD_GRAPH_PATH, self.path)
+        os.makedirs(os.path.dirname(graph_path), exist_ok=True)
+        rrd_paths = []
+        rrd_topics = []
+        rrd_ds_names = []
+        for ds in self.data_sources.all():
+            rrd_paths.append(os.path.join(settings.RRD_DB_PATH, ds.path))
+            rrd_topics.append(ds.topic)
+            rrd_ds_names.append(ds.topic.split("/")[-1])
+        opts = self.rrd_config.format(
+            topics=rrd_topics,
+            ds_names=rrd_ds_names,
+            ds_paths=rrd_paths,
+            title=self.title,
+        ).strip().split('\n')
+        opts = [o.strip() for o in opts]
+        rrdtool.graph(
+            graph_path,
+            * opts
+        )
+
 
 class Dashboard(ModelWithPerms):
     title = models.CharField(max_length=64)
diff --git a/rrd/mqtt.py b/rrd/mqtt.py
index 09589c3..94bd772 100644
--- a/rrd/mqtt.py
+++ b/rrd/mqtt.py
@@ -66,10 +66,9 @@ class MQTTClient:
         topic = msg.topic.removeprefix(settings.MQTT_TOPIC)
         dss = models.DataSource.objects.filter(topic=topic)
         if not dss:
-            dss = [models.DataSource.objects.create(
+            dss = (models.DataSource.objects.create(
                 topic=topic,
                 path=slugify.slugify(topic) + ".rrd",
-            )]
-            dss[0].save()
+            ),)
         for ds in dss:
             ds.update(ts, msg.payload.decode())
diff --git a/rrd/tests/test_graphs.py b/rrd/tests/test_graphs.py
new file mode 100644
index 0000000..f2a4948
--- /dev/null
+++ b/rrd/tests/test_graphs.py
@@ -0,0 +1,39 @@
+import datetime
+import os
+
+from django.conf import settings
+from django.test import TestCase
+
+from .. import models
+
+
+class TestGraphs(TestCase):
+    def setUp(self):
+        self.rrd_file = "test.rrd"
+        try:
+            os.remove(os.path.join(settings.RRD_DB_PATH, self.rrd_file))
+        except FileNotFoundError:
+            pass
+
+    def tearDown(self):
+        try:
+            os.remove(os.path.join(settings.RRD_DB_PATH, self.rrd_file))
+        except FileNotFoundError:
+            pass
+
+    def test_generate_at_ds_update(self):
+        ds = models.DataSource.objects.create(
+            topic="test",
+            path=self.rrd_file,
+        )
+        graph = models.Graph.objects.create(
+            title="Test Graph",
+            path="test/test.png",
+        )
+        graph.data_sources.add(ds)
+        graph.save()
+        now = datetime.datetime.now()
+        ts = int(now.timestamp())
+        ds.update(ts, 10)
+        stat = os.stat(os.path.join(settings.RRD_GRAPH_PATH, "test/test.png"))
+        self.assertGreaterEqual(stat.st_mtime, now.timestamp())