1010from six import raise_from ,iteritems
1111from tempfile import mkstemp ,mkdtemp
1212
13- from .enums import NodeStatus ,ProcessType
13+ from .enums import \
14+ NodeStatus , \
15+ ProcessType , \
16+ DumpFormat
1417
1518from .cache import cached_initdb
1619
5558StartNodeException , \
5659TimeoutException , \
5760InitNodeException , \
58- TestgresException
61+ TestgresException , \
62+ BackupException
5963
6064from .logger import TestgresLogger
6165
@@ -577,12 +581,13 @@ def get_control_data(self):
577581
578582return out_dict
579583
580- def start (self ,params = []):
584+ def start (self ,params = [], wait = True ):
581585"""
582586 Start this node using pg_ctl.
583587
584588 Args:
585589 params: additional arguments for pg_ctl.
590+ wait: wait until operation completes.
586591
587592 Returns:
588593 This instance of :class:`.PostgresNode`.
@@ -592,7 +597,7 @@ def start(self, params=[]):
592597get_bin_path ("pg_ctl" ),
593598"-D" ,self .data_dir ,
594599"-l" ,self .pg_log_file ,
595- "-w" ,# wait
600+ "-w" if wait else '-W' ,#--wait or --no- wait
596601"start"
597602 ]+ params # yapf: disable
598603
@@ -607,12 +612,13 @@ def start(self, params=[]):
607612
608613return self
609614
610- def stop (self ,params = []):
615+ def stop (self ,params = [], wait = True ):
611616"""
612617 Stop this node using pg_ctl.
613618
614619 Args:
615620 params: additional arguments for pg_ctl.
621+ wait: wait until operation completes.
616622
617623 Returns:
618624 This instance of :class:`.PostgresNode`.
@@ -621,7 +627,7 @@ def stop(self, params=[]):
621627_params = [
622628get_bin_path ("pg_ctl" ),
623629"-D" ,self .data_dir ,
624- "-w" ,# wait
630+ "-w" if wait else '-W' ,#--wait or --no- wait
625631"stop"
626632 ]+ params # yapf: disable
627633
@@ -681,6 +687,8 @@ def reload(self, params=[]):
681687
682688execute_utility (_params ,self .utils_log_file )
683689
690+ return self
691+
684692def pg_ctl (self ,params ):
685693"""
686694 Invoke pg_ctl with params.
@@ -812,7 +820,11 @@ def safe_psql(self, query=None, **kwargs):
812820
813821return out
814822
815- def dump (self ,filename = None ,dbname = None ,username = None ):
823+ def dump (self ,
824+ filename = None ,
825+ dbname = None ,
826+ username = None ,
827+ format = DumpFormat .Plain ):
816828"""
817829 Dump database into a file using pg_dump.
818830 NOTE: the file is not removed automatically.
@@ -821,14 +833,27 @@ def dump(self, filename=None, dbname=None, username=None):
821833 filename: database dump taken by pg_dump.
822834 dbname: database name to connect to.
823835 username: database user name.
836+ format: format argument plain/custom/directory/tar.
824837
825838 Returns:
826839 Path to a file containing dump.
827840 """
828841
842+ # Check arguments
843+ if not isinstance (format ,DumpFormat ):
844+ try :
845+ format = DumpFormat (format )
846+ except ValueError :
847+ msg = 'Invalid format "{}"' .format (format )
848+ raise BackupException (msg )
849+
850+ # Generate tmpfile or tmpdir
829851def tmpfile ():
830- fd ,fname = mkstemp (prefix = TMP_DUMP )
831- os .close (fd )
852+ if format == DumpFormat .Directory :
853+ fname = mkdtemp (prefix = TMP_DUMP )
854+ else :
855+ fd ,fname = mkstemp (prefix = TMP_DUMP )
856+ os .close (fd )
832857return fname
833858
834859# Set default arguments
@@ -842,7 +867,8 @@ def tmpfile():
842867"-h" ,self .host ,
843868"-f" ,filename ,
844869"-U" ,username ,
845- "-d" ,dbname
870+ "-d" ,dbname ,
871+ "-F" ,format .value
846872 ]# yapf: disable
847873
848874execute_utility (_params ,self .utils_log_file )
@@ -854,12 +880,29 @@ def restore(self, filename, dbname=None, username=None):
854880 Restore database from pg_dump's file.
855881
856882 Args:
857- filename: database dump taken by pg_dump.
883+ filename: database dump taken by pg_dump in custom/directory/tar formats .
858884 dbname: database name to connect to.
859885 username: database user name.
860886 """
861887
862- self .psql (filename = filename ,dbname = dbname ,username = username )
888+ # Set default arguments
889+ dbname = dbname or default_dbname ()
890+ username = username or default_username ()
891+
892+ _params = [
893+ get_bin_path ("pg_restore" ),
894+ "-p" ,str (self .port ),
895+ "-h" ,self .host ,
896+ "-U" ,username ,
897+ "-d" ,dbname ,
898+ filename
899+ ]# yapf: disable
900+
901+ # try pg_restore if dump is binary formate, and psql if not
902+ try :
903+ execute_utility (_params ,self .utils_log_name )
904+ except ExecUtilException :
905+ self .psql (filename = filename ,dbname = dbname ,username = username )
863906
864907@method_decorator (positional_args_hack (['dbname' ,'query' ]))
865908def poll_query_until (self ,