Connect timeout 5 seconds for psycopg2

Refactored query and is_healthiest_node methods
query method wont retry more then 3 times if OperationalError raised
is_healthiest_node takes into account that one of the nodes could be
still a master
This commit is contained in:
Alexander Kukushkin
2015-05-23 20:21:13 +02:00
parent 34f9b666cf
commit f53c369c69
2 changed files with 78 additions and 52 deletions
+53 -42
View File
@@ -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()
+25 -10
View File
@@ -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:[email protected]:5435/postgres', 28)
me = Member('test0', 'postgres://replicator:[email protected]:5434/postgres', 28)
other = Member('test1', 'postgres://replicator:[email protected]: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):