CHANGELOG.md

+# [1.11.0](https://gitlab.ics.muni.cz/perun-proxy-aai/python/perun-proxy-utils/compare/v1.10.0...v1.11.0) (2023-09-08)
+
+
+### Features
+
+* check_pgsql ([76d83a5](https://gitlab.ics.muni.cz/perun-proxy-aai/python/perun-proxy-utils/commit/76d83a526628daf54cb88eb672c3096e4421d25a))
+
 # [1.10.0](https://gitlab.ics.muni.cz/perun-proxy-aai/python/perun-proxy-utils/compare/v1.9.2...v1.10.0) (2023-08-24)
 
 

README.md

@@ -113,3 +113,19 @@ python3 check_user_logins.py
   ```sh
   python3 check_privacyidea.py --help
   ```
+
+### check_pgsql.py
+
+- check connection to PostgreSQL
+- possible check with configurable query
+- to use this check you must include the postgresql extra, which will install [psycopg2-binary](https://pypi.org/project/psycopg2-binary/):
+
+  ```sh
+  pip install perun.proxy.utils[postgresql]
+  ```
+
+- for usage run:
+
+  ```sh
+  python3 check_pgpsql.py --help
+  ```

perun/proxy/utils/nagios/check_pgsql.py

+#!/usr/bin/env python3
+
+import argparse
+import re
+import sys
+import time
+
+import psycopg2
+
+# nagios return codes
+OK = 0
+WARNING = 1
+CRITICAL = 2
+UNKNOWN = 3
+
+
+def get_args():
+    parser = argparse.ArgumentParser(description="PostgreSQL connection check")
+    parser.add_argument(
+        "-H",
+        "--host",
+        type=str,
+        help="Host name or IP Address (default: 127.0.0.1)",
+        default="127.0.0.1",
+    )
+    parser.add_argument(
+        "-P",
+        "--port",
+        type=str,
+        help="Port number (default: 5432)",
+        default="5432",
+    )
+    parser.add_argument(
+        "-c",
+        "--critical",
+        type=int,
+        help="Response time to result in critical status (default: 8s)",
+        default=8,
+    )
+    parser.add_argument(
+        "-w",
+        "--warning",
+        type=int,
+        help="Response time to result in warning status (default: 2s)",
+        default=2,
+    )
+    parser.add_argument(
+        "-t",
+        "--timeout",
+        type=int,
+        help="Seconds before connection times out (default: 10s)",
+        default=10,
+    )
+    parser.add_argument(
+        "-d",
+        "--database",
+        type=str,
+        help="Database to check (default: template1)",
+        default="template1",
+    )
+    parser.add_argument(
+        "-l",
+        "--logname",
+        type=str,
+        help="Login name of user",
+    )
+    parser.add_argument(
+        "-p",
+        "--password",
+        type=str,
+        help="Password (BIG SECURITY ISSUE)",
+    )
+    parser.add_argument(
+        "-q",
+        "--query",
+        type=str,
+        help="SQL query to run. Only first column in first row will be read",
+    )
+    parser.add_argument(
+        "-o",
+        "--option",
+        type=str,
+        help="Connection parameters (keyword = value)",
+    )
+    parser.add_argument(
+        "-W",
+        "--query-warning",
+        type=threshold_type,
+        dest="query_warning",
+        help="SQL query value to result in warning status (float). "
+        "Single value or range, e.g. '20:50.",
+    )
+    parser.add_argument(
+        "-C",
+        "--query-critical",
+        type=threshold_type,
+        dest="query_critical",
+        help="SQL query value to result in critical status (float). "
+        "Single value or range, e.g. '20:50",
+    )
+
+    return parser.parse_args()
+
+
+def threshold_type(arg_value):
+    threshold_regex = re.compile(r"^\d+(\.\d+)?(:\d+(\.\d+)?)?$")
+    if not threshold_regex.match(arg_value):
+        raise argparse.ArgumentTypeError("Invalid threshold format.")
+    return arg_value
+
+
+def val_in_threshold(threshold_str, value):
+    splitted = threshold_str.split(":")
+    if len(splitted) == 1:
+        return 0 < value < float(splitted[0])
+    else:
+        return float(splitted[0]) < value < float(splitted[1])
+
+
+def do_query(connection, args):
+    try:
+        with connection.cursor() as cursor:
+            cursor.execute(args.query)
+            if len(cursor.description) < 1:
+                print("QUERY WARNING - No columns returned.")
+                return WARNING
+            rows = cursor.fetchone()
+            if len(rows) < 1:
+                print("QUERY WARNING - No rows returned.")
+                return WARNING
+            try:
+                numeric = float(rows[0])
+                if args.query_warning and not val_in_threshold(
+                    args.query_warning, numeric
+                ):
+                    print("QUERY WARNING - ", end="")
+                    state = WARNING
+                elif args.query_critical and not val_in_threshold(
+                    args.query_critical, numeric
+                ):
+                    print("QUERY CRITICAL - ", end="")
+                    state = CRITICAL
+                else:
+                    print("QUERY OK - ", end="")
+                    state = OK
+
+                print(
+                    f"'{args.query}' returned {numeric}"
+                    f"|{args.query_warning if args.query_warning else ''};"
+                    f"{args.query_critical if args.query_critical else ''};;"
+                )
+
+                if len(cursor.description) > 1:
+                    extra_info = rows[1]
+                    if extra_info:
+                        print(f"Extra info: {extra_info}")
+                return state
+
+            except ValueError:
+                print(f"QUERY CRITICAL - Is not a numeric: {rows[0]}")
+                return CRITICAL
+
+    except psycopg2.Error as e:
+        print(f"QUERY CRITICAL - Error with Query: {e}")
+        return CRITICAL
+
+
+def main():
+    args = get_args()
+
+    query_status = UNKNOWN
+    connection_args = {
+        "database": args.database,
+        "host": args.host,
+        "port": args.port,
+        "connect_timeout": args.timeout,
+    }
+    if args.logname:
+        connection_args["user"] = args.logname
+    if args.password:
+        connection_args["password"] = args.password
+    if args.option:
+        connection_args["options"] = args.option
+
+    try:
+        con_start_time = time.time()
+        with psycopg2.connect(**connection_args) as conn:
+            con_end_time = time.time()
+
+            elapsed_time = con_end_time - con_start_time
+            if elapsed_time > args.critical:
+                status = CRITICAL
+            elif elapsed_time > args.warning:
+                status = WARNING
+            else:
+                status = OK
+
+            print(
+                f"{status} check_pgsql - database {args.database} ({elapsed_time:.2f} "
+                f"sec.)|{args.warning};{args.critical};;"
+            )
+
+            if args.query:
+                query_status = do_query(conn, args)
+
+            return query_status if args.query and query_status > status else status
+
+    except psycopg2.Error as e:
+        print(f"{CRITICAL} check_pgsql - connection to {args.database} failed ({e}).")
+        return CRITICAL
+
+
+if __name__ == "__main__":
+    sys.exit(main())

setup.cfg

 [metadata]
-version = 1.10.0
+version = 1.11.0
 license_files = LICENSE
 long_description = file: README.md
 long_description_content_type = text/markdown

setup.py

@@ -22,7 +22,10 @@ setuptools.setup(
     extras_require={
         "ldap": [
             "ldap3~=2.9.1",
-        ]
+        ],
+        "postgresql": [
+            "psycopg2-binary~=2.9",
+        ],
     },
     entry_points={
         "console_scripts": [
@@ -42,6 +45,7 @@ setuptools.setup(
             "check_webserver_availability="
             "perun.proxy.utils.nagios.webserver_availability:main",
             "check_privacyidea=perun.proxy.utils.nagios.check_privacyidea:main",
+            "check_pgsql=perun.proxy.utils.nagios.check_pgsql:main"
             "metadata_expiration=perun.proxy.utils.metadata_expiration:main",
             "print_docker_versions=perun.proxy.utils.print_docker_versions:main",
             "run_version_script=perun.proxy.utils.run_version_script:main",