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

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

Overall cleaning and refactoring

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