tasks/import_reports_to_database.py

   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 market_config_query = "SELECT id FROM market_configs WHERE user_id = %s AND market_name = 'poloniex'"
  22
  23 for report in reports:
  24     with open(report, "rb") as f:
  25         json_content = json.load(f, parse_float=D)
  26         basename = os.path.basename(report)
  27         date, rest = basename.split("_", 1)
  28         user_id, rest = rest.split(".", 1)
  29
  30         date = datetime.strptime(date, "%Y-%m-%dT%H:%M:%S.%f")
  31         cursor.execute(market_config_query, user_id)
  32         market_id = cursor.fetchone()[0]
  33         debug = any("debug" in x and x["debug"] for x in json_content)
  34         print(market_id, date, debug)
  35         cursor.execute(report_query, (date, market_id, debug))
  36         report_id = cursor.fetchone()[0]
  37
  38         for line in json_content:
  39             date = datetime.strptime(line["date"], "%Y-%m-%dT%H:%M:%S.%f")
  40             type_ = line["type"]
  41             del(line["date"])
  42             del(line["type"])
  43
  44             cursor.execute(line_query, (date, report_id, type_, json.dumps(line, indent="  ")))
  45 connection.commit()
  46 cursor.close()
  47 connection.close()