4747pg_config_data = {}
4848
4949
50+ """
51+ Predefined exceptions
52+ """
5053class ClusterException (Exception ):pass
5154class QueryException (Exception ):pass
5255
5356
57+ """
58+ Transaction wrapper returned by Node
59+ """
60+ class NodeConnection (object ):
61+ def __init__ (self ,parent_node ,dbname ):
62+ self .parent_node = parent_node
63+
64+ self .connection = pglib .connect (
65+ database = dbname ,
66+ user = get_username (),
67+ port = parent_node .port ,
68+ host = "127.0.0.1"
69+ )
70+
71+ self .cursor = self .connection .cursor ()
72+
73+ def __enter__ (self ):
74+ return self
75+
76+ def __exit__ (self ,type ,value ,tb ):
77+ self .connection .close ()
78+
79+ def execute (self ,query ):
80+ self .cursor .execute (query )
81+ return self .cursor .fetchall ()
82+
83+ def close (self ):
84+ self .connection .close ()
85+
86+
5487class PostgresNode :
5588def __init__ (self ,name ,port ):
5689self .name = name
@@ -134,6 +167,8 @@ def init(self, allows_streaming=False):
134167
135168self .set_replication_conf ()
136169
170+ return self
171+
137172def init_from_backup (self ,root_node ,backup_name ,has_streaming = False ,hba_permit_replication = True ):
138173"""Initializes cluster from backup, made by another node"""
139174
@@ -175,7 +210,9 @@ def append_conf(self, filename, string):
175210 """
176211config_name = "%s/%s" % (self .data_dir ,filename )
177212with open (config_name ,"a" )as conf :
178- conf .write (string )
213+ conf .write ('' .join ([string ,'\n ' ]))
214+
215+ return self
179216
180217def pg_ctl (self ,command ,params ):
181218"""Runs pg_ctl with specified params
@@ -192,12 +229,17 @@ def pg_ctl(self, command, params):
192229
193230with open (self .output_filename ,"a" )as file_out , \
194231open (self .error_filename ,"a" )as file_err :
195- return subprocess .call (
232+
233+ res = subprocess .call (
196234arguments + [command ],
197235stdout = file_out ,
198236stderr = file_err
199237 )
200238
239+ if res > 0 :
240+ with open (self .error_filename ,"r" )as errfile :
241+ raise ClusterException (errfile .readlines ()[- 1 ])
242+
201243def start (self ):
202244""" Starts cluster """
203245logfile = self .logs_dir + "/postgresql.log"
@@ -206,27 +248,30 @@ def start(self):
206248"-w" :None ,
207249"-l" :logfile ,
208250 }
209- if self .pg_ctl ("start" ,params ):
210- raise ClusterException ("Cluster startup failed" )
251+ self .pg_ctl ("start" ,params )
211252
212253self .working = True
213254
255+ return self
256+
214257def stop (self ):
215258""" Stops cluster """
216259params = {
217260"-D" :self .data_dir ,
218261"-w" :None
219262 }
220- if self .pg_ctl ("stop" ,params ):
221- raise ClusterException ("Cluster stop failed" )
263+ self .pg_ctl ("stop" ,params )
222264
223265self .working = False
224266
267+ return self
268+
225269def reload (self ):
226270"""Reloads config files"""
227271params = {"-D" :self .data_dir }
228- if self .pg_ctl ("reload" ,params ):
229- raise ClusterException ("Cluster reload failed" )
272+ self .pg_ctl ("reload" ,params )
273+
274+ return self
230275
231276def cleanup (self ):
232277"""Stops cluster if needed and removes the data directory"""
@@ -238,6 +283,8 @@ def cleanup(self):
238283# remove data directory
239284shutil .rmtree (self .data_dir )
240285
286+ return self
287+
241288def psql (self ,dbname ,query ):
242289"""Executes a query by the psql
243290
@@ -291,21 +338,8 @@ def poll_query_until(self, dbname, query):
291338
292339def execute (self ,dbname ,query ):
293340"""Executes the query and returns all rows"""
294- connection = pglib .connect (
295- database = dbname ,
296- user = get_username (),
297- port = self .port ,
298- host = "127.0.0.1"
299- )
300- cur = connection .cursor ()
301-
302- cur .execute (query )
303- res = cur .fetchall ()
304-
305- cur .close ()
306- connection .close ()
307-
308- return res
341+ with self .connect (dbname )as node_con :
342+ return node_con .execute (query )
309343
310344def backup (self ,name ):
311345"""Performs pg_basebackup"""
@@ -325,6 +359,9 @@ def backup(self, name):
325359
326360return backup_path
327361
362+ def connect (self ,dbname ):
363+ return NodeConnection (parent_node = self ,dbname = dbname )
364+
328365
329366def get_username ():
330367""" Returns current user name """