Add task to migrate reports to database

1 files changed, 47 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..6031cbe
--- /dev/null
+++ b/tasks/import_reports_to_database.py
@@ -0,0 +1,47 @@
+import sys
+import os
+import simplejson as json
+from datetime import datetime
+from decimal import Decimal as D
+import psycopg2
+
+sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+from main import parse_config
+
+config = sys.argv[1]
+reports = sys.argv[2:]
+
+pg_config, report_path = parse_config(config)
+
+connection = psycopg2.connect(**pg_config)
+cursor = connection.cursor()
+
+report_query = 'INSERT INTO reports("date", "market_config_id", "debug") VALUES (%s, %s, %s) RETURNING id'
+line_query = 'INSERT INTO report_lines("date", "report_id", "type", "payload") VALUES (%s, %s, %s, %s)'
+market_config_query = "SELECT id FROM market_configs WHERE user_id = %s AND market_name = 'poloniex'"
+
+for report in reports:
+    with open(report, "rb") as f:
+        json_content = json.load(f, parse_float=D)
+        basename = os.path.basename(report)
+        date, rest = basename.split("_", 1)
+        user_id, rest = rest.split(".", 1)
+
+        date = datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f")
+        cursor.execute(market_config_query, user_id)
+        market_id = cursor.fetchone()[0]
+        debug = any("debug" in x and x["debug"] for x in json_content)
+        print(market_id, date, debug)
+        cursor.execute(report_query, (date, market_id, debug))
+        report_id = cursor.fetchone()[0]
+
+        for line in json_content:
+            date = datetime.strptime(line["date"], "%Y-%m-%dT%H:%M:%S.%f")
+            type_ = line["type"]
+            del(line["date"])
+            del(line["type"])
+
+            cursor.execute(line_query, (date, report_id, type_, json.dumps(line, indent="  ")))
+connection.commit()
+cursor.close()
+connection.close()