source: TOOLS/ConsoGENCI/trunk/bin/init_conso_tbl.py @ 2775

Last change on this file since 2775 was 2775, checked in by labetoulle, 6 years ago

Overall cleaning and refactoring

File size: 6.1 KB
Line 
1#!/usr/bin/env python
2# -*- coding: utf-8 -*-
3
4# ==================================================================== #
5# Author: Sonia Labetoulle                                             #
6# Contact: sonia.labetoulle _at_ ipsl.jussieu.fr                       #
7# Created: 2016                                                        #
8# History:                                                             #
9# Modification:                                                        #
10# ==================================================================== #
11
12# This must come first
13from __future__ import print_function, unicode_literals, division
14
15# Standard library imports
16import os
17import math
18# import datetime as dt
19from argparse import ArgumentParser
20import pprint
21
22# Application library imports
23import libconso_db as cdb
24import libconso_cpt as ccpt
25import db_data
26
27pp = pprint.PrettyPrinter(indent=2)
28
29
30#######################################################################
31def get_arguments():
32  parser = ArgumentParser()
33  parser.add_argument("project", action="store",
34                      help="Project name")
35  parser.add_argument("center", action="store",
36                      help="Center name (idris/tgcc)")
37
38  parser.add_argument("-v", "--verbose", action="store_true",
39                      help="verbose mode")
40  parser.add_argument("-d", "--dryrun", action="store_true",
41                      help="only print what is to be done")
42  parser.add_argument("-r", "--range", action="store", nargs=2,
43                      help="date range: ssaammjj ssaammjj")
44
45  # parser.add_argument("-f", "--full", action="store_true",
46  #                     help="plot the whole period")
47  # parser.add_argument("-i", "--increment", action="store",
48  #                     type=int, default=1, dest="inc",
49  #                     help="sampling increment")
50  # parser.add_argument("-m", "--max", action="store_true",
51  #                     help="plot with y_max = allocation")
52  # parser.add_argument("-s", "--show", action="store_true",
53  #                     help="interactive mode")
54  # parser.add_argument("-d", "--dods", action="store_true",
55  #                     help="copy output on dods")
56
57  return parser.parse_args()
58
59
60#######################################################################
61if __name__ == "__main__":
62
63  # .. Initialization ..
64  # ====================
65  # ... Command line arguments ...
66  # ------------------------------
67  args = get_arguments()
68  if args.verbose:
69    print(args)
70
71  # ... Files and directories ...
72  # -----------------------------
73  ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
74  DATA_DIR = os.path.join(ROOT_DIR, "data")
75  SUBMIT_DIR = os.getcwd()
76
77  if args.verbose:
78    print("SUBMIT_DIR:", SUBMIT_DIR)
79    print("DATA_DIR:", DATA_DIR)
80
81  pattern = ccpt.cpt_pattern(args.center, args.project)
82  dirin = os.path.join(DATA_DIR, args.center, args.project)
83  filelist = ccpt.find_input_files(dirin, pattern, args.range)
84
85  # .. Connection to database ..
86  # ============================
87  if args.verbose:
88    print("Connection to database")
89  # db_host = "134.157.170.104"
90  # # db_port = "5432"
91  # db_name = "prodiguer"
92  # db_user = "prodiguer_db_user"
93  # # db_pwd = "secret"
94  conn, cursor = cdb.connect_db(
95    db_data.db_host,
96    db_data.db_name,
97    db_data.db_user
98  )
99
100  # .. Extract allocation id from table ..
101  # ======================================
102  table_name = "conso.tbl_allocation"
103  request = (
104    # "SELECT id "
105    "SELECT * "
106    "FROM " + table_name + " "
107    "WHERE project = '" + args.project + "'"
108    "  AND centre = '" + args.center + "'"
109    # "  AND start_date < '" + dt.datetime.strftime(date, "%Y-%m-%d") + "'"
110    # "  AND end_date > '" + dt.datetime.strftime(date, "%Y-%m-%d") + "'"
111    "ORDER BY start_date"
112    ";"
113  )
114
115  cdb.select_db(cursor, request)
116  # print(cursor.rowcount)
117  # print(cursor.fetchall())
118
119  allocs = []
120  for row in cursor:
121    allocs.append(ccpt.AllocRow(row))
122
123  # .. Build dictionary from files ..
124  # =================================
125  conso_per_alloc = {}
126  for filename in filelist:
127    date, blocs = ccpt.parse_input_cpt(
128      filename, args.project, args.center, mode_conso=True
129    )
130
131    for bloc in blocs:
132      if not bloc.alloc_id:
133        bloc.alloc_id = ccpt.get_project_id(bloc, allocs)
134        if not bloc.alloc_id:
135          print("no alloc id found, skip bloc")
136          print(bloc.machine, bloc.node, bloc.cpt_date)
137          pp.pprint(allocs)
138          continue
139      if bloc.alloc_id not in conso_per_alloc:
140        conso_per_alloc[bloc.alloc_id] = set()
141      for conso in bloc.consos:
142        conso_per_alloc[bloc.alloc_id].add(conso)
143
144  # .. Insert data in table, one alloc_id at a time ..
145  # ==================================================
146  for alloc_id, consos in conso_per_alloc.iteritems():
147    # ... Create request sub string ...
148    # ---------------------------------
149    lines_req = [
150      (
151        "('{alloc}', "
152        "'{date}', "
153        "{total_hrs}, "
154        "{login}, "
155        "{create})"
156      ) .format(
157        alloc=alloc_id,
158        date=item.date,
159        total_hrs=item.conso if not math.isnan(item.conso) else "'NaN'",
160        login="'"+item.login+"'" if item.login != "total" else "NULL",
161        create="CURRENT_TIMESTAMP",
162      ) for item in consos
163    ]
164
165    # print(lines_req)
166
167    # exit()
168
169    # ... Create full request ...
170    # ---------------------------
171    table_name = "conso.tbl_consumption"
172    request = (
173      "INSERT INTO " + table_name + " ("
174      "  allocation_id, "
175      "  date, "
176      "  total_hrs, "
177      "  login, "
178      "  row_create_date "
179      ") "
180      "VALUES "
181    )
182    request = request + ", ".join(lines_req)
183
184    # ... Execute request ...
185    # -----------------------
186    if args.verbose:
187      print("Execute request for alloc_id = {}".format(alloc_id))
188    cdb.insert_db(cursor, request)
189
190    # ... Commit inserts ...
191    # ----------------------
192    if not args.dryrun:
193      if args.verbose:
194        print("Commit inserts")
195      cdb.commit_db(conn)
196
197  # .. Close connection to database ..
198  # ==================================
199  if args.verbose:
200    print("Close connection")
201  cdb.close_db(conn)
Note: See TracBrowser for help on using the repository browser.