-
Notifications
You must be signed in to change notification settings - Fork 3
/
run.py
62 lines (46 loc) · 1.67 KB
/
run.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
import os
import logging
import py2neo
import json
logging.basicConfig(level=logging.DEBUG)
logging.getLogger('py2neo.client.bolt').setLevel(logging.WARNING)
logging.getLogger('py2neo.client').setLevel(logging.WARNING)
logging.getLogger('urllib3.connectionpool').setLevel(logging.WARNING)
logging.getLogger('graphio').setLevel(logging.WARNING)
log = logging.getLogger(__name__)
# import and setup
from covid_graph import helper, post, jhu, unwpp
ROOT_DIR = os.getenv('ROOT_DIR', '/download')
RUN_MODE = os.getenv('RUN_MODE', 'prod')
NEO4J_CONFIG_STRING = os.getenv("NEO4J")
log.debug(f'NEO4J config string from env (not changed): {NEO4J_CONFIG_STRING}')
log.debug(f'Type of config string: {type(NEO4J_CONFIG_STRING)}')
try:
NEO4J_CONFIG_DICT = json.loads(NEO4J_CONFIG_STRING)
except json.decoder.JSONDecodeError:
# try to replace single quotes with double quotes
# JSON always expects double quotes, common mistake when writing JSON strings
NEO4J_CONFIG_STRING = NEO4J_CONFIG_STRING.replace("'", '"')
NEO4J_CONFIG_DICT = json.loads(NEO4J_CONFIG_STRING)
log.debug("NEO4J_CONFIG_DICT: ")
log.debug(NEO4J_CONFIG_DICT)
if RUN_MODE.lower() == 'test':
import pytest
log.info("Run tests")
pytest.main()
else:
graph = py2neo.Graph(**NEO4J_CONFIG_DICT)
# setup DB
helper.setup_db(graph)
# download data
jhu_zip_file = jhu.download_jhu(ROOT_DIR)
jhu_dir = helper.unzip_file(jhu_zip_file)
wpp_csv_file = unwpp.download_population_data(ROOT_DIR, skip_existing=True)
###
# load to Neo4j
jhu.read_daily_report_JHU(jhu_dir, graph)
unwpp.load_wpp_data(ROOT_DIR, graph)
###
# post process
post.set_latest_update(graph)
###