"""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-2009 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 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.dbapi import get_repository, repo_connect

from cubicweb.common.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)

        if not interactive:

            assert cnx

            assert repo

        if cnx is not None:

            assert repo

            self._cnx = cnx

            self.repo = repo

        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 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)

            else:

                return super(ServerMigrationHelper, self).process_script(

                    migrscript, funcname, *args, **kwargs)

            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)

        for source in repo.sources:

            try:

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

            except Exception, exc:

                print '-> error trying to backup [%s]' % exc

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

                    raise SystemExit(1)

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

        for filename in os.listdir(tmpdir):

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

        bkup.close()

        shutil.rmtree(tmpdir)

        # call hooks

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

        # done

        print '-> backup file',  backupfile

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

                         askconfirm=True):

        config = self.config

        repo = self.repo_connect()

        # 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 ?'

                                           % (config.appid, backupfile)):

            return

        # unpack backup

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

        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')

        tmpdir = tempfile.mkdtemp()

        bkup.extractall(path=tmpdir)

        for source in repo.sources:

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

                continue

            try:

                source.restore(osp.join(tmpdir, source.uri), drop=drop)

            except Exception, exc:

                print '-> error trying to restore [%s]' % exc

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

                    raise SystemExit(1)

        bkup.close()

        shutil.rmtree(tmpdir)

        # call hooks

        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, 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')

            return self._cnx

    @property

    def session(self):

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

    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({'checkpoint': self.checkpoint,

                        'sql': self.sqlexec,

                        'rql': self.rqlexec,

                        'rqliter': self.rqliter,

                        'schema': self.repo.schema,

                        'fsschema': self.fs_schema,

                        'session' : self.session,

                        'repo' : self.repo,

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

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

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

                        })

        return context

    @cached

    def group_mapping(self):

        """cached group mapping"""

        self.session.set_pool()

        return ss.group_mapping(self.session)

    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.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.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.reactivate_verification_hooks()

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

    def _synchronize_permissions(self, ertype):

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

        if ertype in VIRTUAL_RTYPES:

            return

        newrschema = self.fs_schema[ertype]

        teid = self.repo.schema[ertype].eid

        if 'update' in newrschema.ACTIONS or newrschema.is_final():

            # entity type

            exprtype = u'ERQLExpression'

        else:

            # relation type

            exprtype = u'RRQLExpression'

        assert teid, ertype

        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 newrschema.ACTIONS:

            perm = '%s_permission' % action

            # handle groups

            newgroups = list(newrschema.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, ertype, 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, ertype, 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 newrschema.get_rqlexprs(action))

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

                                                    'T eid %s' % (perm, teid),

                                                    ask_confirm=False):

                if not expression in newexprs:

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

                                                   % (expression, action, ertype)):

                        # deleting the relation will delete the expression entity

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

                                     % (perm, teid),

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

                else:

                    newexprs.pop(expression)

            for expression in newexprs.values():

                expr = expression.expression

                if not confirm or self.confirm('add %s expression for %s permission of %s?'

                                               % (expr, action, ertype)):

                    self.rqlexec('INSERT RQLExpression X: X exprtype %%(exprtype)s, '

                                 'X expression %%(expr)s, X mainvars %%(vars)s, T %s X '

                                 'WHERE T eid %%(x)s' % perm,

                                 {'expr': expr, 'exprtype': exprtype,

                                  'vars': expression.mainvars, 'x': teid}, 'x',

                                 ask_confirm=False)

    def _synchronize_rschema(self, rtype, syncrdefs=True, syncperms=True):

        """synchronize properties of the persistent relation schema against its

        current definition:

        * description

        * symetric, meta

        * inlined

        * relation definitions if `syncrdefs`

        * permissions if `syncperms`

        physical schema changes should be handled by repository's schema hooks

        """

        rtype = str(rtype)

        if rtype in self._synchronized:

            return

        self._synchronized.add(rtype)

        rschema = self.fs_schema.rschema(rtype)

        self.rqlexecall(ss.updaterschema2rql(rschema),

                        ask_confirm=self.verbosity>=2)

        reporschema = self.repo.schema.rschema(rtype)

        if syncrdefs:

            for subj, obj in rschema.iter_rdefs():

                if not reporschema.has_rdef(subj, obj):