1919
2020from .config import testgres_config
2121
22- from .connection import \
23- NodeConnection , \
24- InternalError , \
25- ProgrammingError
22+ from .connection import NodeConnection
2623
2724from .consts import \
2825DATA_DIR , \
@@ -541,10 +538,10 @@ def append_conf(self, line='', filename=PG_CONF_FILE, **kwargs):
541538 This instance of :class:`.PostgresNode`.
542539
543540 Examples:
544- append_conf(fsync=False)
545- append_conf('log_connections = yes')
546- append_conf(random_page_cost=1.5, fsync=True, ...)
547- append_conf('postgresql.conf', 'synchronous_commit = off')
541+ >>> append_conf(fsync=False)
542+ >>> append_conf('log_connections = yes')
543+ >>> append_conf(random_page_cost=1.5, fsync=True, ...)
544+ >>> append_conf('postgresql.conf', 'synchronous_commit = off')
548545 """
549546
550547lines = [line ]
@@ -815,7 +812,8 @@ def psql(self,
815812filename = None ,
816813dbname = None ,
817814username = None ,
818- input = None ):
815+ input = None ,
816+ ** variables ):
819817"""
820818 Execute a query using psql.
821819
@@ -825,9 +823,15 @@ def psql(self,
825823 dbname: database name to connect to.
826824 username: database user name.
827825 input: raw input to be passed.
826+ **variables: vars to be set before execution.
828827
829828 Returns:
830829 A tuple of (code, stdout, stderr).
830+
831+ Examples:
832+ >>> psql('select 1')
833+ >>> psql('postgres', 'select 2')
834+ >>> psql(query='select 3', ON_ERROR_STOP=1)
831835 """
832836
833837# Set default arguments
@@ -846,6 +850,10 @@ def psql(self,
846850dbname
847851 ]# yapf: disable
848852
853+ # set variables before execution
854+ for key ,value in iteritems (variables ):
855+ psql_params .extend (["--set" ,'{}={}' .format (key ,value )])
856+
849857# select query source
850858if query :
851859psql_params .extend (("-c" ,query ))
@@ -877,10 +885,15 @@ def safe_psql(self, query=None, **kwargs):
877885 username: database user name.
878886 input: raw input to be passed.
879887
888+ **kwargs are passed to psql().
889+
880890 Returns:
881891 psql's output as str.
882892 """
883893
894+ # force this setting
895+ kwargs ['ON_ERROR_STOP' ]= 1
896+
884897ret ,out ,err = self .psql (query = query ,** kwargs )
885898if ret :
886899raise QueryException ((err or b'' ).decode ('utf-8' ),query )
@@ -980,8 +993,7 @@ def poll_query_until(self,
980993sleep_time = 1 ,
981994expected = True ,
982995commit = True ,
983- raise_programming_error = True ,
984- raise_internal_error = True ):
996+ suppress = None ):
985997"""
986998 Run a query once per second until it returns 'expected'.
987999 Query should return a single value (1 row, 1 column).
@@ -994,13 +1006,13 @@ def poll_query_until(self,
9941006 sleep_time: how much should we sleep after a failure?
9951007 expected: what should be returned to break the cycle?
9961008 commit: should (possible) changes be committed?
997- raise_programming_error: enable ProgrammingError?
998- raise_internal_error: enable InternalError?
1009+ suppress: a collection of exceptions to be suppressed.
9991010
10001011 Examples:
1001- poll_query_until('select true')
1002- poll_query_until('postgres', "select now() > '01.01.2018'")
1003- poll_query_until('select false', expected=True, max_attempts=4)
1012+ >>> poll_query_until('select true')
1013+ >>> poll_query_until('postgres', "select now() > '01.01.2018'")
1014+ >>> poll_query_until('select false', expected=True, max_attempts=4)
1015+ >>> poll_query_until('select 1', suppress={testgres.OperationalError})
10041016 """
10051017
10061018# sanity checks
@@ -1032,13 +1044,8 @@ def poll_query_until(self,
10321044elif expected is None :
10331045return # done
10341046
1035- except ProgrammingError as e :
1036- if raise_programming_error :
1037- raise e
1038-
1039- except InternalError as e :
1040- if raise_internal_error :
1041- raise e
1047+ except tuple (suppress or []):
1048+ pass # we're suppressing them
10421049
10431050time .sleep (sleep_time )
10441051attempts += 1
@@ -1229,13 +1236,14 @@ def pgbench_run(self, dbname=None, username=None, options=[], **kwargs):
12291236 options: additional options for pgbench (list).
12301237
12311238 **kwargs: named options for pgbench.
1232- Examples:
1233- pgbench_run(initialize=True, scale=2)
1234- pgbench_run(time=10)
12351239 Run pgbench --help to learn more.
12361240
12371241 Returns:
12381242 Stdout produced by pgbench.
1243+
1244+ Examples:
1245+ >>> pgbench_run(initialize=True, scale=2)
1246+ >>> pgbench_run(time=10)
12391247 """
12401248
12411249# Set default arguments