1 files changed, 6 insertions, 9 deletions

diff --git a/tasks/import_reports_to_database.py b/tasks/import_reports_to_database.py
index 152c762..6031cbe 100644
--- a/tasks/import_reports_to_database.py
+++ b/tasks/import_reports_to_database.py
@@ -16,13 +16,9 @@ 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);'
-
-user_id_to_market_id = {
-        2: 1,
-        1: 3,
-        }
+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:
@@ -32,7 +28,8 @@ for report in reports:
         user_id, rest = rest.split(".", 1)
 
         date = datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f")
-        market_id = user_id_to_market_id[int(user_id)]
+        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))
@@ -44,7 +41,7 @@ for report in reports:
             del(line["date"])
             del(line["type"])
 
-            cursor.execute(line_query, (date, report_id, type_, json.dumps(line)))
+            cursor.execute(line_query, (date, report_id, type_, json.dumps(line, indent="  ")))
 connection.commit()
 cursor.close()
 connection.close()