source: TOOLS/ConsoGENCI/trunk/bin/insert_conso_tbl.py @ 2783

#!/usr/bin/env python
# -*- coding: utf-8 -*-

# ==================================================================== #
# Author: Sonia Labetoulle                                             #
# Contact: sonia.labetoulle _at_ ipsl.jussieu.fr                       #
# Created: 2016                                                        #
# History:                                                             #
# Modification:                                                        #
# ==================================================================== #

# This must come first
from __future__ import print_function, unicode_literals, division

# Standard library imports
import os
import math
from argparse import ArgumentParser
import shutil
import pprint

# Application library imports
import libconso_db as cdb
import libconso_cpt as ccpt
import db_data

pp = pprint.PrettyPrinter(indent=2)


#######################################################################
def get_arguments():
  parser = ArgumentParser()
  parser.add_argument("project", action="store",
                      help="Project name")
  parser.add_argument("center", action="store",
                      help="Center name (idris/tgcc)")

  parser.add_argument("-v", "--verbose", action="store_true",
                      help="verbose mode")
  parser.add_argument("-d", "--dryrun", action="store_true",
                      help="only print what is to be done")
  parser.add_argument("-r", "--range", action="store", nargs=2,
                      help="date range: ssaammjj ssaammjj")
  return parser.parse_args()


#######################################################################
if __name__ == "__main__":

  # .. Initialization ..
  # ====================
  # ... Command line arguments ...
  # ------------------------------
  args = get_arguments()
  if args.verbose:
    print(args)

  # ... Files and directories ...
  # -----------------------------
  ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
  DATA_DIR = os.path.join(ROOT_DIR, "data")
  SUBMIT_DIR = os.getcwd()

  pattern = ccpt.cpt_pattern(center=args.center, project=args.project)
  dirin = os.path.join(DATA_DIR, args.center, "tmp")
  dirout = os.path.join(DATA_DIR, args.center, args.project)
  filelist = ccpt.find_input_files(dirin, pattern, args.range)

  if not filelist:
    print("No files found, exit")
    exit(1)

  if args.verbose:
    print("SUBMIT_DIR:", SUBMIT_DIR)
    print("DATA_DIR:", DATA_DIR)
    print("dirin:", dirin)

  # .. Connection to database ..
  # ============================
  if args.verbose:
    print("Connection to database")
  conn, cursor = cdb.connect_db(
    db_data.db_host,
    db_data.db_name,
    db_data.db_user,
    db_data.db_pwd,
  )

  # .. Extract allocation id from table ..
  # ======================================
  table_name = "conso.tbl_allocation"
  request = (
    "SELECT * "
    "FROM " + table_name + " "
    "WHERE project = '" + args.project + "'"
    "  AND centre = '" + args.center + "'"
    "ORDER BY start_date"
    ";"
  )

  cdb.select_db(cursor, request)
  # print(cursor.rowcount)
  # print(cursor.fetchall())

  allocs = []
  for row in cursor:
    allocs.append(ccpt.AllocRow(row))

  # .. Process files ..
  # ===================
  for filename in filelist:

    fileout = os.path.join(dirout, os.path.basename(filename))
    if os.path.isfile(fileout):
      print(
        "file {} already processed, remove file".format(
          os.path.basename(filename)
        )
      )
      if not args.dryrun:
        try:
          os.remove(filename)
        except Exception as rc:
          print("Could not move {}:\n{}".format(filename, rc))
      break

    # .. Build dictionary from files ..
    # =================================
    conso_per_alloc = {}

    date, blocs = ccpt.parse_input_cpt(
      filename, args.project, args.center, mode_conso=True
    )

    for bloc in blocs:
      if not bloc.alloc_id:
        bloc.alloc_id = ccpt.get_project_id(bloc, allocs)
        if not bloc.alloc_id:
          print("no alloc id found, skip bloc")
          print(bloc.machine, bloc.node, bloc.cpt_date)
          pp.pprint(allocs)
          continue
      if bloc.alloc_id not in conso_per_alloc:
        conso_per_alloc[bloc.alloc_id] = set()
      for conso in bloc.consos:
        conso_per_alloc[bloc.alloc_id].add(conso)

    # .. Insert data in table, one alloc_id at a time ..
    # ==================================================
    for alloc_id, consos in conso_per_alloc.iteritems():
      # ... Create request sub string ...
      # ---------------------------------
      lines_req = [
        (
          "('{alloc}', "
          "'{date}', "
          "{total_hrs}, "
          "{login}, "
          "{create})"
        ) .format(
          alloc=alloc_id,
          date=item.date,
          total_hrs=item.conso if not math.isnan(item.conso) else "'NaN'",
          login="'"+item.login+"'" if item.login != "total" else "NULL",
          create="CURRENT_TIMESTAMP",
        ) for item in consos
      ]

      # ... Create full request ...
      # ---------------------------
      table_name = "conso.tbl_consumption"
      request = (
        "INSERT INTO " + table_name + " ("
        "  allocation_id, "
        "  date, "
        "  total_hrs, "
        "  login, "
        "  row_create_date "
        ") "
        "VALUES "
      )
      request = request + ", ".join(lines_req)

      # ... Execute request ...
      # -----------------------
      if args.verbose:
        print("Execute request for alloc_id = {}".format(alloc_id))
        print(request)
      cdb.insert_db(cursor, request)

      # ... Commit inserts ...
      # ----------------------
      if not args.dryrun:
        if args.verbose:
          print("Commit inserts")
        cdb.commit_db(conn)

    # .. Move processed file to project directory ..
    # ===============================================
    if args.verbose:
      print("Move processed file:")
      print("{} => {}".format(filename, dirout))
    if not args.dryrun:
      try:
        shutil.move(filename, dirout)
      except Exception as rc:
        print("Could not move {}:\n{}".format(filename, rc))

  # .. Close connection to database ..
  # ==================================
  if args.verbose:
    print("Close connection")
  cdb.close_db(conn)

  exit(0)