from pgtoolkit.toolbase import SimpleTool

from pgtoolkit import pgbrowser

from pycolib.prettysize import prettysize_short

from pycolib.ansicolor import highlight

class BigTablesTool(SimpleTool):

    """List largest tables.

    Reads size statistics of tables and indexes from pgcatalog.

    """

    def __init__(self):

        SimpleTool.__init__(self, name='bigtables')

    def specify_args(self):

        SimpleTool.specify_args(self)

        self.parser.add_argument('-n', '--limit', metavar='NUM', dest='limit', type=int, default=5, help='Show NUM biggest tables.')

        self.parser.add_argument('-v', '--details', dest='details', action='store_true', help='Show sizes of data and individual indexes.')

    def main(self):

        browser = pgbrowser.PgBrowser(self.pgm.get_conn('target'))

        # scan all tables from all shemas, remember names and sizes

        all_tables = []

        all_indexes = []

        schemas = browser.list_schemas()

        for schema in schemas:

            tables = browser.list_tables(schema['name'])

            for table in tables:

                table_name = '%s.%s' % (schema['name'], table['name'])

                indexes = browser.list_indexes(table['name'], schema['name'])

                for index in indexes:

                    all_indexes.append({'name': index['name'], 'table': table_name, 'size': index['size']})

                all_tables.append({'name': table_name, 'size': table['size'], 'indexes': indexes})

        # print names and sizes of 20 largest tables

        for table in sorted(all_tables, reverse=True, key=lambda x: x['size'])[:self.args.limit]:

            size_of_indexes = sum(index['size'] for index in table['indexes'])

            print(highlight(1) + prettysize_short(table['size'] + size_of_indexes, trailing_zeros=True).rjust(8) + highlight(0),

                  '(total)'.ljust(8),

                  highlight(1) + table['name'] + highlight(0), sep='  ')

            if self.args.details:

                print(prettysize_short(table['size'], trailing_zeros=True).rjust(8),

                      '(data)'.ljust(8), sep='  ')

                for index in sorted(table['indexes'], reverse=True, key=lambda x: x['size']):

                    print(prettysize_short(index['size'], trailing_zeros=True).rjust(8),

                          '(index)'.ljust(8), index['name'], sep='  ')

            print()

cls = BigTablesTool