From 0a05102bc16a4f6a992e38cbed03d18b1731602f Mon Sep 17 00:00:00 2001 From: James Campbell Date: Wed, 8 Jan 2025 17:01:11 -0500 Subject: [PATCH] Build out more unit tests and fix bugs --- pgmon.py | 70 +++++++-- pgmon.yml | 27 ++-- test_pgmon.py | 420 +++++++++++++++++++++++++++++++++++++++++++++++++- 3 files changed, 483 insertions(+), 34 deletions(-) diff --git a/pgmon.py b/pgmon.py index c232dd2..33e9848 100755 --- a/pgmon.py +++ b/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': 'debug', + # Log level for stderr logging + '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: - cfg = yaml.safe_load(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): diff --git a/pgmon.yml b/pgmon.yml index 9b3fd42..ba746b6 100644 --- a/pgmon.yml +++ b/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: 'debug', +# Log level for stderr logging (see https://docs.python.org/3/library/logging.html#logging-levels) +# 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: {} diff --git a/test_pgmon.py b/test_pgmon.py index e9b630b..4b35fa3 100644 --- a/test_pgmon.py +++ b/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') +