11# coding: utf-8
22
3+ import atexit
4+ import copy
5+ import shutil
6+ import tempfile
37
4- class TestgresConfig :
8+ from contextlib import contextmanager
9+
10+
11+ class GlobalConfig (object ):
512 """
613 Global config (override default settings).
714
815 Attributes:
9- cache_initdb: shall we use cached initdb instance?
10- cached_initdb_dir: shall we create a temp dir for cached initdb?
16+ cache_initdb: shall we use cached initdb instance?
17+ cached_initdb_dir: shall we create a temp dir for cached initdb?
1118
12- cache_pg_config: shall we cache pg_config results?
19+ cache_pg_config: shall we cache pg_config results?
1320
14- use_python_logging: use python logging configuration for all nodes.
15- error_log_lines: N of log lines to be included into exception (0=inf).
21+ use_python_logging: use python logging configuration for all nodes.
22+ error_log_lines: N of log lines to be shown in exception (0=inf).
1623
17- node_cleanup_full: shall we remove EVERYTHING (including logs)?
24+ node_cleanup_full: shall we remove EVERYTHING (including logs)?
1825 node_cleanup_on_good_exit: remove base_dir on nominal __exit__().
1926 node_cleanup_on_bad_exit: remove base_dir on __exit__() via exception.
27+
28+ NOTE: attributes must not be callable or begin with __.
2029 """
2130
2231 cache_initdb = True
23- cached_initdb_dir = None
32+ _cached_initdb_dir = None
2433
2534 cache_pg_config = True
2635
@@ -31,12 +40,128 @@ class TestgresConfig:
3140 node_cleanup_on_good_exit = True
3241 node_cleanup_on_bad_exit = False
3342
43+ @property
44+ def cached_initdb_dir (self ):
45+ return self ._cached_initdb_dir
46+
47+ @cached_initdb_dir .setter
48+ def cached_initdb_dir (self , value ):
49+ self ._cached_initdb_dir = value
50+
51+ if value :
52+ cached_initdb_dirs .add (value )
53+
54+ def __init__ (self , ** options ):
55+ self .update (options )
56+
57+ def __setitem__ (self , key , value ):
58+ setattr (self , key , value )
59+
60+ def __getitem__ (self , key ):
61+ return getattr (self , key )
62+
63+ def __setattr__ (self , name , value ):
64+ if name not in self .keys ():
65+ raise TypeError ('Unknown option {}' .format (name ))
66+
67+ super (GlobalConfig , self ).__setattr__ (name , value )
68+
69+ def keys (self ):
70+ keys = []
71+
72+ for key in dir (GlobalConfig ):
73+ if not key .startswith ('__' ) and not callable (self [key ]):
74+ keys .append (key )
75+
76+ return keys
77+
78+ def items (self ):
79+ return ((key , self [key ]) for key in self .keys ())
80+
81+ def update (self , config ):
82+ for key , value in config .items ():
83+ self [key ] = value
84+
85+ return self
86+
87+ def copy (self ):
88+ return copy .copy (self )
89+
90+
91+ # cached dirs to be removed
92+ cached_initdb_dirs = set ()
93+
94+ # default config object
95+ testgres_config = GlobalConfig ()
96+
97+ # NOTE: for compatibility
98+ TestgresConfig = testgres_config
99+
100+ # stack of GlobalConfigs
101+ config_stack = [testgres_config ]
102+
103+
104+ def rm_cached_initdb_dirs ():
105+ for d in cached_initdb_dirs :
106+ shutil .rmtree (d , ignore_errors = True )
107+
108+
109+ def push_config (** options ):
110+ """
111+ Permanently set custom GlobalConfig options
112+ and put previous settings on top of stack.
113+ """
114+
115+ # push current config to stack
116+ config_stack .append (testgres_config .copy ())
117+
118+ return testgres_config .update (options )
119+
120+
121+ def pop_config ():
122+ """
123+ Set previous GlobalConfig options from stack.
124+ """
125+
126+ if len (config_stack ) <= 1 :
127+ raise IndexError ('Reached initial config' )
128+
129+ # restore popped config
130+ return testgres_config .update (config_stack .pop ())
131+
132+
133+ @contextmanager
134+ def scoped_config (** options ):
135+ """
136+ Temporarily set custom GlobalConfig options for this context.
137+
138+ Example:
139+ >>> with scoped_config(cache_initdb=False):
140+ ... with get_new_node().init().start() as node:
141+ ... print(node.execute('select 1'))
142+ """
143+
144+ # set a new config with options
145+ config = push_config (** options )
146+ try :
147+ # return it
148+ yield config
149+ finally :
150+ # restore previous config
151+ pop_config ()
152+
34153
35154def configure_testgres (** options ):
36155 """
37- Configure testgres .
38- Look at TestgresConfig to check what can be changed .
156+ Adjust current global options .
157+ Look at GlobalConfig to learn what can be set .
39158 """
40159
41- for key , option in options .items ():
42- setattr (TestgresConfig , key , option )
160+ testgres_config .update (options )
161+
162+
163+ # NOTE: to be executed at exit()
164+ atexit .register (rm_cached_initdb_dirs )
165+
166+ # NOTE: assign initial cached dir for initdb
167+ testgres_config .cached_initdb_dir = tempfile .mkdtemp ()
0 commit comments