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
  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()