Unverified Commit 2ad85cd0 authored by Jin Hyuk Chang's avatar Jin Hyuk Chang Committed by GitHub

[DPTOOLS-1858] Add shutdown hook for Closer (#3)

* [DPTOOLS-1858] Add shutdown hook for deleting tmp folder in loader

* Update
parent 50cd5554
...@@ -26,10 +26,12 @@ class FsNeo4jCSVLoader(Loader): ...@@ -26,10 +26,12 @@ class FsNeo4jCSVLoader(Loader):
# Config keys # Config keys
NODE_DIR_PATH = 'node_dir_path' NODE_DIR_PATH = 'node_dir_path'
RELATION_DIR_PATH = 'relationship_dir_path' RELATION_DIR_PATH = 'relationship_dir_path'
FORCE_CREATE_DIR = 'force_create_directory'
SHOULD_DELETE_CREATED_DIR = 'delete_created_directories' SHOULD_DELETE_CREATED_DIR = 'delete_created_directories'
_DEFAULT_CONFIG = ConfigFactory.from_dict({ _DEFAULT_CONFIG = ConfigFactory.from_dict({
SHOULD_DELETE_CREATED_DIR: True SHOULD_DELETE_CREATED_DIR: True,
FORCE_CREATE_DIR: False
}) })
def __init__(self): def __init__(self):
...@@ -55,6 +57,7 @@ class FsNeo4jCSVLoader(Loader): ...@@ -55,6 +57,7 @@ class FsNeo4jCSVLoader(Loader):
self._delete_created_dir = \ self._delete_created_dir = \
conf.get_bool(FsNeo4jCSVLoader.SHOULD_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._node_dir)
self._create_directory(self._relation_dir) self._create_directory(self._relation_dir)
...@@ -67,8 +70,11 @@ class FsNeo4jCSVLoader(Loader): ...@@ -67,8 +70,11 @@ class FsNeo4jCSVLoader(Loader):
:return: :return:
""" """
if os.path.exists(path): if os.path.exists(path):
raise RuntimeError( if self._force_create_dir:
'Directory should not exist: {}'.format(path)) LOGGER.info('Directory exist. Deleting directory {}'.format(path))
shutil.rmtree(path)
else:
raise RuntimeError('Directory should not exist: {}'.format(path))
os.makedirs(path) os.makedirs(path)
......
import atexit
from typing import Callable, List # noqa: F401 from typing import Callable, List # noqa: F401
...@@ -13,6 +15,7 @@ class Closer(object): ...@@ -13,6 +15,7 @@ class Closer(object):
def __init__(self): def __init__(self):
# type: () -> None # type: () -> None
self._stack = [] # type: List self._stack = [] # type: List
atexit.register(self.close)
def register(self, close_callable): def register(self, close_callable):
# type: (Callable) -> None # type: (Callable) -> None
......
from setuptools import setup, find_packages from setuptools import setup, find_packages
__version__ = '1.0.0' __version__ = '1.0.1'
setup( setup(
name='amundsen-databuilder', name='amundsen-databuilder',
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment