"""a class implementing basic actions used in migration scripts.

The following schema actions are supported for now:

* add/drop/rename attribute

* add/drop entity/relation type

* rename entity type

The following data actions are supported for now:

* add an entity

* execute raw RQL queries

:organization: Logilab

:copyright: 2001-2010 LOGILAB S.A. (Paris, FRANCE), license is LGPL v2.

:contact: http://www.logilab.fr/ -- mailto:contact@logilab.fr

:license: GNU Lesser General Public License, v2.1 - http://www.gnu.org/licenses

"""

__docformat__ = "restructuredtext en"

import sys

import os

import tarfile

import tempfile

import shutil

import os.path as osp

from datetime import datetime

from glob import glob

from warnings import warn

from logilab.common.deprecation import deprecated

from logilab.common.decorators import cached, clear_cache

from logilab.common.adbh import get_adv_func_helper

from yams.constraints import SizeConstraint

from yams.schema2sql import eschema2sql, rschema2sql

from cubicweb import AuthenticationError, ETYPE_NAME_MAP

from cubicweb.schema import (META_RTYPES, VIRTUAL_RTYPES,

                             CubicWebRelationSchema, order_eschemas)

from cubicweb.dbapi import get_repository, repo_connect

from cubicweb.migration import MigrationHelper, yes

try:

    from cubicweb.server import SOURCE_TYPES, schemaserial as ss

    from cubicweb.server.utils import manager_userpasswd, ask_source_config

    from cubicweb.server.sqlutils import sqlexec, SQL_PREFIX

except ImportError: # LAX

    pass

class ServerMigrationHelper(MigrationHelper):

    """specific migration helper for server side  migration scripts,

    providind actions related to schema/data migration

    """

    def __init__(self, config, schema, interactive=True,

                 repo=None, cnx=None, verbosity=1, connect=True):

        MigrationHelper.__init__(self, config, interactive, verbosity)

        # no config on shell to a remote instance

        if not interactive:

            assert cnx

            assert repo

        if cnx is not None:

            assert repo

            self._cnx = cnx

            self.repo = repo

            if config is not None:

                self.session.data['rebuild-infered'] = False

        elif connect:

            self.repo_connect()

        if not schema:

            schema = config.load_schema(expand_cubes=True)

        self.fs_schema = schema

        self._synchronized = set()

    # overriden from base MigrationHelper ######################################

    @cached

    def repo_connect(self):

        self.repo = get_repository(method='inmemory', config=self.config)

        return self.repo

    def cube_upgraded(self, cube, version):

        self.cmd_set_property('system.version.%s' % cube.lower(),

                              unicode(version))

        self.commit()

    def shutdown(self):

        if self.repo is not None:

            self.repo.shutdown()

    def migrate(self, vcconf, toupgrade, options):

        if not options.fs_only:

            if options.backup_db is None:

                self.backup_database()

            elif options.backup_db:

                self.backup_database(askconfirm=False)

        super(ServerMigrationHelper, self).migrate(vcconf, toupgrade, options)

    def cmd_process_script(self, migrscript, funcname=None, *args, **kwargs):

        """execute a migration script

        in interactive mode,  display the migration script path, ask for

        confirmation and execute it if confirmed

        """

        try:

            if migrscript.endswith('.sql'):

                if self.execscript_confirm(migrscript):

                    sqlexec(open(migrscript).read(), self.session.system_sql)

            elif migrscript.endswith('.py') or migrscript.endswith('.txt'):

                return super(ServerMigrationHelper, self).cmd_process_script(

                    migrscript, funcname, *args, **kwargs)

            else:

                print

                print ('-> ignoring %s, only .py .sql and .txt scripts are considered' %

                       migrscript)

                print

            self.commit()

        except:

            self.rollback()

            raise

    # server specific migration methods ########################################

    def backup_database(self, backupfile=None, askconfirm=True):

        config = self.config

        repo = self.repo_connect()

        # paths

        timestamp = datetime.now().strftime('%Y-%m-%d_%H-%M-%S')

        instbkdir = osp.join(config.appdatahome, 'backup')

        if not osp.exists(instbkdir):

            os.makedirs(instbkdir)

        backupfile = backupfile or osp.join(instbkdir, '%s-%s.tar.gz'

                                            % (config.appid, timestamp))

        # check backup has to be done

        if osp.exists(backupfile) and not \

                self.confirm('Backup file %s exists, overwrite it?' % backupfile):

            print '-> no backup done.'

            return

        elif askconfirm and not self.confirm('Backup %s database?' % config.appid):

            print '-> no backup done.'

            return

        open(backupfile,'w').close() # kinda lock

        os.chmod(backupfile, 0600)

        # backup

        tmpdir = tempfile.mkdtemp(dir=instbkdir)

        try:

            for source in repo.sources:

                try:

                    source.backup(osp.join(tmpdir, source.uri))

                except Exception, exc:

                    print '-> error trying to backup %s [%s]' % (source.uri, exc)

                    if not self.confirm('Continue anyway?', default='n'):

                        raise SystemExit(1)

                    else:

                        break

            else:

                bkup = tarfile.open(backupfile, 'w|gz')

                for filename in os.listdir(tmpdir):

                    bkup.add(osp.join(tmpdir,filename), filename)

                bkup.close()

                # call hooks

                repo.hm.call_hooks('server_backup', repo=repo, timestamp=timestamp)

                # done

                print '-> backup file',  backupfile

        finally:

            shutil.rmtree(tmpdir)

    def restore_database(self, backupfile, drop=True, systemonly=True,

                         askconfirm=True):

        # check

        if not osp.exists(backupfile):

            raise Exception("Backup file %s doesn't exist" % backupfile)

            return

        if askconfirm and not self.confirm('Restore %s database from %s ?'

                                           % (self.config.appid, backupfile)):

            return

        # unpack backup

        tmpdir = tempfile.mkdtemp()

        try:

            bkup = tarfile.open(backupfile, 'r|gz')

        except tarfile.ReadError:

            # assume restoring old backup

            shutil.copy(backupfile, osp.join(tmpdir, 'system'))

        else:

            for name in bkup.getnames():

                if name[0] in '/.':

                    raise Exception('Security check failed, path starts with "/" or "."')

            bkup.close() # XXX seek error if not close+open !?!

            bkup = tarfile.open(backupfile, 'r|gz')

            bkup.extractall(path=tmpdir)

            bkup.close()

        self.config.open_connections_pools = False

        repo = self.repo_connect()

        for source in repo.sources:

            if systemonly and source.uri != 'system':

                continue

            try:

                source.restore(osp.join(tmpdir, source.uri), self.confirm, drop)

            except Exception, exc:

                print '-> error trying to restore %s [%s]' % (source.uri, exc)

                if not self.confirm('Continue anyway?', default='n'):

                    raise SystemExit(1)

        shutil.rmtree(tmpdir)

        # call hooks

        repo.open_connections_pools()

        repo.hm.call_hooks('server_restore', repo=repo, timestamp=backupfile)

        print '-> database restored.'

    @property

    def cnx(self):

        """lazy connection"""

        try:

            return self._cnx

        except AttributeError:

            sourcescfg = self.repo.config.sources()

            try:

                login = sourcescfg['admin']['login']

                pwd = sourcescfg['admin']['password']

            except KeyError:

                login, pwd = manager_userpasswd()

            while True:

                try:

                    self._cnx = repo_connect(self.repo, login, password=pwd)

                    if not 'managers' in self._cnx.user(self.session).groups:

                        print 'migration need an account in the managers group'

                    else:

                        break

                except AuthenticationError:

                    print 'wrong user/password'

                except (KeyboardInterrupt, EOFError):

                    print 'aborting...'

                    sys.exit(0)

                try:

                    login, pwd = manager_userpasswd()

                except (KeyboardInterrupt, EOFError):

                    print 'aborting...'

                    sys.exit(0)

            self.session.keep_pool_mode('transaction')

            self.session.data['rebuild-infered'] = False

            return self._cnx

    @property

    def session(self):

        if self.config is not None:

            return self.repo._get_session(self.cnx.sessionid)

        # no access to session on remote instance

        return None

    def commit(self):

        if hasattr(self, '_cnx'):

            self._cnx.commit()

    def rollback(self):

        if hasattr(self, '_cnx'):

            self._cnx.rollback()

    def rqlexecall(self, rqliter, cachekey=None, ask_confirm=True):

        for rql, kwargs in rqliter:

            self.rqlexec(rql, kwargs, cachekey, ask_confirm)

    @cached

    def _create_context(self):

        """return a dictionary to use as migration script execution context"""

        context = super(ServerMigrationHelper, self)._create_context()

        context.update({'commit': self.checkpoint,

                        'rollback': self.rollback,

                        'checkpoint': deprecated('[3.6] use commit')(self.checkpoint),

                        'sql': self.sqlexec,

                        'rql': self.rqlexec,

                        'rqliter': self.rqliter,

                        'schema': self.repo.get_schema(),

                        'cnx': self.cnx,

                        'fsschema': self.fs_schema,

                        'session' : self.session,

                        'repo' : self.repo,

                        'synchronize_schema': deprecated()(self.cmd_sync_schema_props_perms), # 3.4

                        'synchronize_eschema': deprecated()(self.cmd_sync_schema_props_perms), # 3.4

                        'synchronize_rschema': deprecated()(self.cmd_sync_schema_props_perms), # 3.4

                        })

        return context

    @cached

    def group_mapping(self):

        """cached group mapping"""

        return ss.group_mapping(self._cw)

    def exec_event_script(self, event, cubepath=None, funcname=None,

                          *args, **kwargs):

        if cubepath:

            apc = osp.join(cubepath, 'migration', '%s.py' % event)

        else:

            apc = osp.join(self.config.migration_scripts_dir(), '%s.py' % event)

        if osp.exists(apc):

            if self.config.free_wheel:

                from cubicweb.server.hooks import setowner_after_add_entity

                self.repo.hm.unregister_hook(setowner_after_add_entity,

                                             'after_add_entity', '')

                self.cmd_deactivate_verification_hooks()

            self.info('executing %s', apc)

            confirm = self.confirm

            execscript_confirm = self.execscript_confirm

            self.confirm = yes

            self.execscript_confirm = yes

            try:

                return self.cmd_process_script(apc, funcname, *args, **kwargs)

            finally:

                self.confirm = confirm

                self.execscript_confirm = execscript_confirm

                if self.config.free_wheel:

                    self.repo.hm.register_hook(setowner_after_add_entity,

                                               'after_add_entity', '')

                    self.cmd_reactivate_verification_hooks()

    def install_custom_sql_scripts(self, directory, driver):

        self.session.set_pool() # ensure pool is set

        for fpath in glob(osp.join(directory, '*.sql.%s' % driver)):

            newname = osp.basename(fpath).replace('.sql.%s' % driver,

                                                  '.%s.sql' % driver)

            warn('[3.5.6] rename %s into %s' % (fpath, newname),

                 DeprecationWarning)

            print '-> installing', fpath

            sqlexec(open(fpath).read(), self.session.system_sql, False,

                    delimiter=';;')

        for fpath in glob(osp.join(directory, '*.%s.sql' % driver)):

            print '-> installing', fpath

            sqlexec(open(fpath).read(), self.session.system_sql, False,

                    delimiter=';;')

    # schema synchronization internals ########################################

    def _synchronize_permissions(self, erschema, teid):

        """permission synchronization for an entity or relation type"""

        assert teid, erschema

        if 'update' in erschema.ACTIONS or erschema.final:

            # entity type

            exprtype = u'ERQLExpression'

        else:

            # relation type

            exprtype = u'RRQLExpression'

        gm = self.group_mapping()

        confirm = self.verbosity >= 2

        # * remove possibly deprecated permission (eg in the persistent schema

        #   but not in the new schema)

        # * synchronize existing expressions

        # * add new groups/expressions

        for action in erschema.ACTIONS:

            perm = '%s_permission' % action

            # handle groups

            newgroups = list(erschema.get_groups(action))

            for geid, gname in self.rqlexec('Any G, GN WHERE T %s G, G name GN, '

                                            'T eid %%(x)s' % perm, {'x': teid}, 'x',

                                            ask_confirm=False):

                if not gname in newgroups:

                    if not confirm or self.confirm('remove %s permission of %s to %s?'

                                                   % (action, erschema, gname)):

                        self.rqlexec('DELETE T %s G WHERE G eid %%(x)s, T eid %s'

                                     % (perm, teid),

                                     {'x': geid}, 'x', ask_confirm=False)

                else:

                    newgroups.remove(gname)

            for gname in newgroups:

                if not confirm or self.confirm('grant %s permission of %s to %s?'

                                               % (action, erschema, gname)):

                    self.rqlexec('SET T %s G WHERE G eid %%(x)s, T eid %s'

                                 % (perm, teid),

                                 {'x': gm[gname]}, 'x', ask_confirm=False)

            # handle rql expressions

            newexprs = dict((expr.expression, expr) for expr in erschema.get_rqlexprs(action))

            for expreid, expression in self.rqlexec('Any E, EX WHERE T %s E, E expression EX, '