from pgtoolkit.toolbase import SimpleTool

from pgtoolkit import pgbrowser

class AnalyzeAllTool(SimpleTool):

    """

    Analyze/vacuum all tables in selected schemas.

    Partially emulates VACUUM ANALYZE VERBOSE query.

    But this program is more configurable, skips pg_catalog etc.

    """

    def __init__(self):

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

        self.target_isolation_level = 'autocommit'

    def specify_args(self):

        SimpleTool.specify_args(self)

        self.parser.add_argument('-s', dest='schema', nargs='*', help='Schema filter')

        self.parser.add_argument('--vacuum', action='store_true', help='Call VACUUM ANALYZE')

        self.parser.add_argument('--vacuum-full', action='store_true', help='Call VACUUM FULL ANALYZE')

        self.parser.add_argument('--reindex', action='store_true', help='Call REINDEX TABLE')

    def main(self):

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

        query_patterns = ['ANALYZE %s.%s;']

        if self.args.vacuum:

            query_patterns = ['VACUUM ANALYZE %s.%s;']

        if self.args.vacuum_full:

            query_patterns = ['VACUUM FULL ANALYZE %s.%s;']

        if self.args.reindex:

            query_patterns += ['REINDEX TABLE %s.%s;']

        schema_list = self.args.schema

        if not schema_list:

            schema_list = [schema['name'] for schema in browser.list_schemas() if not schema['system']]

        for schema in schema_list:

            tables = browser.list_tables(schema=schema)

            with self.pgm.cursor('target') as curs:

                for table in tables:

                    for query_pattern in query_patterns:

                        query = query_pattern % (schema, table['name'])

                        self.log.info(query)

                        curs.execute(query, [])

cls = AnalyzeAllTool