diff options
-rw-r--r-- | tasks/import_reports_to_database.py | 50 |
1 files changed, 50 insertions, 0 deletions
diff --git a/tasks/import_reports_to_database.py b/tasks/import_reports_to_database.py new file mode 100644 index 0000000..152c762 --- /dev/null +++ b/tasks/import_reports_to_database.py | |||
@@ -0,0 +1,50 @@ | |||
1 | import sys | ||
2 | import os | ||
3 | import simplejson as json | ||
4 | from datetime import datetime | ||
5 | from decimal import Decimal as D | ||
6 | import psycopg2 | ||
7 | |||
8 | sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) | ||
9 | from main import parse_config | ||
10 | |||
11 | config = sys.argv[1] | ||
12 | reports = sys.argv[2:] | ||
13 | |||
14 | pg_config, report_path = parse_config(config) | ||
15 | |||
16 | connection = psycopg2.connect(**pg_config) | ||
17 | cursor = connection.cursor() | ||
18 | |||
19 | report_query = 'INSERT INTO reports("date", "market_config_id", "debug") VALUES (%s, %s, %s) RETURNING id;' | ||
20 | line_query = 'INSERT INTO report_lines("date", "report_id", "type", "payload") VALUES (%s, %s, %s, %s);' | ||
21 | |||
22 | user_id_to_market_id = { | ||
23 | 2: 1, | ||
24 | 1: 3, | ||
25 | } | ||
26 | |||
27 | for report in reports: | ||
28 | with open(report, "rb") as f: | ||
29 | json_content = json.load(f, parse_float=D) | ||
30 | basename = os.path.basename(report) | ||
31 | date, rest = basename.split("_", 1) | ||
32 | user_id, rest = rest.split(".", 1) | ||
33 | |||
34 | date = datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f") | ||
35 | market_id = user_id_to_market_id[int(user_id)] | ||
36 | debug = any("debug" in x and x["debug"] for x in json_content) | ||
37 | print(market_id, date, debug) | ||
38 | cursor.execute(report_query, (date, market_id, debug)) | ||
39 | report_id = cursor.fetchone()[0] | ||
40 | |||
41 | for line in json_content: | ||
42 | date = datetime.strptime(line["date"], "%Y-%m-%dT%H:%M:%S.%f") | ||
43 | type_ = line["type"] | ||
44 | del(line["date"]) | ||
45 | del(line["type"]) | ||
46 | |||
47 | cursor.execute(line_query, (date, report_id, type_, json.dumps(line))) | ||
48 | connection.commit() | ||
49 | cursor.close() | ||
50 | connection.close() | ||