diff --git a/databuilder/loader/file_system_neo4j_csv_loader.py b/databuilder/loader/file_system_neo4j_csv_loader.py index 1a1b3ca8b..860af185b 100644 --- a/databuilder/loader/file_system_neo4j_csv_loader.py +++ b/databuilder/loader/file_system_neo4j_csv_loader.py @@ -26,10 +26,12 @@ class FsNeo4jCSVLoader(Loader): # Config keys NODE_DIR_PATH = 'node_dir_path' RELATION_DIR_PATH = 'relationship_dir_path' + FORCE_CREATE_DIR = 'force_create_directory' SHOULD_DELETE_CREATED_DIR = 'delete_created_directories' _DEFAULT_CONFIG = ConfigFactory.from_dict({ - SHOULD_DELETE_CREATED_DIR: True + SHOULD_DELETE_CREATED_DIR: True, + FORCE_CREATE_DIR: False }) def __init__(self): @@ -55,6 +57,7 @@ def init(self, conf): self._delete_created_dir = \ conf.get_bool(FsNeo4jCSVLoader.SHOULD_DELETE_CREATED_DIR) + self._force_create_dir = conf.get_bool(FsNeo4jCSVLoader.FORCE_CREATE_DIR) self._create_directory(self._node_dir) self._create_directory(self._relation_dir) @@ -67,8 +70,11 @@ def _create_directory(self, path): :return: """ if os.path.exists(path): - raise RuntimeError( - 'Directory should not exist: {}'.format(path)) + if self._force_create_dir: + LOGGER.info('Directory exist. Deleting directory {}'.format(path)) + shutil.rmtree(path) + else: + raise RuntimeError('Directory should not exist: {}'.format(path)) os.makedirs(path) diff --git a/databuilder/utils/closer.py b/databuilder/utils/closer.py index f04231487..49eef4328 100644 --- a/databuilder/utils/closer.py +++ b/databuilder/utils/closer.py @@ -1,3 +1,5 @@ +import atexit + from typing import Callable, List # noqa: F401 @@ -13,6 +15,7 @@ class Closer(object): def __init__(self): # type: () -> None self._stack = [] # type: List + atexit.register(self.close) def register(self, close_callable): # type: (Callable) -> None diff --git a/setup.py b/setup.py index 86796318e..b2a698b64 100644 --- a/setup.py +++ b/setup.py @@ -1,7 +1,7 @@ from setuptools import setup, find_packages -__version__ = '1.0.0' +__version__ = '1.0.1' setup( name='amundsen-databuilder',