# Copyright (C) 2014, 2015, 2018 The Meme Factory, Inc.  http://www.meme.com/
#
#    This file is part of Gombe-MI.
#
#    Gombe-MI is free software; you can redistribute it and/or modify
#    it under the terms of the GNU General Public License as published by
#    the Free Software Foundation; either version 3 of the License, or
#    (at your option) any later version.
#
#    This program is distributed in the hope that it will be useful,
#    but WITHOUT ANY WARRANTY; without even the implied warranty of
#    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#    GNU General Public License for more details.
#
#    You should have received a copy of the GNU General Public License
#    along with Gombe-MI.  If not, see <http://www.gnu.org/licenses/>.
#
# Copy a schema from one db to another.
#
# Karl O. Pinc <kop@meme.com>
#

# Write python 3 compatible code.
from __future__ import print_function
from __future__ import unicode_literals
from __future__ import absolute_import
from __future__ import division

import tempfile
import os
import subprocess
import string
import random
import markupsafe
import logging
import psycopg2

from wtforms import (
    BooleanField,
    StringField)

from pyramid.view import view_config

from gmi_pyramid.lib.upload import (
    UploadEngine
    , NoTransactionEngine
    , Error
    , DBError
    , AuthFailError
    , UserWTForm
    , UserInitialPost
    , CredsLoadedForm
    , SessionDBHandler
    , LogSQLCommand
    , SQLData
    , NoOpProcessor
    , ExecuteSQL
    , format_exception
    )

from gmi_pyramid.lib.form_constants import(
    LIVE_DB
    , TEST_DB
    , DEFAULT_SCHEMA
    , CHECKED
    , UNCHECKED
    )

BIN = '/usr/bin'

log = logging.getLogger(__name__)

class CopySchemaInitialPost(UserInitialPost):
    schema     = DEFAULT_SCHEMA
    from_db    = LIVE_DB
    to_db      = TEST_DB
    force      = False


class CopySchemaWTForm(UserWTForm):
    '''The wtform used to connect to the "gombemi" db to authenticate .'''
    # We don't actually use the labels, wanting the template to
    # look (and render) like html, but I'll define them anyway
    # just to keep my hand in.
    schema       = StringField('Schema To Copy:')
    from_db      = StringField('From Database:')
    to_db        = StringField('To Database:')
    overwrite    = BooleanField('Overwrite Destination Schema:')
    force        = BooleanField('Allow Copy To The {0} DB:'.format(LIVE_DB))

    ivals        = CopySchemaInitialPost


class CopySchemaForm(CredsLoadedForm):
    '''
    Acts like a dict, but with extra methods.

    Attributes:
      uh      The UploadHandler instance using the form
    '''
    def __init__(self, uh, fc=CopySchemaWTForm, data={}, **kwargs):
        data.update(kwargs)
        super(CopySchemaForm, self).__init__(uh, fc, data)

    def read(self):
        '''
        Read form data from the client
        '''

        # Read parent's data
        super(CopySchemaForm, self).read()

        # Read our own data
        self['schema']          = self._form.schema.data.lower().strip()
        self['from_db']         = self._form.from_db.data.lower().strip()
        self['to_db']           = self._form.to_db.data.lower().strip()
        self['overwrite']       = self._form.overwrite.data
        self['force']           = self._form.force.data


    def write(self, result, errors):
        '''
        Produces the dict pyramid will use to render the form.
        '''
        response = super(CopySchemaForm, self).write(result, errors)

        response['schema']  = self['schema']
        response['from_db'] = self['from_db']
        response['to_db']   = self['to_db']

        if self['force']:
            force_checked = CHECKED
        else:
            force_checked = UNCHECKED
        response['force']   = force_checked

        if self['overwrite']:
            overwrite_checked = CHECKED
        else:
            overwrite_checked = UNCHECKED
        response['overwrite']   = overwrite_checked

        return response


# Setup errors
class NoSchemaError(Error):
    def __init__(self, e, descr='', detail=''):
        super(NoSchemaError, self).__init__(e, descr, detail)

class LiveDBTargetError(Error):
    def __init__(self, e, descr='', detail=''):
        super(LiveDBTargetError, self).__init__(e, descr, detail)

class NonAdminUserError(Error):
    def __init__(self, e, descr='', detail=''):
        super(NonAdminUserError, self).__init__(e, descr, detail)

class NoFromDBError(Error):
    def __init__(self, e, descr='', detail=''):
        super(NoFromDBError, self).__init__(e, descr, detail)

class NoToDBError(Error):
    def __init__(self, e, descr='', detail=''):
        super(NoToDBError, self).__init__(e, descr, detail)

class SameDBError(Error):
    def __init__(self, e, descr='', detail=''):
        super(SameDBError, self).__init__(e, descr, detail)

class InvalidSchemaError(Error):
    def __init__(self, e, descr='', detail=''):
        super(InvalidSchemaError, self).__init__(e, descr, detail)

class CopyToDefaultError(Error):
    def __init__(self, e, descr='', detail=''):
        super(CopyToDefaultError, self).__init__(e, descr, detail)

class SchemaExistsError(Error):
    def __init__(self, e, descr='', detail=''):
        super(SchemaExistsError, self).__init__(e, descr, detail)

class InconsistentDBError(Error):
    def __init__(self, e, descr='', detail=''):
        super(InconsistentDBError, self).__init__(e, descr, detail)

class SchemaDroppedError(InconsistentDBError):
    def __init__(self, e, descr='', detail=''):
        super(SchemaDroppedError, self).__init__(e, descr, detail)

class SchemaCopiedError(InconsistentDBError):
    def __init__(self, e, descr='', detail=''):
        super(SchemaCopiedError, self).__init__(e, descr, detail)

class CopyCommandError(Error):
    def __init__(self, e, descr='', detail=''):
        super(CopyCommandError, self).__init__(e, descr, detail)

# Errors copying using pg_dump/pg_restore.

class CopyOSError(CopyCommandError):
    '''The copy of the schema failed with an OSError.'''
    def __init__(self, e, ex, descr=''):
        detail = ('<p>The operating system reports: {0}</p>'
                  .format(markupsafe.escape(ex.strerror)))
        super(CopyOSError,self).__init__(e, descr, detail)

class PGDumpOSError(CopyOSError):
    def __init__(self, e, ex, descr=''):
        super(PGDumpOSError, self).__init__(e, ex, descr)

class PGRestoreOSError(CopyOSError):
    def __init__(self, e, ex, descr=''):
        super(PGRestoreOSError, self).__init__(e, ex, descr)

class CopyProcessError(CopyCommandError):
    '''The process which copies the schema failed.'''
    def __init__(self, e, retcode, descr=''):
        detail = ('<p>The process reports: {0}</p>'
                  .format(markupsafe.escape(translate_retcode(retcode))))
        super(CopyProcessError,self).__init__(e, descr, detail)

class PGDumpError(CopyProcessError):
    def __init__(self, e, retcode, descr=''):
        super(PGDumpError, self).__init__(e, retcode, descr)

class PGRestoreError(CopyProcessError):
    def __init__(self, e, retcode, descr=''):
        super(PGRestoreError, self).__init__(e, retcode, descr)


# Database statement execution errors

class DropSchemaError(DBError):
    def __init__(self, pgexc, schema, db):
        '''
        pgexc  The psycopg2 exception object
        '''
        super(DropSchemaError, self).__init__(
            pgexc, 'drop the ({0}) schema in the ({1}) db'.format(schema, db))

class VacuumFullError(DBError):
    def __init__(self, pgexc, db):
        '''
        pgexc  The psycopg2 exception object
        '''
        super(VacuumFullError, self).__init__(
            pgexc, 'VACUUM FULL the ({0}) db'.format(db))

class VacuumAnalyzeError(DBError):
    def __init__(self, pgexc, db):
        '''
        pgexc  The psycopg2 exception object
        '''
        super(VacuumAnalyzeError, self).__init__(
            pgexc, 'VACUUM ANALYZE the ({0}) db'.format(db))


# Utility functions

def schema_exists(cur, schema):
    '''Does the schema exist?'''
    cur.execute('SELECT 1 FROM pg_namespace'
                '  WHERE nspname = %s',
                (schema,))
    return cur.fetchone() != None

def translate_retcode(retcode):
    '''Translate a subprocess.returncode into a descriptive string.'''
    if retcode < 0:
        return 'Terminated by signal {0}'.format(-retcode)
    else:
        return os.strerror(retcode)


class CopySchemaHandler(SessionDBHandler):
    '''
    Abstract class to copy a schema from one db to another.
    '''
    def __init__(self, request):
        '''
        request A pyramid request instance
        '''
        super(CopySchemaHandler, self).__init__(request)

    def make_form(self):
        return CopySchemaForm(self)

    def write(self, result, errors):
        '''
        Setup dict to render resulting html form

        Returns:
          Dict pyramid will use to render the resulting form
          Reserved keys:
            errors      A list of UploadError exceptions.
            db_changed  Boolean. Whether the db was changed.
        '''
        response = super(CopySchemaHandler, self).write(result, errors)

        response['db_changed'] = (not response['errors']
                                  and self.uf['action'] != '')

        return response

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.

        Input:

        Side Effects:
        Yes, lots.
        '''
        return ExecuteSQL(ue, self)

    def render(self, errors, result):
        '''Instead of rendering, just our results so we can
        decide what to do next.

        Input:
          errors    List of Error instances
          result    Db connection result dict
        '''
        response = super(CopySchemaHandler, self).render(errors, result)
        return (response, errors)
            
    def to_db_detail(self):
        '''Return a string for the log detailing an SQL execution attempt.'''
        return ('In database {db}: By user ({user})'
                .format(db=self.uf['to_db'], user=self.uf['user']))

    def log_sql_success(self, msg):
        '''Write a success message to the log when executing sql.'''
        log.warning('{0}: {1}'
                   .format(msg, self.to_db_detail()))

    def log_sql_failure(self, msg, ex):
        '''Write a failture message to the log when executing sql.'''
        log.info('{0}: {1}{2}'
                 .format(msg, self.to_db_detail(), format_exception(ex)))


class FromSchemaHandler(CopySchemaHandler):
    '''Handler to check the source schema and do initial validation.
    
    Attributes:
       cur
    '''
    def __init__(self, request):
        '''
        request A pyramid request instance
        '''
        super(FromSchemaHandler, self).__init__(request)

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.
        Having the UploadEngine we can test the db content.

        Input:

        Side Effects:
        Yes, lots.
        '''
        self.cur = ue.cur
        
        uf = self.uf
        cur = self.cur

        # Is the logged in user a superuser?
        user = uf['user']
        cur.execute('SELECT 1 FROM pg_roles'
                    '  WHERE rolname = %s AND rolsuper',
                    (user,))
        if cur.fetchone() == None:
            ex = NonAdminUserError(
                    'Must be a superuser',
                    ('The supplied user ({0}) is not an administrator'
                     .format(user)))
            self.log_sql_failure('User not a superuser', ex)
            raise ex

        schema = uf['schema']
        if not schema_exists(cur, schema):
            ex = InvalidSchemaError(
                    'Source schema does not exist',
                    ('The supplied source schema ({0}) does not exist'
                     .format(schema)))
            self.log_sql_failure('Source schema does not exist', ex)
            raise ex
        
        return super(FromSchemaHandler, self).factory(ue)

    def get_data(self):
        '''We do all the work validating input so need no data.'''
        self.data = ()

    def val_input(self):
        '''
        Validate input needed beyond that required to connect to the db.

        Returns:
          A list of Error instances
        '''
        uf = self.uf
        errors = []

        schema = uf['schema']

        # Do we have the essential basic data?
        if schema == '':
            ex = NoSchemaError('No schema supplied')
            self.log_sql_failure('No schema supplied', ex)
            errors.append(ex)

        if uf['to_db'] == '':
            ex = NoToDBError('No to db supplied')
            self.log_sql_failure('No to db supplied', ex)
            errors.append(ex)
        elif uf['to_db'] == LIVE_DB and not uf['force']:
            ex = LiveDBTargetError(
                    'Cannot write into the production database',
                    'The "Allow ..." box must be checked to copy into'
                    'the {0} database.'
                    .format(LIVE_DB))
            self.log_sql_failure('Non-standard schema supplied', ex)
            errors.append(ex)

        if uf['from_db'] == '':
            ex = NoFromDBError('No from db supplied')
            self.log_sql_failure('No from db supplied', ex)
            errors.append(ex)
        elif uf['from_db'] == uf['to_db']:
            ex = SameDBError('The from db is the same as the to db',
                             'Both the from db and the to db are ({0}).'
                             .format(markupsafe.escape(uf['from_db'])))
            self.log_sql_failure('From db and to db are the same', ex)
            errors.append(ex)
        return errors


class DropSchemaHandler(CopySchemaHandler):
    '''
    Handler to drop the target schema from the target db.

    (Extra) Attributes:
      drop_attempted  True when the handler will attempt to drop the
                      schema.
    '''
    def __init__(self, request):
        '''
        request A pyramid request instance
        '''
        super(DropSchemaHandler, self).__init__(request)
        self.drop_attempted = False

    def factory(self, ue):
        '''Make a db loader function from an UploadEngine.
        Having the UploadEngine we can test the db content.

        Input:

        Side Effects:
        Yes, lots.
        '''
        cur = ue.cur
        schema = self.uf['schema']
        if schema_exists(cur, schema):
            if self.uf['overwrite']:
                self.drop_attempted = True
                return super(DropSchemaHandler, self).factory(ue)
            else:
                log.info('Abort:'
                         ' Schema ({0}) already exists in the ({1}) db'
                         .format(schema, self.uf['to_db']))
                raise SchemaExistsError(
                    'The target schema exists',
                    'The schema ({0}) exists in the {1} db.'
                    .format(markupsafe.escape(schema), self.uf['to_db']))
        else:
            log.info('Schema ({0}) does not exist in the ({1}) db'
                     .format(schema, self.uf['to_db']))
            return NoOpProcessor(ue, self)

    def log_dropschema_success(self):
        '''Write a success message to the log when dropping the schema.'''
        self.log_sql_success('Success dropping schema ({0})'
                             .format(self.uf['schema']))

    def log_dropschema_failure(self, ex):
        '''Write a failure message to the log when dropping the schema.'''
        self.log_sql_failure('Failure dropping schema ({0})'
                             .format(self.uf['schema']),
                             ex)

    def log_vacuumfull_success(self):
        '''Write a success message to the log when vacuum full succeeds.'''
        self.log_sql_success('Successful VACUUM FULL')

    def log_vacuumfull_failure(self, ex):
        '''Write a failure message to the log when vacuum full fails.'''
        self.log_sql_failure('Failed VACUUM FULL', ex)

    def get_data(self):
        '''
        Build and stash the SQL to be executed.
        
        Returns:
          List of SQCommand instances
        '''
        uf = self.uf
        sql = []

        schema = uf['schema']
        to_db  = uf['to_db']

        # Drop the schema
        sql.append(LogSQLCommand(
                'DROP SCHEMA {0} CASCADE'.format(schema)
                , ()
                , lambda ex: DropSchemaError(ex, schema, to_db)
                , log_success=self.log_dropschema_success
                , log_failure=self.log_dropschema_failure))
                   
        # Vacuum the target db
        sql.append(LogSQLCommand(
                'VACUUM FULL'
                , ()
                , lambda ex: VacuumFullError(ex, to_db)
                , log_success=self.log_vacuumfull_success
                , log_failure=self.log_vacuumfull_failure))

        self.data = SQLData(sql)


class VacuumHandler(CopySchemaHandler):
    '''
    Handler to add a user schema to the test db.

    Attributes:
    '''
    def __init__(self, request):
        '''
        request A pyramid request instance
        '''
        super(VacuumHandler, self).__init__(request)

    def log_vacuumanalyze_success(self):
        self.log_sql_success('Successful VACUUM ANALYZE')

    def log_vacuumanalyze_failure(self, ex):
        self.log_sql_failure('Failed VACUUM ANALYZE', ex)

    def get_data(self):
        '''
        Build and stash the SQL to be executed.
        
        Returns:
          List of SQCommand instances
        '''

        uf = self.uf
        sql = []

        schema = uf['schema']
        to_db  = uf['to_db']

        # Vacuum the target db
        sql.append(LogSQLCommand(
                'VACUUM ANALYZE'
                , ()
                , lambda ex: VacuumAnalyzeError(ex, to_db)
                , log_success=self.log_vacuumanalyze_success
                , log_failure=self.log_vacuumanalyze_failure))

        self.data = SQLData(sql)


class CheckFromSchemaEngine(UploadEngine):
    '''Customize credential failure error message.'''
    def authfailerror_factory(self):
        orig = super(CheckFromSchemaEngine, self).authfailerror_factory()
        return AuthFailError(
                ('Unable to login to the ({0}) db'
                 .format(self.uh.uf['from_db'])),
                orig.descr)

    def read_db(self):
        '''The handler has two db's, save the right one.'''
        self.db = self.uh.uf['from_db']


class ToDBEngine(NoTransactionEngine):
    '''Use the to_db to login.'''
    def read_db(self):
        self.db = self.uh.uf['to_db']

def copy_schema(uf, settings):
    '''Copy the desired schema using pg_dump/pg_restore.

    Input:
      uf        The upload form, holding some of the info we need
                to do the copy.
      settings  The configuration settings, holding some of the info we need
                to do the copy.

    Returns:
      Exception, or None.
    '''
    host     = settings['pg_host']
    port     = settings['pg_port']
    user     = uf['user']
    password = uf['password']
    schema   = uf['schema']
    from_db  = uf['from_db']
    to_db    = uf['to_db']
    
    # Create the service file.
    sfile = tempfile.NamedTemporaryFile(mode='w',
                                        prefix='copy_schema_service_file')
    servicefile = sfile.name
    # No point in writing the dbname in [to_db] since (with 9.2 anyway)
    # pg_restore does not use it.
    sfile.writelines(
        '[from_db]\n'
        'host={host}\n'
        'port={port}\n'
        'dbname={from_db}\n'
        'user={user}\n'
        'password={password}\n'
        '\n'
        '[to_db]\n'
        'host={host}\n'
        'port={port}\n'
        'user={user}\n'
        'password={password}\n'
        .format(from_db=from_db,
                to_db=to_db,
                user=user,
                password=password,
                host=host,
                port=port))
    sfile.flush()
        
    # Run the necessary commands.
    log_info = ('User ({0}): From db ({1}): To db ({2}): Schema ({3})'
                .format(user, from_db, to_db, schema))
    try:
        pg_dump = subprocess.Popen(
            ('{0}/pg_dump'.format(BIN),
             '--format', 'c', '--compress', '0',
             '--no-password',
             '--schema', schema,
             '--blobs'),
            bufsize=-1, close_fds=True,
            env={'PGSERVICEFILE': servicefile,
                 'PGSERVICE': 'from_db'},
            stdout=subprocess.PIPE)
    except OSError as ex:
        log.warn('Failed to execute pg_dump command: {0}: OS error {1}'
                 .format(log_info, ex.strerror))
        raise PGDumpOSError('Failed to execute pg_dump command', ex)

    try:
        pg_restore = subprocess.Popen(
            ('{0}/pg_restore'.format(BIN),
             '--single-transaction',
             '--no-password',
             '--dbname', to_db),
            bufsize=-1, close_fds=True,
            env={'PGSERVICEFILE': servicefile,
                 'PGSERVICE': 'to_db'},
            stdin=pg_dump.stdout)
    except OSError as ex:
        log.warn('Failed to execute pg_restore command: {0}: OS error {1}'
                 .format(log_info, ex.strerror))
        raise PGDumpOSError('Failed to execute pg_restore command', ex)

    retcode = pg_dump.wait()
    if retcode:
        log.warn('pg_dump process failed: {0}: {1}'
                 .format(log_info, translate_retcode(retcode)))
        raise PGDumpError('Failure of pg_dump process', retcode)
            
    retcode = pg_restore.wait()
    if retcode:
        log.warn('pg_restore process failed: {0}: {1}'
                 .format(log_info, translate_retcode(retcode)))
        raise PGRestoreError('Failure of pg_restore process', retcode)

    log.warn('Successfully copied schema: {0}'.format(log_info))

    # Get rid of the service file.
    sfile.close()


@view_config(route_name='copy_schema'
             , renderer='gmi_pyramid:templates/copy_schema.mak')
def copy_schema_view(request):

    # Log that we started doing something.
    # We don't have a form yet, so get data right out of POST.
    post =request.POST
    form_submitted = 'action' in post and post['action'] != ''
    if form_submitted:
        # The form is submitted.
        log.debug('Processing form submission')

    # Check that the source schema exists and can be copied.
    # This is a normal engine and checks for csrf.
    from_handler = FromSchemaHandler(request)
    response, errors = CheckFromSchemaEngine(from_handler).run()
    if errors:
        log.debug('Failed feasiblity test')
        return response

    # Drop schema from the target db.
    drop_handler = DropSchemaHandler(request)
    response, errors = ToDBEngine(drop_handler).run()
    if response['db_changed']:
        # Copy schema to target db.
        try:
            copy_schema(from_handler.uf, request.registry.settings)
        except CopyCommandError as ex:
            copyerrors = [ex]
            if drop_handler.drop_attempted:
                copyerrors.append(SchemaDroppedError(
                        'Target schema dropped but copy failed',
                        'The ({schema}) was dropped from the {db} db'
                        ' but the copy failed so now it no longer'
                        ' exists in the {db} db.'
                        .format(
                            schema=markupsafe.escape(
                                from_handler.uf['schema']),
                            db=markupsafe.escape(from_handler.uf['to_db']))))
            response['errors'].extend(copyerrors)
            response['db_changed'] = False
        else:
            # Vacuum analyze in the target db.
            vacuumout, vacuumerrors = (
                ToDBEngine(VacuumHandler(request)).run())
            # Merge results of modifying twice
            if vacuumerrors:
                vacuumerrors.append(SchemaCopiedError(
                        'VACUUM ANALYZE failed',
                        'VACUUM ANALYZE failed in the {db} db,'
                        ' performance may suffer'
                        .format(
                            db=markupsafe.escape(from_handler.uf['to_db']))))
                response['errors'].extend(vacuumerrors)
                del vacuumout['errors']
            response.update(vacuumout)
            response['db_changed'] &= not vacuumerrors
    if form_submitted:
        log.debug('Finished processing request')
    return response