# Copyright (c) 2016-2017 VMware, Inc.  All rights reserved. -- VMware Confidential

"""find-columns.py

This script gets a VCDB dump generated with pg_dump and finds the VMODL
tables and columns inside it. In VCDB any text field can contain a
serialized VMODL object which needs to be upgraded during B2B.
"""

import re
import sys
import pprint


TABLE_START = re.compile('COPY (?P<name>\w+) \((?P<columns>.*)\) FROM stdin;')
VMODL_FIELD = re.compile('.*versionId="(?P<id>\w+)".*')
TABLE_END = re.compile('\\\.$')


def ParseDump(vcdbDumpFileName):
   """Identifies all the tables in the VCDB dump file.
   """

   result = {}

   with open(vcdbDumpFileName) as f:
      for l in f:
         startMatch = TABLE_START.match(l)
         if startMatch:
            name = startMatch.group('name')
            result[name] = {
               'columns': startMatch.group('columns').split(', '),
               'values': [],
            }

            for v in f:
               if TABLE_END.match(v):
                  break

               result[name]['values'].append(v)

   return result


def FilterVmodlTables(tables):
   """Identifies all the tables that have VMODL fields.
   """
   result = []

   for t, info in tables.items():
      for v in info['values']:
         if VMODL_FIELD.match(v):
            result.append(t)

   return list(set(result))


def FilterVmodlColumns(tables, vmodlTableNames):
   """Identifies the columns that have VMODL fields.
   """

   vmodlColumns = {}
   versionIds = []

   for t in vmodlTableNames:
      vmodlColumns[t] = []
      for v in tables[t]['values']:
         if VMODL_FIELD.match(v):
            cols = v.split('\t')
            for index, c in enumerate(cols):
               m = VMODL_FIELD.match(c)
               if m:
                  vmodlColumns[t].append(tables[t]['columns'][index])
                  versionIds.append(m.group('id'))
      vmodlColumns[t] = list(set(vmodlColumns[t]))

   return vmodlColumns, list(set(versionIds))


def main():
   """Prints the identified VMODL columns and the version IDs found.
   """
   tables = ParseDump(sys.argv[1])
   vmodlTableNames = FilterVmodlTables(tables)
   vmodlColumns, versionIds = FilterVmodlColumns(tables, vmodlTableNames)

   print("Tables and their VMODL columns:")
   pprint.pprint(vmodlColumns)

   print("\nIdentified version IDs:")
   pprint.pprint(versionIds)


if __name__ == '__main__':
   sys.exit(main())

