2727Copyright (c) 2016, Postgres Professional
2828"""
2929
30+ import atexit
31+ import logging
3032import os
31- import subprocess
3233import pwd
34+ import select
3335import shutil
34- import time
3536import six
36- import port_for
37-
38- import threading
39- import logging
40- import select
37+ import subprocess
4138import tempfile
39+ import threading
40+ import time
41+
42+ import port_for
4243
4344from enum import Enum
4445from distutils .version import LooseVersion
4546
47+
4648# Try to use psycopg2 by default. If psycopg2 isn't available then use
4749# pg8000 which is slower but much more portable because uses only
4850# pure-Python code
6062# threads for loggers
6163util_threads = []
6264
63- # chached initdb dir
64- cached_data_dir = None
65-
6665# rows returned by PG_CONFIG
6766pg_config_data = {}
6867
@@ -80,9 +79,18 @@ class TestgresConfig:
8079 Global config (override default settings)
8180 """
8281
82+ # shall we cache pg_config results?
8383cache_pg_config = True
84+
85+ # shall we use cached initdb instance?
8486cache_initdb = True
8587
88+ # shall we create a temp dir for cached initdb?
89+ cached_initdb_dir = None
90+
91+ # shall we remove EVERYTHING (including logs)?
92+ node_cleanup_full = False
93+
8694
8795class TestgresException (Exception ):
8896"""
@@ -427,7 +435,7 @@ def __init__(self,
427435self .port = port or reserve_port ()
428436self .should_free_port = port is None
429437self .base_dir = base_dir or tempfile .mkdtemp ()
430- self .should_rm_base_dir = base_dir is None
438+ self .should_rm_dirs = base_dir is None
431439self .use_logging = use_logging
432440self .logger = None
433441
@@ -741,8 +749,6 @@ def reload(self, params=[]):
741749_params = ["reload" ,"-D" ,self .data_dir ,"-w" ]+ params
742750_execute_utility ("pg_ctl" ,_params ,self .utils_logname )
743751
744- return self
745-
746752def pg_ctl (self ,params ):
747753"""
748754 Invoke pg_ctl with params.
@@ -784,9 +790,16 @@ def cleanup(self, max_attempts=3):
784790
785791attempts += 1
786792
787- # remove data directory if necessary
788- if self .should_rm_base_dir :
789- shutil .rmtree (self .data_dir ,ignore_errors = True )
793+ # remove directory tree if necessary
794+ if self .should_rm_dirs :
795+
796+ # choose directory to be removed
797+ if TestgresConfig .node_cleanup_full :
798+ rm_dir = self .base_dir # everything
799+ else :
800+ rm_dir = self .data_dir # just data, save logs
801+
802+ shutil .rmtree (rm_dir ,ignore_errors = True )
790803
791804return self
792805
@@ -1101,16 +1114,29 @@ def call_initdb(_data_dir):
11011114call_initdb (data_dir )
11021115# Else we can use cached dir
11031116else :
1104- global cached_data_dir
1117+ # Set default temp dir for cached initdb
1118+ if TestgresConfig .cached_initdb_dir is None :
1119+ def rm_cached_data_dir (rm_dir ):
1120+ shutil .rmtree (rm_dir ,ignore_errors = True )
1121+
1122+ # Create default temp dir
1123+ TestgresConfig .cached_initdb_dir = tempfile .mkdtemp ()
11051124
1106- # Initialize cached initdb
1107- if cached_data_dir is None :
1108- cached_data_dir = tempfile .mkdtemp ()
1109- call_initdb (cached_data_dir )
1125+ # Schedule cleanup
1126+ atexit .register (rm_cached_data_dir ,
1127+ TestgresConfig .cached_initdb_dir )
11101128
11111129try :
1130+ # Fetch cached initdb dir
1131+ cached_data_dir = TestgresConfig .cached_initdb_dir
1132+
1133+ # Initialize cached initdb
1134+ if not os .listdir (cached_data_dir ):
1135+ call_initdb (cached_data_dir )
1136+
11121137# Copy cached initdb to current data dir
11131138shutil .copytree (cached_data_dir ,data_dir )
1139+
11141140except Exception as e :
11151141raise InitNodeException (str (e ))
11161142