[2775] | 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 |
---|
| 13 | from __future__ import print_function, unicode_literals, division |
---|
| 14 | |
---|
| 15 | # Standard library imports |
---|
| 16 | import os |
---|
| 17 | import math |
---|
| 18 | from argparse import ArgumentParser |
---|
| 19 | import shutil |
---|
| 20 | import pprint |
---|
| 21 | |
---|
| 22 | # Application library imports |
---|
| 23 | import libconso_db as cdb |
---|
| 24 | import libconso_cpt as ccpt |
---|
| 25 | import db_data |
---|
| 26 | |
---|
| 27 | pp = pprint.PrettyPrinter(indent=2) |
---|
| 28 | |
---|
| 29 | |
---|
| 30 | ####################################################################### |
---|
| 31 | def 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 | return parser.parse_args() |
---|
| 45 | |
---|
| 46 | |
---|
| 47 | ####################################################################### |
---|
| 48 | if __name__ == "__main__": |
---|
| 49 | |
---|
| 50 | # .. Initialization .. |
---|
| 51 | # ==================== |
---|
| 52 | # ... Command line arguments ... |
---|
| 53 | # ------------------------------ |
---|
| 54 | args = get_arguments() |
---|
| 55 | if args.verbose: |
---|
| 56 | print(args) |
---|
| 57 | |
---|
| 58 | # ... Files and directories ... |
---|
| 59 | # ----------------------------- |
---|
| 60 | ROOT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
---|
| 61 | DATA_DIR = os.path.join(ROOT_DIR, "data") |
---|
| 62 | SUBMIT_DIR = os.getcwd() |
---|
| 63 | |
---|
| 64 | pattern = ccpt.cpt_pattern(center=args.center, project=args.project) |
---|
| 65 | dirin = os.path.join(DATA_DIR, args.center, "tmp") |
---|
| 66 | dirout = os.path.join(DATA_DIR, args.center, args.project) |
---|
| 67 | filelist = ccpt.find_input_files(dirin, pattern, args.range) |
---|
| 68 | |
---|
| 69 | if not filelist: |
---|
| 70 | print("No files found, exit") |
---|
[2783] | 71 | exit(1) |
---|
[2775] | 72 | |
---|
| 73 | if args.verbose: |
---|
| 74 | print("SUBMIT_DIR:", SUBMIT_DIR) |
---|
| 75 | print("DATA_DIR:", DATA_DIR) |
---|
| 76 | print("dirin:", dirin) |
---|
| 77 | |
---|
| 78 | # .. Connection to database .. |
---|
| 79 | # ============================ |
---|
| 80 | if args.verbose: |
---|
| 81 | print("Connection to database") |
---|
| 82 | conn, cursor = cdb.connect_db( |
---|
| 83 | db_data.db_host, |
---|
| 84 | db_data.db_name, |
---|
[2783] | 85 | db_data.db_user, |
---|
| 86 | db_data.db_pwd, |
---|
[2775] | 87 | ) |
---|
| 88 | |
---|
| 89 | # .. Extract allocation id from table .. |
---|
| 90 | # ====================================== |
---|
| 91 | table_name = "conso.tbl_allocation" |
---|
| 92 | request = ( |
---|
| 93 | "SELECT * " |
---|
| 94 | "FROM " + table_name + " " |
---|
| 95 | "WHERE project = '" + args.project + "'" |
---|
| 96 | " AND centre = '" + args.center + "'" |
---|
| 97 | "ORDER BY start_date" |
---|
| 98 | ";" |
---|
| 99 | ) |
---|
| 100 | |
---|
| 101 | cdb.select_db(cursor, request) |
---|
| 102 | # print(cursor.rowcount) |
---|
| 103 | # print(cursor.fetchall()) |
---|
| 104 | |
---|
| 105 | allocs = [] |
---|
| 106 | for row in cursor: |
---|
| 107 | allocs.append(ccpt.AllocRow(row)) |
---|
| 108 | |
---|
[2776] | 109 | # .. Process files .. |
---|
| 110 | # =================== |
---|
[2775] | 111 | for filename in filelist: |
---|
[2776] | 112 | |
---|
| 113 | fileout = os.path.join(dirout, os.path.basename(filename)) |
---|
| 114 | if os.path.isfile(fileout): |
---|
| 115 | print( |
---|
| 116 | "file {} already processed, remove file".format( |
---|
| 117 | os.path.basename(filename) |
---|
| 118 | ) |
---|
| 119 | ) |
---|
| 120 | if not args.dryrun: |
---|
| 121 | try: |
---|
| 122 | os.remove(filename) |
---|
| 123 | except Exception as rc: |
---|
| 124 | print("Could not move {}:\n{}".format(filename, rc)) |
---|
| 125 | break |
---|
| 126 | |
---|
| 127 | # .. Build dictionary from files .. |
---|
| 128 | # ================================= |
---|
[2775] | 129 | conso_per_alloc = {} |
---|
| 130 | |
---|
| 131 | date, blocs = ccpt.parse_input_cpt( |
---|
| 132 | filename, args.project, args.center, mode_conso=True |
---|
| 133 | ) |
---|
| 134 | |
---|
| 135 | for bloc in blocs: |
---|
| 136 | if not bloc.alloc_id: |
---|
| 137 | bloc.alloc_id = ccpt.get_project_id(bloc, allocs) |
---|
| 138 | if not bloc.alloc_id: |
---|
| 139 | print("no alloc id found, skip bloc") |
---|
| 140 | print(bloc.machine, bloc.node, bloc.cpt_date) |
---|
| 141 | pp.pprint(allocs) |
---|
| 142 | continue |
---|
| 143 | if bloc.alloc_id not in conso_per_alloc: |
---|
| 144 | conso_per_alloc[bloc.alloc_id] = set() |
---|
| 145 | for conso in bloc.consos: |
---|
| 146 | conso_per_alloc[bloc.alloc_id].add(conso) |
---|
| 147 | |
---|
| 148 | # .. Insert data in table, one alloc_id at a time .. |
---|
| 149 | # ================================================== |
---|
| 150 | for alloc_id, consos in conso_per_alloc.iteritems(): |
---|
| 151 | # ... Create request sub string ... |
---|
| 152 | # --------------------------------- |
---|
| 153 | lines_req = [ |
---|
| 154 | ( |
---|
| 155 | "('{alloc}', " |
---|
| 156 | "'{date}', " |
---|
| 157 | "{total_hrs}, " |
---|
| 158 | "{login}, " |
---|
| 159 | "{create})" |
---|
| 160 | ) .format( |
---|
| 161 | alloc=alloc_id, |
---|
| 162 | date=item.date, |
---|
| 163 | total_hrs=item.conso if not math.isnan(item.conso) else "'NaN'", |
---|
| 164 | login="'"+item.login+"'" if item.login != "total" else "NULL", |
---|
| 165 | create="CURRENT_TIMESTAMP", |
---|
| 166 | ) for item in consos |
---|
| 167 | ] |
---|
| 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 | print(request) |
---|
| 189 | cdb.insert_db(cursor, request) |
---|
| 190 | |
---|
| 191 | # ... Commit inserts ... |
---|
| 192 | # ---------------------- |
---|
| 193 | if not args.dryrun: |
---|
| 194 | if args.verbose: |
---|
| 195 | print("Commit inserts") |
---|
| 196 | cdb.commit_db(conn) |
---|
| 197 | |
---|
| 198 | # .. Move processed file to project directory .. |
---|
| 199 | # =============================================== |
---|
| 200 | if args.verbose: |
---|
| 201 | print("Move processed file:") |
---|
| 202 | print("{} => {}".format(filename, dirout)) |
---|
| 203 | if not args.dryrun: |
---|
| 204 | try: |
---|
| 205 | shutil.move(filename, dirout) |
---|
| 206 | except Exception as rc: |
---|
| 207 | print("Could not move {}:\n{}".format(filename, rc)) |
---|
| 208 | |
---|
| 209 | # .. Close connection to database .. |
---|
| 210 | # ================================== |
---|
| 211 | if args.verbose: |
---|
| 212 | print("Close connection") |
---|
| 213 | cdb.close_db(conn) |
---|
[2783] | 214 | |
---|
| 215 | exit(0) |
---|