tools/upgrade-db.py

   1 #!/usr/bin/python
   2 #
   3 # Tool for upgrading a db based on db version #
   4 import sys
   5 import os
   6 import getopt
   7 import pgdb
   8
   9 config = {}
  10 config_file = "/etc/planetlab/plc_config"
  11 execfile(config_file, config)
  12 upgrade_config_file = "plcdb.3-4.conf"
  13 schema_file = "planetlab4.sql"
  14 temp_dir = "/tmp"
  15
  16
  17 def usage():
  18         print "Usage: %s [OPTION] UPGRADE_CONFIG_FILE " % sys.argv[0]
  19         print "Options:"
  20         print "     -s, --schema=FILE       Upgraded Database Schema"
  21         print "     -t, --temp-dir=DIR      Temp Directory"
  22         print "     --help                  This message"
  23         sys.exit(1)
  24
  25 try:
  26         (opts, argv) = getopt.getopt(sys.argv[1:],
  27                                      "s:d:",
  28                                      ["schema=",
  29                                       "temp-dir=",
  30                                       "help"])
  31 except getopt.GetoptError, err:
  32         print "Error: ", err.msg
  33         usage()
  34
  35 for (opt, optval) in opts:
  36         if opt == "-s" or opt == "--schema":
  37                 schema_file = optval
  38         elif opt == "-d" or opt == "--temp-dir":
  39                 temp_dir = optval
  40         elif opt == "--help":
  41                 usage()
  42 try:
  43         upgrade_config_file = argv[0]
  44 except IndexError:
  45         print "Error: too few arguments"
  46         usage()
  47
  48 schema = {}
  49 inserts = []
  50 schema_items_ordered = []
  51 temp_tables = {}
  52
  53
  54 # load conf file for this upgrade
  55 try:
  56         upgrade_config = {}
  57         execfile(upgrade_config_file, upgrade_config)
  58         upgrade_config.pop('__builtins__')
  59         db_version_previous = upgrade_config['DB_VERSION_PREVIOUS']
  60         db_version_new = upgrade_config['DB_VERSION_NEW']
  61
  62 except IOError, fault:
  63         print "Error: upgrade config file (%s) not found. Exiting" % \
  64                 (fault)
  65         sys.exit(1)
  66 except KeyError, fault:
  67         print "Error: %s not set in upgrade confing (%s). Exiting" % \
  68                 (fault, upgrade_config_file)
  69         sys.exit(1)
  70
  71
  72
  73
  74 def connect():
  75         db = pgdb.connect(user = config['PLC_DB_USER'],
  76                   database = config['PLC_DB_NAME'])
  77         return db
  78
  79 def archive_db(database, archived_database):
  80
  81         archive_db = " dropdb -U postgres %s > /dev/null 2>&1;" \
  82                      " psql template1 postgres -qc " \
  83                      " 'ALTER DATABASE %s RENAME TO %s;';" \
  84                      " createdb -U postgres %s > /dev/null; " % \
  85                      (archived_database, database, archived_database, database)
  86         exit_status = os.system(archive_db)
  87         if exit_status:
  88                 print "Error: unable to archive database. Upgrade failed"
  89                 sys.exit(1)
  90         #print "Status: %s has been archived. now named %s" % (database, archived_database)
  91
  92
  93 def encode_utf8(inputfile_name, outputfile_name):
  94         # rewrite a iso-8859-1 encoded file in utf8
  95         try:
  96                 inputfile = open(inputfile_name, 'r')
  97                 outputfile = open(outputfile_name, 'w')
  98                 for line in inputfile:
  99                         outputfile.write(unicode(line, 'iso-8859-1').encode('utf8'))
 100                 inputfile.close()
 101                 outputfile.close()
 102         except:
 103                 print 'error encoding file'
 104                 raise
 105
 106 def create_item_from_schema(item_name):
 107
 108         try:
 109                 (type, body_list) = schema[item_name]
 110                 exit_status = os.system('psql %s %s -qc "%s" > /dev/null 2>&1' % \
 111                             (config['PLC_DB_NAME'], config['PLC_DB_USER'],"".join(body_list) ) )
 112                 if exit_status:
 113                         raise Exception
 114         except Exception, fault:
 115                 print 'Error: create %s failed. Check schema.' % item_name
 116                 sys.exit(1)
 117                 raise fault
 118
 119         except KeyError:
 120                 print "Error: cannot create %s. definition not found in %s" % \
 121                         (key, schema_file)
 122                 return False
 123
 124 def fix_row(row, table_name, table_fields):
 125
 126         if table_name in ['nodenetworks']:
 127                 # convert str bwlimit to bps int
 128                 bwlimit_index = table_fields.index('bwlimit')
 129                 if isinstance(row[bwlimit_index], int):
 130                         pass
 131                 elif row[bwlimit_index].find('mbit') > -1:
 132                         row[bwlimit_index] = int(row[bwlimit_index].split('mbit')[0]) \
 133                                             * 1000000
 134                 elif row[bwlimit_index].find('kbit') > -1:
 135                         row[bwlimit_index] = int(row[bwlimit_index].split('kbit')[0]) \
 136                                              * 1000
 137         elif table_name in ['slice_attribute']:
 138                 # modify some invalid foreign keys
 139                 attribute_type_index = table_fields.index('attribute_type_id')
 140                 if row[attribute_type_index] == 10004:
 141                         row[attribute_type_index] = 10016
 142                 elif row[attribute_type_index] == 10006:
 143                         row[attribute_type_index] = 10017
 144         elif table_name in ['slice_attribute_types']:
 145                 type_id_index = table_fields.index('attribute_type_id')
 146                 if row[type_id_index] in [10004, 10006]:
 147                         return None
 148         return row
 149
 150 def fix_table(table, table_name, table_fields):
 151         if table_name in ['slice_attribute_types']:
 152                 # remove duplicate/redundant primary keys
 153                 type_id_index = table_fields.index('attribute_type_id')
 154                 for row in table:
 155                         if row[type_id_index] in [10004, 10006]:
 156                                 table.remove(row)
 157         return table
 158
 159 def remove_temp_tables():
 160         # remove temp_tables
 161         try:
 162                 for temp_table in temp_tables:
 163                         os.remove(temp_tables[temp_table])
 164         except:
 165                 raise
 166
 167 def generate_temp_table(table_name, db):
 168         cursor = db.cursor()
 169         try:
 170                 # get upgrade directions
 171                 table_def = upgrade_config[table_name].replace('(', '').replace(')', '').split(',')
 172                 table_fields, old_fields, required_joins = [], [], set()
 173                 for field in table_def:
 174                         field_parts = field.strip().split(':')
 175                         table_fields.append(field_parts[0])
 176                         old_fields.append(field_parts[1])
 177                         if field_parts[2:]:
 178                                 required_joins.update(set(field_parts[2:]))
 179
 180                 # get indices of fields that cannot be null
 181                 (type, body_list) = schema[table_name]
 182                 not_null_indices = []
 183                 for field in table_fields:
 184                         for body_line in body_list:
 185                                 if body_line.find(field) > -1 and \
 186                                    body_line.upper().find("NOT NULL") > -1:
 187                                         not_null_indices.append(table_fields.index(field))
 188
 189                 # get index of primary key
 190                 primary_key_indices = []
 191                 for body_line in body_list:
 192                         if body_line.find("PRIMARY KEY") > -1:
 193                                 primary_key = body_line
 194                                 for field in table_fields:
 195                                         if primary_key.find(field) > -1:
 196                                                 primary_key_indices.append(table_fields.index(field))
 197                                 break
 198
 199                 # get old data
 200                 get_old_data = "SELECT DISTINCT %s FROM %s" % \
 201                       (", ".join(old_fields), old_fields[0].split(".")[0])
 202                 for join in required_joins:
 203                         get_old_data = get_old_data + " INNER JOIN %s USING (%s) " % \
 204                                        (join.split('.')[0], join.split('.')[1])
 205
 206                 cursor.execute(get_old_data)
 207                 rows = cursor.fetchall()
 208
 209                 # write data to a temp file
 210                 temp_file_name = '%s/%s.tmp' % (temp_dir, table_name)
 211                 temp_file = open(temp_file_name, 'w')
 212                 for row in rows:
 213                         # attempt to make any necessary fixes to data
 214                         row = fix_row(row, table_name, table_fields)
 215                         # do not attempt to write null rows
 216                         if row == None:
 217                                 continue
 218                         # do not attempt to write rows with null primary keys
 219                         if filter(lambda x: row[x] == None, primary_key_indices):
 220                                 continue
 221                         for i in range(len(row)):
 222                                 # convert nulls into something pg can understand
 223                                 if row[i] == None:
 224                                         if i in not_null_indices:
 225                                                 # XX doesnt work if column is int type
 226                                                 row[i] = ""
 227                                         else:
 228                                                 row[i] = "\N"
 229                                 if isinstance(row[i], int) or isinstance(row[i], float):
 230                                         row[i] = str(row[i])
 231                                 # escape whatever can mess up the data format
 232                                 if isinstance(row[i], str):
 233                                         row[i] = row[i].replace('\t', '\\t')
 234                                         row[i] = row[i].replace('\n', '\\n')
 235                                         row[i] = row[i].replace('\r', '\\r')
 236                         data_row = "\t".join(row)
 237                         temp_file.write(data_row + "\n")
 238                 temp_file.write("\.\n")
 239                 temp_file.close()
 240                 temp_tables[table_name] = temp_file_name
 241
 242         except KeyError:
 243                 #print "WARNING: cannot upgrade %s. upgrade def not found. skipping" % \
 244                 #       (table_name)
 245                 return False
 246         except IndexError, fault:
 247                 print "Error: error found in upgrade config file. " \
 248                       "check %s configuration. Aborting " % \
 249                       (table_name)
 250                 sys.exit(1)
 251         except:
 252                 print "Error: configuration for %s doesnt match db schema. " \
 253                       " Aborting" % (table_name)
 254                 try:
 255                         db.rollback()
 256                 except:
 257                         pass
 258                 raise
 259
 260
 261 # Connect to current db
 262 db = connect()
 263 cursor = db.cursor()
 264
 265 # determin current db version
 266 try:
 267         cursor.execute("SELECT relname from pg_class where relname = 'plc_db_version'")
 268         rows = cursor.fetchall()
 269         if not rows:
 270                 print "Warning: current db has no version. Unable to validate config file."
 271         else:
 272                 cursor.execute("SELECT version FROM plc_db_version")
 273                 rows = cursor.fetchall()
 274                 if not rows or not rows[0]:
 275                         print "Warning: current db has no version. Unable to validate config file."
 276                 elif rows[0][0] == db_version_new:
 277                         print "Status: Versions are the same. No upgrade necessary."
 278                         sys.exit()
 279                 elif not rows[0][0] == db_version_previous:
 280                         print "Stauts: DB_VERSION_PREVIOUS in config file (%s) does not" \
 281                               " match current db version %d" % (upgrade_config_file, rows[0][0])
 282                         sys.exit()
 283                 else:
 284                         print "STATUS: attempting upgrade from %d to %d" % \
 285                                 (db_version_previous, db_version_new)
 286
 287         # check db encoding
 288         sql = " SELECT pg_catalog.pg_encoding_to_char(d.encoding)" \
 289               " FROM pg_catalog.pg_database d " \
 290               " LEFT JOIN pg_catalog.pg_user u ON d.datdba = u.usesysid " \
 291               " WHERE d.datname = '%s' " % config['PLC_DB_NAME']
 292         cursor.execute(sql)
 293         rows = cursor.fetchall()
 294         if rows[0][0] not in ['UTF8']:
 295                 print "WARNING: db encoding is not utf8. Attempting to encode"
 296                 db.close()
 297                 # generate db dump
 298                 dump_file = '%s/dump.sql' % (temp_dir)
 299                 dump_file_encoded = dump_file + ".utf8"
 300                 dump_cmd = 'pg_dump -i %s -U %s -f %s > /dev/null 2>&1' % \
 301                            (config['PLC_DB_NAME'], config['PLC_DB_USER'], dump_file)
 302                 if os.system(dump_cmd):
 303                         print "ERROR: during db dump. Exiting."
 304                         sys.exit(1)
 305                 # encode dump to utf8
 306                 print "Status: encoding database dump"
 307                 encode_utf8(dump_file, dump_file_encoded)
 308                 # archive original db
 309                 archive_db(config['PLC_DB_NAME'], config['PLC_DB_NAME']+'_sqlascii_archived')
 310                 # create a utf8 database and upload encoded data
 311                 recreate_cmd = 'createdb -U postgres -E UTF8 %s > /dev/null 2>&1; ' \
 312                                'psql -a -U  %s %s < %s > /dev/null 2>&1;'   % \
 313                           (config['PLC_DB_NAME'], config['PLC_DB_USER'], \
 314                            config['PLC_DB_NAME'], dump_file_encoded)
 315                 print "Status: recreating database as utf8"
 316                 if os.system(recreate_cmd):
 317                         print "Error: database encoding failed. Aborting"
 318                         sys.exit(1)
 319
 320                 os.remove(dump_file_encoded)
 321                 os.remove(dump_file)
 322 except:
 323         raise
 324
 325
 326 db = connect()
 327 cursor = db.cursor()
 328
 329 # parse the schema user wishes to upgrade to
 330 try:
 331         file = open(schema_file, 'r')
 332         index = 0
 333         lines = file.readlines()
 334         while index < len(lines):
 335                 line = lines[index]
 336                 # find all created objects
 337                 if line.startswith("CREATE"):
 338                         line_parts = line.split(" ")
 339                         item_type = line_parts[1]
 340                         item_name = line_parts[2]
 341                         schema_items_ordered.append(item_name)
 342                         if item_type in ['INDEX']:
 343                                 schema[item_name] = (item_type, line)
 344
 345                         # functions, tables, views span over multiple lines
 346                         # handle differently than indexes
 347                         elif item_type in ['AGGREGATE', 'TABLE', 'VIEW']:
 348                                 fields = [line]
 349                                 while index < len(lines):
 350                                         index = index + 1
 351                                         nextline =lines[index]
 352                                         fields.append(nextline)
 353                                         if nextline.find(";") >= 0:
 354                                                 break
 355                                 schema[item_name] = (item_type, fields)
 356                         else:
 357                                 print "Error: unknown type %s" % item_type
 358                 elif line.startswith("INSERT"):
 359                         inserts.append(line)
 360                 index = index + 1
 361
 362 except:
 363         raise
 364
 365 print "Status: generating temp tables"
 366 # generate all temp tables
 367 for key in schema_items_ordered:
 368         (type, body_list) = schema[key]
 369         if type == 'TABLE':
 370                 generate_temp_table(key, db)
 371
 372 # disconenct from current database and archive it
 373 cursor.close()
 374 db.close()
 375
 376 print "Status: archiving database"
 377 archive_db(config['PLC_DB_NAME'], config['PLC_DB_NAME']+'_archived')
 378
 379
 380 print "Status: upgrading database"
 381 # attempt to create and load all items from schema into temp db
 382 try:
 383         for key in schema_items_ordered:
 384                 (type, body_list) = schema[key]
 385                 create_item_from_schema(key)
 386                 if type == 'TABLE':
 387                         if upgrade_config.has_key(key):
 388                                 table_def = upgrade_config[key].replace('(', '').replace(')', '').split(',')
 389                                 table_fields = [field.strip().split(':')[0] for field in table_def]
 390                                 insert_cmd = "psql %s %s -c " \
 391                                              " 'COPY %s (%s) FROM stdin;' < %s " % \
 392                                              (config['PLC_DB_NAME'], config['PLC_DB_USER'], key, ", ".join(table_fields), temp_tables[key] )
 393                                 exit_status = os.system(insert_cmd)
 394                                 if exit_status:
 395                                         print "Error: upgrade %s failed" % key
 396                                         raise
 397                         else:
 398                                 # check if there are any insert stmts in schema for this table
 399                                 print "Warning: %s has no temp data file. Unable to populate with old data" % key
 400                                 for insert_stmt in inserts:
 401                                         if insert_stmt.find(key) > -1:
 402                                                 insert_cmd = 'psql %s postgres -qc "%s;" > /dev/null 2>&1' % \
 403                                                 (config['PLC_DB_NAME'], insert_stmt)
 404                                                 os.system(insert_cmd)
 405 except:
 406         print "Error: failed to populate db. Unarchiving original database and aborting"
 407         undo_command = "dropdb -U postgres %s > /dev/null; psql template1 postgres -qc" \
 408                        " 'ALTER DATABASE %s RENAME TO %s;';  > /dev/null" % \
 409                        (config['PLC_DB_NAME'], config['PLC_DB_NAME']+'_archived', config['PLC_DB_NAME'])
 410         os.system(undo_command)
 411         remove_temp_tables()
 412         raise
 413
 414 remove_temp_tables()
 415
 416 print "upgrade complete"