source: TOOLS/ConsoGENCI/bin/libconsodb.py @ 2713

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

# ==================================================================== #
# ssh readonly@prodiguer-test-db.ipsl.upmc.fr                          #
# psql -U prodiguer_db_user prodiguer                                  #
# ==================================================================== #


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

# Standard library imports
import os
import glob
import psycopg2
import psycopg2.extras

# Application library imports


########################################
def connect_db(db_host, db_name, db_user):

  conn_string = "host={} dbname={} user={}".format(
    db_host, db_name, db_user
  )

  # Print the connection string we will use to connect
  print("Connecting to database\n  -> {}".format(conn_string))

  # Get a connection, if a connect cannot be made an exception will be
  # raised here
  try :
    conn = psycopg2.connect(conn_string)
  except Exception as rc:
    print("Problem with DB connect:\n{}".format(rc))
    exit(1)

  # conn.cursor will return a cursor object, you can use this cursor
  # to perform queries
  # cursor = conn.cursor()
  cursor = conn.cursor(cursor_factory=psycopg2.extras.DictCursor)
  print("Connected!\n")

  return conn, cursor


########################################
def select_db(cursor, request):

  # Execute query
  try:
    cursor.execute(request)
  except Exception as rc:
    print("Problem with select:\n{}".format(rc))
    exit(1)

  # # retrieve the records from the database
  # records = cursor.fetchall()

  # # print out the records using pretty print
  # # note that the NAMES of the columns are not shown, instead just indexes.
  # # for most people this isn't very useful so we'll show you how to return
  # # columns as a dictionary (hash) in the next example.
  # pprint(records)


########################################
def insert_db(cursor, request):
  """
  """

  # Execute query
  try:
    cursor.execute(request)
  except Exception as rc:
    print("Problem with insert:\n{}".format(rc))
    exit(1)


########################################
def commit_db(conn):
  """
  """

  # Execute query
  try:
    conn.commit()
  except Exception as rc:
    print("Problem with commit:\n{}".format(rc))
    exit(1)


########################################
def parse_input_cpt(filename, project, center):

  if center == "idris":
    return parse_idris_cpt(filename, project)
  elif center == "tgcc":
    return parse_tgcc_cpt(filename, project)
  else:
    print("Unknown center {}".format(center))
    exit()


########################################
def parse_tgcc_cpt(filename, project):

  import datetime as dt

  consos = set()

  with open(filename, "r") as filein:
    # Skip lines until we find project name.
    # Then extract script date
    for ligne in filein:
      if project["name"] in ligne:
        date = ligne.split()[-1]
        date = dt.datetime.strptime(
          "{} {}".format(date), "%d-%m-%Y"
        )
        project["machine"] = ligne.split()[5]
        project["nodes"]   = ligne.split()[6]
        break

    # Skip next two lines : first is blank and second is login titles
    for _ in xrange(1):
      next(filein)

    # Login list, until blank line
    for ligne in filein:
      if not ligne.strip():
        break
      login, conso = ligne.split()
      # logins[login] = float(conso)
      try:
        conso = float(conso)
      except ValueError:
        conso = float("nan")
      consos.add((login, conso))

    # Skip until we find consumed time (hours)
    for ligne in filein:
      if "Total" in ligne:
        total = float(ligne.split()[-1])
        break

    # Skip until we find allocated time (hours)
    for ligne in filein:
      if "Allocated" in ligne:
        alloc = float(ligne.split()[-1])
        break

    # Skip until we find theoratical use (%)
    for ligne in filein:
      if "Suggested use at this time" in ligne:
        utheo = float(ligne.split()[-1].strip("%"))
        break

    # Skip until we find real use (%)
    for ligne in filein:
      if "Real use at this time" in ligne:
        ureal = float(ligne.split()[-1].strip("%"))
        break

    # Skip until we find deadline
    for ligne in filein:
      if "Project deadline" in ligne:
        deadline = ligne.split()[-1]
        break

  # return project, logins, today, total, utheo, ureal
  return date, alloc, consos


########################################
def parse_idris_cpt(filename, project):

  import datetime as dt

  # print(filename)

  with open(filename, "r") as filein:
    lignes = filein.readlines()

  consos = set()

  for ligne in lignes:
    if "mise a jour" in ligne:
      jour, heure = ligne.strip().split()[-2:]
    elif "Heures" in ligne:
      alloc = int(ligne.strip().split()[-1])
    elif "Totaux" in ligne:
      conso = float(ligne.strip().split()[1])
      consos.add(("total", conso))
    elif project in ligne:
      login = ligne.strip().split()[0]
      try:
        conso = float(ligne.strip().split()[2])
      except ValueError:
        conso = float("nan")
      consos.add((login, conso))

  date = dt.datetime.strptime(
    "{} {}".format(jour, heure), "%d/%m/%Y %H:%M"
  )

  # return jour, heure, alloc, consos
  return date, alloc, consos


########################################
def find_input_files(data_dir, pattern, date_range):
  """
  """
  # os.chdir(DATA_DIR)
  # pattern = "compta_*.dat"
  file_list = glob.glob(os.path.join(data_dir, pattern))
  # if file_list:
  #   res = sorted(file_list)
  # else:
  #   res = None
  # os.chdir(SUBMIT_DIR)

  if date_range:
    res = []
    for filename in file_list:
      file_date = os.path.basename(filename).split("_")[1]
      if file_date >= min(date_range) and \
         file_date <= max(date_range):
        res.append(filename)
  else:
    res = file_list

  return set(res)


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