diff --git a/helpers/postgresql.py b/helpers/postgresql.py index 8d32d65d..11be38c0 100644 --- a/helpers/postgresql.py +++ b/helpers/postgresql.py @@ -18,63 +18,69 @@ logger = logging.getLogger(__name__) def parseurl(url): r = urlparse(url) return { - 'hostname': r.hostname, + 'host': r.hostname, 'port': r.port or 5432, - 'username': r.username, + 'user': r.username, 'password': r.password, + 'database': r.path[1:], + 'fallback_application_name': 'Governor', + 'connect_timeout': 5, } class Postgresql: def __init__(self, config): + self.config = config self.name = config['name'] self.listen_addresses, self.port = config['listen'].split(':') self.data_dir = config['data_dir'] self.replication = config['replication'] self.recovery_conf = os.path.join(self.data_dir, 'recovery.conf') + self.pid_path = os.path.join(self.data_dir, 'postmaster.pid') self._pg_ctl = 'pg_ctl -w -D ' + self.data_dir - self.config = config - - connect_address = config.get('connect_address', config['listen']) or config['listen'] + self.local_address = self.get_local_address() + connect_address = config.get('connect_address', None) or self.local_address self.connection_string = 'postgres://{username}:{password}@{connect_address}/postgres'.format( connect_address=connect_address, **self.replication) - self.conn = None - self.cursor_holder = None + self._connection = None + self._cursor_holder = None self.members = [] # list of already existing replication slots - def cursor(self): - if not self.cursor_holder: - self.conn = psycopg2.connect('postgres://{}:{}/postgres'.format( - self.listen_addresses.split(',')[0].strip(), self.port)) - self.conn.autocommit = True - self.cursor_holder = self.conn.cursor() + def get_local_address(self): + # TODO: try to get unix_socket_directory from postmaster.pid + return self.listen_addresses.split(',')[0].strip() + ':' + self.port - return self.cursor_holder + def connection(self): + if not self._connection or self._connection.closed: + self._connection = psycopg2.connect('postgres://{}/postgres'.format(self.local_address)) + self._connection.autocommit = True + return self._connection + + def _cursor(self): + if not self._cursor_holder or self._cursor_holder.closed != 0: + self._cursor_holder = self.connection().cursor() + return self._cursor_holder def disconnect(self): - try: - self.conn.close() - except: - logger.exception('Error disconnecting') + self._connection and self._connection.close() + self._connection = self._cursor_holder = None def query(self, sql, *params): max_attempts = 0 while True: try: - self.cursor().execute(sql, params) - break - except psycopg2.OperationalError as e: - if self.conn: - self.disconnect() - self.cursor_holder = None - if max_attempts > 4: - raise e + cursor = self._cursor() + cursor.execute(sql, params) + return cursor + except (psycopg2.OperationalError, psycopg2.InterfaceError) as e: + self.disconnect() max_attempts += 1 + if max_attempts >= 3: + raise e time.sleep(5) - return self.cursor() def data_directory_empty(self): return not os.path.exists(self.data_dir) or os.listdir(self.data_dir) == [] @@ -90,11 +96,11 @@ class Postgresql: pgpass = 'pgpass' with open(pgpass, 'w') as f: os.fchmod(f.fileno(), 0o600) - f.write('{hostname}:{port}:*:{username}:{password}\n'.format(**r)) + f.write('{host}:{port}:*:{user}:{password}\n'.format(**r)) try: os.environ['PGPASSFILE'] = pgpass - return os.system('pg_basebackup -R -D {data_dir} --host={hostname} --port={port} -U {username}'.format( + return os.system('pg_basebackup -R -D {data_dir} --host={host} --port={port} -U {user}'.format( data_dir=self.data_dir, **r)) == 0 finally: os.environ.pop('PGPASSFILE') @@ -111,10 +117,9 @@ class Postgresql: logger.error('Cannot start PostgreSQL because one is already running.') return False - pid_path = os.path.join(self.data_dir, 'postmaster.pid') - if os.path.exists(pid_path): - os.remove(pid_path) - logger.info('Removed %s', pid_path) + if os.path.exists(self.pid_path): + os.remove(self.pid_path) + logger.info('Removed %s', self.pid_path) ret = os.system(self._pg_ctl + ' start -o "{}"'.format(self.server_options())) == 0 ret and self.load_replication_slots() @@ -142,6 +147,9 @@ class Postgresql: return True def is_healthiest_node(self, cluster): + if self.is_leader(): + return True + if cluster.last_leader_operation - self.xlog_position() > self.config.get('maximum_lag_on_failover', 0): return False @@ -149,18 +157,19 @@ class Postgresql: if member.hostname == self.name: continue try: - member_conn = psycopg2.connect(member.address) + member_conn = psycopg2.connect(parseurl(member.address)) member_conn.autocommit = True member_cursor = member_conn.cursor() member_cursor.execute( - "SELECT %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)", (self.xlog_position(), )) - xlog_diff = member_cursor.fetchone()[0] - logger.info([self.name, member.hostname, xlog_diff]) + "SELECT pg_is_in_recovery(), %s - (pg_last_xlog_replay_location() - '0/0000000'::pg_lsn)", + (self.xlog_position(), )) + row = member_cursor.fetchone() member_cursor.close() member_conn.close() - if xlog_diff < 0: + logger.error([self.name, member.hostname, row]) + if not row[0] or row[1] < 0: return False - except psycopg2.OperationalError: + except psycopg2.Error: continue return True @@ -171,7 +180,7 @@ class Postgresql: @staticmethod def primary_conninfo(leader_url): r = parseurl(leader_url) - return 'user={username} password={password} host={hostname} port={port} sslmode=prefer sslcompression=1'.format(**r) + return 'user={user} password={password} host={host} port={port} sslmode=prefer sslcompression=1'.format(**r) def check_recovery_conf(self, leader): if not os.path.isfile(self.recovery_conf): @@ -218,7 +227,9 @@ primary_conninfo = '{}' self.replication['username']), self.replication['password']) def xlog_position(self): - return self.query("SELECT pg_last_xlog_replay_location() - '0/0000000'::pg_lsn").fetchone()[0] + return self.query("""SELECT CASE WHEN pg_is_in_recovery() + THEN pg_last_xlog_replay_location() - '0/0000000'::pg_lsn + ELSE pg_current_xlog_location() - '0/00000'::pg_lsn END""").fetchone()[0] def load_replication_slots(self): cursor = self.query("SELECT slot_name FROM pg_replication_slots WHERE slot_type='physical'") @@ -239,4 +250,4 @@ primary_conninfo = '{}' self.members = members def last_operation(self): - return self.query("SELECT pg_current_xlog_location() - '0/00000'::pg_lsn").fetchone()[0] + return self.xlog_position() diff --git a/tests/test_postgresql.py b/tests/test_postgresql.py index 297bd929..5442e58c 100644 --- a/tests/test_postgresql.py +++ b/tests/test_postgresql.py @@ -11,9 +11,18 @@ def os_system(cmd): return 0 +def false(*args, **kwargs): + return False + + +def xlog_position(): + return 1 + + class MockCursor: def __init__(self): + self.closed = 0 self.current = 0 self.results = [] @@ -21,15 +30,19 @@ class MockCursor: if sql.startswith('blabla'): raise psycopg2.OperationalError() elif sql.startswith('SELECT slot_name'): - self.results = [('blabla'), ('foobar')] + self.results = [('blabla',), ('foobar',)] elif sql.startswith('SELECT pg_current_xlog_location()'): self.results = [(0,)] - elif sql.startswith('SELECT %s - (pg_last_xlog_replay_location()'): - self.results = [(0,)] - elif sql.startswith('SELECT pg_last_xlog_replay_location()'): + elif sql.startswith('SELECT pg_is_in_recovery(), %s'): + if params[0][0] != 0: + raise psycopg2.OperationalError() + self.results = [(False, 0)] + elif sql.startswith('SELECT CASE WHEN pg_is_in_recovery()'): self.results = [(0,)] elif sql.startswith('SELECT pg_is_in_recovery()'): self.results = [(False, )] + elif sql.startswith('SELECT to_char(pg_postmaster_start_time'): + self.results = [('', True, '', '', '', False)] else: self.results = [] @@ -48,13 +61,13 @@ class MockConnect: def __init__(self): self.autocommit = False + self.closed = False def cursor(self): return MockCursor() def close(self): - if not self.autocommit: - raise psycopg2.OperationalError() + pass def psycopg2_connect(*args, **kwargs): @@ -115,17 +128,19 @@ class TestPostgresql(unittest.TestCase): def test_query(self): self.p.query('select 1') - self.p.conn.autocommit = False self.assertRaises(psycopg2.OperationalError, self.p.query, 'blabla') - self.p.query('select %s', 1) def test_is_healthiest_node(self): leader = Member('leader', 'postgres://replicator:rep-pass@127.0.0.1:5435/postgres', 28) me = Member('test0', 'postgres://replicator:rep-pass@127.0.0.1:5434/postgres', 28) other = Member('test1', 'postgres://replicator:rep-pass@127.0.0.1:5433/postgres', 28) - cluster = Cluster(leader, 0, [leader, me, other]) + cluster = Cluster(leader, 0, [me, other, leader]) self.assertTrue(self.p.is_healthiest_node(cluster)) - self.p.config['maximum_lag_on_failover'] = -1 + self.p.is_leader = false + self.assertFalse(self.p.is_healthiest_node(cluster)) + self.p.xlog_position = xlog_position + self.assertTrue(self.p.is_healthiest_node(cluster)) + self.p.config['maximum_lag_on_failover'] = -2 self.assertFalse(self.p.is_healthiest_node(cluster)) def test_is_leader(self):