Add options and pull counts

ombi.py

@@ -1,13 +1,24 @@
 # Do not edit this script. Edit configuration.py
+import sys
 import requests
 from datetime import datetime, timezone
 from influxdb import InfluxDBClient
-
+import argparse
+from argparse import RawTextHelpFormatter
 import configuration
 
-current_time = datetime.now(timezone.utc).astimezone().isoformat()
-
 headers = {'Apikey': configuration.ombi_api_key}
+
+def now_iso():
+    now_iso = datetime.now(timezone.utc).astimezone().isoformat()
+    return now_iso
+
+def influx_sender(influx_payload):
+    influx = InfluxDBClient(configuration.influxdb_url, configuration.influxdb_port, configuration.influxdb_username,
+                            configuration.influxdb_password, configuration.ombi_influxdb_db_name)
+    influx.write_points(influx_payload)
+
+def get_total_requests():
     get_tv_requests = requests.get('{}/api/v1/Request/tv'.format(configuration.ombi_url), headers=headers).json()
     get_movie_requests = requests.get('{}/api/v1/Request/movie'.format(configuration.ombi_url), headers=headers).json()
 
@@ -24,16 +35,53 @@ influx_payload = [
         {
             "measurement": "Ombi",
             "tags": {
-            "type": "Requests"
+                "type": "Request_Total"
             },
-        "time": current_time,
+            "time": now_iso(),
             "fields": {
                 "total": count_movie_requests + count_tv_requests
             }
         }
     ]
+    return influx_payload
 
-influx = InfluxDBClient(configuration.grafana_url, configuration.grafana_port, configuration.grafana_username,
+def get_request_counts():
-                        configuration.grafana_password, configuration.ombi_grafana_db_name)
+    get_request_counts = requests.get('{}/api/v1/Request/count'.format(configuration.ombi_url), headers=headers).json()
-influx.write_points(influx_payload)
 
+    influx_payload = [
+        {
+        "measurement": "Ombi",
+            "tags": {
+                "type": "Request_Counts"
+            },
+            "time": now_iso(),
+            "fields": {
+                "pending": int(get_request_counts['pending']),
+                "approved": int(get_request_counts['approved']),
+                "available": int(get_request_counts['available'])
+            }
+        }
+    ]
+    return influx_payload
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(prog='Ombi stats operations',
+        description='Script to aid in data gathering from Ombi', formatter_class=RawTextHelpFormatter)
+
+    parser.add_argument("--total",  action='store_true',
+        help='Get the total count of all requests')
+
+    parser.add_argument("--counts",  action='store_true',
+        help='Get the count of pending, approved, and available requests')
+
+    opts = parser.parse_args()
+
+    if opts.total:
+        influx_sender(get_total_requests())
+
+    elif opts.counts:
+        influx_sender(get_request_counts())
+
+    elif len(sys.argv) == 1:
+        parser.print_help(sys.stderr)
+        sys.exit(1)