Build out more unit tests and fix bugs

pgmon.py

@@ -3,6 +3,7 @@
 import yaml
 import json
 import time
+import os
 
 import argparse
 import logging
@@ -64,6 +65,8 @@ class DisconnectedError(Exception):
     pass
 class UnhappyDBError(Exception):
     pass
+class MetricVersionError(Exception):
+    pass
 
 # Default config settings
 default_config = {
@@ -76,8 +79,8 @@ default_config = {
     # How long a connection can sit idle in the pool before it's removed (seconds)
     'max_idle_time': 30,
 
-    # Log level for stderr logging (or 'off')
+    # Log level for stderr logging
-    'log_level': 'debug',
+    'log_level': 'error',
 
     # Database user to connect as
     'dbuser': 'postgres',
@@ -149,10 +152,51 @@ def read_config(path, included = False):
     # Read config file
     log.info(f"Reading log file: {path}")
     with open(path, 'r') as f:
+        try:
             cfg = yaml.safe_load(f)
+        except yaml.parser.ParserError as e:
+            raise ConfigError(f"Inavlid config file: {path}: {e}")
+
+    # Since we use it a few places, get the base directory from the config
+    config_base = os.path.dirname(path)
+
+    # Read any external queries and validate metric definitions
+    for name, metric in cfg.get('metrics', {}).items():
+        # Validate return types
+        try:
+            if metric['type'] not in ['value', 'row', 'column', 'set']:
+                raise ConfigError(f"Invalid return type: {metric['type']} for metric {name} in {path}")
+        except KeyError:
+            raise ConfigError(f"No type specified for metric {name} in {path}")
+
+        # Ensure queries exist
+        query_dict = metric.get('query', {})
+        if type(query_dict) is not dict:
+            raise ConfigError(f"Query definition should be a dictionary, got: {query_dict} for metric {name} in {path}")
+
+        if len(query_dict) == 0:
+            raise ConfigError(f"Missing queries for metric {name} in {path}")
+
+        # Read external sql files and validate version keys
+        for vers, query in metric['query'].items():
+            try:
+                int(vers)
+            except:
+                raise ConfigError(f"Invalid version: {vers} for metric {name} in {path}")
+
+            if query.startswith('file:'):
+                query_path = query[5:]
+                if not query_path.startswith('/'):
+                  query_path = os.path.join(config_base, query_path)
+                with open(query_path, 'r') as f:
+                    metric['query'][vers] = f.read()
+
 
     # Read any included config files
     for inc in cfg.get('include', []):
+        # Prefix relative paths with the directory from the current config
+        if not inc.startswith('/'):
+          inc = os.path.join(config_base, inc)
         update_deep(cfg, read_config(inc, included=True))
 
     # Return the config we read if this is an include, otherwise set the final
@@ -161,21 +205,17 @@ def read_config(path, included = False):
         return cfg
     else:
         new_config = {}
-        new_config.update(default_config)
+        update_deep(new_config, default_config)
         update_deep(new_config, cfg)
 
-        # Read any external queries
-        for metric in new_config.get('metrics', {}).values():
-            for vers, query in metric['query'].items():
-                if query.startswith('file:'):
-                    path = query[5:]
-                    with open(path, 'r') as f:
-                        metric['query'][vers] = f.read()
-
         # Minor sanity checks
         if len(new_config['metrics']) == 0:
             log.error("No metrics are defined")
-            raise ConfigError()
+            raise ConfigError("No metrics defined")
+
+        # Validate the new log level before changing the config
+        if new_config['log_level'].upper() not in ['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL']:
+            raise ConfigError(f"Invalid log level: {new_config['log_level']}")
 
         global config
         config = new_config
@@ -280,9 +320,11 @@ def get_query(metric, version):
     # Select the correct query
     for v in reversed(sorted(metric['query'].keys())):
         if version >= v:
+            if len(metric['query'][v].strip()) == 0:
+                raise MetricVersionError("Metric no longer applies to PostgreSQL {version}")
             return metric['query'][v]
 
-    raise Exception('Missing metric query')
+    raise MetricVersionError('Missing metric query for PostgreSQL {version}')
 
 
 def run_query_no_retry(pool, return_type, query, args):

pgmon.yml

@@ -1,39 +1,40 @@
 # Min PostgreSQL connection pool size (per database)
-#min_pool_size: 0,
+#min_pool_size: 0
 
 # Max PostgreSQL connection pool size (per database)
-#max_pool_size: 4,
+#max_pool_size: 4
 
 # How long a connection can sit idle in the pool before it's removed (seconds)
-#max_idle_time: 30,
+#max_idle_time: 30
 
-# Log level for stderr logging (or 'off')
+# Log level for stderr logging (see https://docs.python.org/3/library/logging.html#logging-levels)
-#log_level: 'debug',
+# Possible values are: debug, info, warning, error, critical
+#log_level: 'error'
 
 # Database user to connect as
-#dbuser: 'postgres',
+#dbuser: 'postgres'
 
 # Database host
-#dbhost: '/var/run/postgresql',
+#dbhost: '/var/run/postgresql'
 
 # Database port
-#dbport: 5432,
+#dbport: 5432
 
 # Default database to connect to when none is specified for a metric
-#dbname: 'postgres',
+#dbname: 'postgres'
 
 # Timeout for getting a connection slot from a pool
-#pool_slot_timeout: 5,
+#pool_slot_timeout: 5
 
 # PostgreSQL connection timeout (seconds)
 # Note: It can actually be double this because of retries
-#connect_timeout: 5,
+#connect_timeout: 5
 
 # Time to wait before trying to reconnect again after a reconnect failure (seconds)
-#reconnect_cooldown: 30,
+#reconnect_cooldown: 30
 
 # How often to check the version of PostgreSQL (seconds)
-#version_check_period: 300,
+#version_check_period: 300
 
 # Metrics
 #metrics: {}

test_pgmon.py

@@ -1,9 +1,15 @@
 import unittest
 
 from datetime import datetime, timedelta
+import tempfile
+
+import logging
 
 import pgmon
 
+# Silence most logging output
+logging.disable(logging.CRITICAL)
+
 class TestPgmonMethods(unittest.TestCase):
   ##
   # update_deep
@@ -197,14 +203,414 @@ class TestPgmonMethods(unittest.TestCase):
     }
     self.assertEqual(pgmon.get_query(metric, 100000), 'OLD')
 
-  def test_get_query__missing_metric(self):
-    # Test getting a metric that doesn't exist
-    pass
-
   def test_get_query__missing_version(self):
+    metric = {
+      'type': 'value',
+      'query': {
+        96000: 'OLD',
+        110000: 'NEW',
+        150000: ''
+      }
+    }
+
     # Test getting a metric that only exists for newer versions
-    pass
+    self.assertRaises(pgmon.MetricVersionError, pgmon.get_query, metric, 80000)
 
     # Test getting a metric that only exists for older versions
-    pass
+    self.assertRaises(pgmon.MetricVersionError, pgmon.get_query, metric, 160000)
+
+
+  ##
+  # read_config
+  ##
+
+  def test_read_config__simple(self):
+    pgmon.config = {}
+
+    # Test reading just a metric and using the defaults for everything else
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+# This is a comment!
+metrics:
+  test1:
+    type: value
+    query:
+      0: TEST1
+""")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+    self.assertEqual(pgmon.config['max_pool_size'], pgmon.default_config['max_pool_size'])
+    self.assertEqual(pgmon.config['dbuser'], pgmon.default_config['dbuser'])
+
+    pgmon.config = {}
+
+    # Test reading a basic config
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+# This is a comment!
+min_pool_size: 1
+max_pool_size: 2
+max_idle_time: 10
+log_level: debug
+dbuser: someone
+dbhost: localhost
+dbport: 5555
+dbname: template0
+pool_slot_timeout: 1
+connect_timeout: 1
+reconnect_cooldown: 15
+version_check_period: 3600
+metrics:
+  test1:
+    type: value
+    query:
+      0: TEST1
+  test2:
+    type: set
+    query:
+      0: TEST2
+  test3:
+    type: row
+    query:
+      0: TEST3
+  test4:
+    type: column
+    query:
+      0: TEST4
+""")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+    self.assertEqual(pgmon.config['dbuser'], 'someone')
+    self.assertEqual(pgmon.config['metrics']['test1']['type'], 'value')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+    self.assertEqual(pgmon.config['metrics']['test2']['query'][0], 'TEST2')
+
+  def test_read_config__include(self):
+    pgmon.config = {}
+
+    # Test reading a config that includes other files (absolute and relative paths, multiple levels)
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write(f"""---
+# This is a comment!
+min_pool_size: 1
+max_pool_size: 2
+max_idle_time: 10
+log_level: debug
+pool_slot_timeout: 1
+connect_timeout: 1
+reconnect_cooldown: 15
+version_check_period: 3600
+include:
+  - dbsettings.yml
+  - {tmpdirname}/metrics.yml
+""")
+    
+      with open(f"{tmpdirname}/dbsettings.yml", 'w') as f:
+        f.write(f"""---
+dbuser: someone
+dbhost: localhost
+dbport: 5555
+dbname: template0
+""")
+
+      with open(f"{tmpdirname}/metrics.yml", 'w') as f:
+        f.write(f"""---
+metrics:
+  test1:
+    type: value
+    query:
+      0: TEST1
+  test2:
+    type: value
+    query:
+      0: TEST2
+include:
+  - more_metrics.yml
+""")
+
+      with open(f"{tmpdirname}/more_metrics.yml", 'w') as f:
+        f.write(f"""---
+metrics:
+  test3:
+    type: value
+    query:
+      0: TEST3
+""")
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+    
+    self.assertEqual(pgmon.config['max_idle_time'], 10)
+    self.assertEqual(pgmon.config['dbuser'], 'someone')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+    self.assertEqual(pgmon.config['metrics']['test2']['query'][0], 'TEST2')
+    self.assertEqual(pgmon.config['metrics']['test3']['query'][0], 'TEST3')
+
+  def test_read_config__reload(self):
+    pgmon.config = {}
+
+    # Test rereading a config to update an existing config
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+# This is a comment!
+min_pool_size: 1
+max_pool_size: 2
+max_idle_time: 10
+log_level: debug
+dbuser: someone
+dbhost: localhost
+dbport: 5555
+dbname: template0
+pool_slot_timeout: 1
+connect_timeout: 1
+reconnect_cooldown: 15
+version_check_period: 3600
+metrics:
+  test1:
+    type: value
+    query:
+      0: TEST1
+  test2:
+    type: value
+    query:
+      0: TEST2
+""")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+      # Just make sure the first config was read
+      self.assertEqual(len(pgmon.config['metrics']), 2)
+
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+# This is a comment!
+min_pool_size: 7
+metrics:
+  test1:
+    type: value
+    query:
+      0: NEW1
+""")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+    self.assertEqual(pgmon.config['min_pool_size'], 7)
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'NEW1')
+    self.assertEqual(len(pgmon.config['metrics']), 1)
+
+  def test_read_config__query_file(self):
+    pgmon.config = {}
+
+    # Read a config file that reads a query from a file
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+metrics:
+  test1:
+    type: value
+    query:
+      0: file:some_query.sql
+""")
+
+      with open(f"{tmpdirname}/some_query.sql", 'w') as f:
+        f.write("This is a query")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'This is a query')
+
+  def test_read_config__invalid(self):
+    pgmon.config = {}
+
+    # For all of these tests, we start with a valid config and also ensure that
+    #  it is not modified when a new config read fails
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+metrics:
+  test1:
+    type: value
+    query:
+      0: TEST1
+""")
+
+      pgmon.read_config(f"{tmpdirname}/config.yml")
+
+    # Just make sure the config was read
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test reading a nonexistant config file
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      self.assertRaises(FileNotFoundError, pgmon.read_config, f'{tmpdirname}/missing.yml')
+
+    # Test reading an invalid config file
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""[default]
+This looks a lot like an ini file to me
+
+Or maybe a TOML?
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+
+    # Test reading a config that includes an invalid file
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query:
+      0: EVIL1
+include:
+  - missing_file.yml
+""")
+      self.assertRaises(FileNotFoundError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test invalid log level
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+log_level: noisy
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query:
+      0: EVIL1
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test invalid query return type
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: lots_of_data
+    query:
+      0: EVIL1
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test invalid query dict type
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: lots_of_data
+    query: EVIL1
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test incomplete metric: missing type
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    query:
+      0: EVIL1
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test incomplete metric: missing queries
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test incomplete metric: empty queries
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query: {}
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test incomplete metric: query dict is None
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query:
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test reading a config with no metrics
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test reading a query defined in a file but the file is missing
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query: 
+      0: file:missing.sql
+""")
+      self.assertRaises(FileNotFoundError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')
+
+    # Test invalid query versions
+    with tempfile.TemporaryDirectory() as tmpdirname:
+      with open(f"{tmpdirname}/config.yml", 'w') as f:
+        f.write("""---
+dbuser: evil
+metrics:
+  test1:
+    type: value
+    query: 
+      default: EVIL1
+""")
+      self.assertRaises(pgmon.ConfigError, pgmon.read_config, f'{tmpdirname}/config.yml')
+    self.assertEqual(pgmon.config['dbuser'], 'postgres')
+    self.assertEqual(pgmon.config['metrics']['test1']['query'][0], 'TEST1')