from lettuce import world, before, after import os.path import psycopg2 import requests import subprocess import shutil import tempfile import time import yaml class PatroniController(object): PATRONI_CONFIG = '{}.yml' """ starts and stops individual patronis""" def __init__(self): self._output_dir = None self._patroni_path = None self._connections = {} self._config = {} self._connstring = {} self._cursors = {} self._log = {} self._processes = {} @property def patroni_path(self): if self._patroni_path is None: cwd = os.path.realpath(__file__) while True: path, entry = os.path.split(cwd) cwd = path if entry == 'features' or cwd == '/': break self._patroni_path = cwd return self._patroni_path def start(self, pg_name, max_wait_limit=15): if not self._is_running(pg_name): if pg_name in self._processes: del self._processes[pg_name] cwd = self.patroni_path self._log[pg_name] = open(os.path.join(self._output_dir, 'patroni_{0}.log'.format(pg_name)), 'a') self._config[pg_name] = self._make_patroni_test_config(pg_name) p = subprocess.Popen(['python', 'patroni.py', self._config[pg_name]], stdout=self._log[pg_name], stderr=subprocess.STDOUT, cwd=cwd) if not (p and p.pid and p.poll() is None): assert False, "PostgreSQL {0} is not running after being started".format(pg_name) self._processes[pg_name] = p # wait while patroni is available for queries, but not more than 10 seconds. for _ in range(max_wait_limit): if self.query(pg_name, "SELECT 1", fail_ok=True) is not None: break time.sleep(1) else: assert False,\ "Patroni instance is not available for queries after {0} seconds".format(max_wait_limit) def stop(self, pg_name, kill=False): while self._is_running(pg_name): if not kill: self._processes[pg_name].terminate() else: self._processes[pg_name].kill() time.sleep(1) if self._log.get('pg_name') and not self._log['pg_name'].closed: self._log[pg_name].close() if pg_name in self._processes: del self._processes[pg_name] def query(self, pg_name, query, fail_ok=False): try: cursor = self._cursor(pg_name) cursor.execute(query) return cursor except psycopg2.Error: if fail_ok: return None else: raise def check_role_has_changed_to(self, pg_name, new_role, timeout=10): bound_time = time.time() + timeout recovery_status = False if new_role == 'primary' else True role_has_changed = False while not role_has_changed: cur = self.query(pg_name, "SELECT pg_is_in_recovery()", fail_ok=True) if cur: row = cur.fetchone() if row and len(row) > 0 and row[0] == recovery_status: role_has_changed = True if time.time() > bound_time: break time.sleep(1) return role_has_changed def stop_all(self): for patroni in self._processes.copy(): self.stop(patroni) def create_and_set_output_directory(self, feature_name): feature_dir = os.path.join(pctl.patroni_path, "features", "output", feature_name.encode('utf-8').replace(' ', '_')) if os.path.exists(feature_dir): shutil.rmtree(feature_dir) os.makedirs(feature_dir) self._output_dir = feature_dir def _is_running(self, pg_name): return pg_name in self._processes and self._processes[pg_name].pid and (self._processes[pg_name].poll() is None) def _make_patroni_test_config(self, pg_name): patroni_config_name = PatroniController.PATRONI_CONFIG.format(pg_name) patroni_config_path = os.path.join(self._output_dir, patroni_config_name) with open(patroni_config_name) as f: config = yaml.load(f) postgresql = config['postgresql'] postgresql['name'] = pg_name.encode('utf-8') postgresql_params = postgresql['parameters'] postgresql_params['logging_collector'] = 'on' postgresql_params['log_destination'] = 'csvlog' postgresql_params['log_directory'] = self._output_dir postgresql_params['log_filename'] = '{0}.log'.format(pg_name) postgresql_params['log_statement'] = 'all' postgresql_params['log_min_messages'] = 'debug1' with open(patroni_config_path, 'w') as f: yaml.dump(config, f, default_flow_style=False) return patroni_config_path def _make_connstring(self, pg_name): if pg_name in self._connstring: return self._connstring[pg_name] try: patroni_path = self.patroni_path with open(os.path.join(patroni_path, PatroniController.PATRONI_CONFIG.format(pg_name)), 'r') as f: config = yaml.load(f) except OSError: return None connstring = config['postgresql']['connect_address'] if ':' in connstring: address, port = connstring.split(':') else: address = connstring port = '5432' user = "postgres" dbname = "postgres" self._connstring[pg_name] = "host={0} port={1} dbname={2} user={3}".format(address, port, dbname, user) return self._connstring[pg_name] def _connection(self, pg_name): if pg_name not in self._connections or self._connections[pg_name].closed: conn = psycopg2.connect(self._make_connstring(pg_name)) conn.autocommit = True self._connections[pg_name] = conn return self._connections[pg_name] def _cursor(self, pg_name): if pg_name not in self._cursors or self._cursors[pg_name].closed: cursor = self._connection(pg_name).cursor() self._cursors[pg_name] = cursor return self._cursors[pg_name] class EtcdController(object): """ handles all etcd related tasks, used for the tests setup and cleanup """ ETCD_VERSION_URL = 'http://127.0.0.1:2379/version' ETCD_CLEANUP_URL = 'http://127.0.0.1:2379/v2/keys/service/batman?recursive=true' def __init__(self): self.handle = None self.work_directory = None self.pid = None self.start_timeot = 5 def start(self): """ start etcd if it's not already running """ if self._is_running(): return True self.work_directory = tempfile.mkdtemp() self.handle =\ subprocess.Popen(["etcd", "--data-dir", self.work_directory], stdout=subprocess.PIPE, stderr=subprocess.PIPE) start_time = time.time() while (not self._is_running()): if time.time() - start_time > self.start_timeout: assert False, "Failed to start etcd" time.sleep(1) return True def query(self, key): """ query etcd for a value of a given key """ r = requests.get("http://127.0.0.1:2379/v2/keys/service/batman/{0}".format(key)) if r.ok: content = r.json() if content: return content.get('node', {}).get('value', None) return None def stop_and_remove_work_directory(self): """ terminate etcd and wipe out the temp work directory, but only if we actually started it""" if self._is_running() and self.handle: self.handle.terminate() self.handle = None if self.work_directory: shutil.rmtree(self.work_directory) self.work_directory = None @staticmethod def cleanup_service_tree(): """ clean all contents stored in the tree used for the tests """ r = None try: r = requests.delete(EtcdController.ETCD_CLEANUP_URL) if r and not r.ok: assert False,\ "request to cleanup the etcd contents was not successfull: status code {0}".format(r.status_code) except requests.exceptions.RequestException as e: assert False, "exception when cleanin up etcd contents: {0}".format(e) def _is_running(self): # if we have already started etcd if self.handle and self.handle.pid and (self.handle.poll() is None): return True # if etcd is running, but we didn't start it try: r = requests.get(EtcdController.ETCD_VERSION_URL) running = (r and r.ok and 'etcdserver' in r.content) except requests.ConnectionError: running = False return running pctl = PatroniController() etcd_ctl = EtcdController() # export pctl to manage patroni from scenario files world.pctl = pctl world.etcd_ctl = etcd_ctl # actions to execute on start/stop of the tests and before running invidual features @before.all def start_etcd(): etcd_ctl.start() etcd_ctl.cleanup_service_tree() @after.all def stop_etcd(*args, **kwargs): etcd_ctl.stop_and_remove_work_directory() @before.each_feature def make_test_output_dir(feature): """ create per-feature output directory to collect Patroni and PostgreSQL logs """ pctl.create_and_set_output_directory(feature.name) @after.each_feature def cleanup(*args, **kwargs): """ stop all Patronis, remove their data directory and cleanup the keys in etcd """ pctl.stop_all() shutil.rmtree(os.path.join(pctl.patroni_path, 'data')) etcd_ctl.cleanup_service_tree()