From 3696f815977e56216635e38430f6c45ec3b86f48 Mon Sep 17 00:00:00 2001
From: "Joshua M. Boniface" <joshua@boniface.me>
Date: Wed, 22 Feb 2023 13:08:54 -0500
Subject: [PATCH] Add PostgreSQL monitoring check

---
 node-daemon/plugins/psql | 139 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 139 insertions(+)
 create mode 100644 node-daemon/plugins/psql

diff --git a/node-daemon/plugins/psql b/node-daemon/plugins/psql
new file mode 100644
index 00000000..ec37a4e3
--- /dev/null
+++ b/node-daemon/plugins/psql
@@ -0,0 +1,139 @@
+#!/usr/bin/env python3
+
+# psql.py - PVC Monitoring example plugin for Postgres/Patroni
+# Part of the Parallel Virtual Cluster (PVC) system
+#
+#    Copyright (C) 2018-2022 Joshua M. Boniface <joshua@boniface.me>
+#
+#    This program is free software: you can redistribute it and/or modify
+#    it under the terms of the GNU General Public License as published by
+#    the Free Software Foundation, version 3.
+#
+#    This program is distributed in the hope that it will be useful,
+#    but WITHOUT ANY WARRANTY; without even the implied warranty of
+#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+#    GNU General Public License for more details.
+#
+#    You should have received a copy of the GNU General Public License
+#    along with this program.  If not, see <https://www.gnu.org/licenses/>.
+#
+###############################################################################
+
+# This script provides an example of a PVC monitoring plugin script. It will create
+# a simple plugin to check the Patroni PostgreSQL instance on the node for operation.
+
+# This script can thus be used as an example or reference implementation of a
+# PVC monitoring pluginscript and expanded upon as required.
+
+# A monitoring plugin script must implement the class "MonitoringPluginScript" which
+# extends "MonitoringPlugin", providing the 3 functions indicated. Detailed explanation
+# of the role of each function is provided in context of the example; see the other
+# examples for more potential uses.
+
+# WARNING:
+#
+# This script will run in the context of the node daemon keepalives as root.
+# DO NOT install untrusted, unvetted plugins under any circumstances.
+
+
+# This import is always required here, as MonitoringPlugin is used by the
+# MonitoringPluginScript class
+from pvcnoded.objects.MonitoringInstance import MonitoringPlugin
+
+
+# A monitoring plugin script must always expose its nice name, which must be identical to
+# the file name
+PLUGIN_NAME = "psql"
+
+
+# The MonitoringPluginScript class must be named as such, and extend MonitoringPlugin.
+class MonitoringPluginScript(MonitoringPlugin):
+    def setup(self):
+        """
+        setup(): Perform special setup steps during node daemon startup
+
+        This step is optional and should be used sparingly.
+        """
+
+        pass
+
+    def run(self):
+        """
+        run(): Perform the check actions and return a PluginResult object
+        """
+
+        # Run any imports first
+        from psycopg2 import connect
+
+        conn_metadata = None
+        cur_metadata = None
+        conn_dns = None
+        cur_dns = None
+
+        # Set the health delta to 0 (no change)
+        health_delta = 0
+        # Craft a message that can be used by the clients
+        message = "Successfully connected to PostgreSQL databases on localhost"
+
+        # Check the Metadata database (primary)
+        try:
+            conn_metadata = connect(
+                host='127.0.0.1',
+                port=self.config["metadata_postgresql_port"],
+                dbname=self.config["metadata_postgresql_dbname"],
+                user=self.config["metadata_postgresql_user"],
+                password=self.config["metadata_postgresql_password"],
+            )
+            cur_metadata = conn_metadata.cursor()
+            cur_metadata.execute("""SELECT * FROM alembic_version""")
+            data = cur_metadata.fetchone()
+        except Exception as e:
+            health_delta = 50
+            err = str(e).split('\n')[0]
+            message = f"Failed to connect to PostgreSQL database {self.config['metadata_postgresql_dbname']}: {err}"
+        finally:
+            if cur_metadata is not None:
+                cur_metadata.close()
+            if conn_metadata is not None:
+                conn_metadata.close()
+
+        if health_delta == 0:
+            # Check the PowerDNS database (secondary)
+            try:
+                conn_pdns = connect(
+                    host='127.0.0.1',
+                    port=self.config["pdns_postgresql_port"],
+                    dbname=self.config["pdns_postgresql_dbname"],
+                    user=self.config["pdns_postgresql_user"],
+                    password=self.config["pdns_postgresql_password"],
+                )
+                cur_pdns = conn_pdns.cursor()
+                cur_pdns.execute("""SELECT * FROM supermasters""")
+                data = cur_pdns.fetchone()
+            except Exception as e:
+                health_delta = 50
+                err = str(e).split('\n')[0]
+                message = f"Failed to connect to PostgreSQL database {self.config['pdns_postgresql_dbname']}: {err}"
+            finally:
+                if cur_pdns is not None:
+                    cur_pdns.close()
+                if conn_pdns is not None:
+                    conn_pdns.close()
+
+        # Set the health delta in our local PluginResult object
+        self.plugin_result.set_health_delta(health_delta)
+
+        # Set the message in our local PluginResult object
+        self.plugin_result.set_message(message)
+
+        # Return our local PluginResult object
+        return self.plugin_result
+
+    def cleanup(self):
+        """
+        cleanup(): Perform special cleanup steps during node daemon termination
+
+        This step is optional and should be used sparingly.
+        """
+
+        pass