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

from os.path import join, exists

from datetime import datetime

from logilab.common.deprecation import deprecated_function, obsolete

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 CubicWebRelationSchema

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

        source = config.sources()['system']

        helper = get_adv_func_helper(source['db-driver'])

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

        app = config.appid

        backupfile = backupfile or join(config.backup_dir(),

                                        '%s-%s.dump' % (app, date))

        if exists(backupfile):

            if not self.confirm('a backup already exists for %s, overwrite it?' % app):

                return

        elif askconfirm and not self.confirm('backup %s database?' % app):

            return

        cmd = helper.backup_command(source['db-name'], source.get('db-host'),

                                    source.get('db-user'), backupfile,

                                    keepownership=False)

        while True:

            print cmd

            if os.system(cmd):

                print 'error while backuping the base'

                answer = self.confirm('continue anyway?',

                                      shell=False, abort=False, retry=True)

                if not answer:

                    raise SystemExit(1)

                if answer == 1: # 1: continue, 2: retry

                    break

            else:

                from cubicweb.toolsutils import restrict_perms_to_user

                print 'database backup:', backupfile

                restrict_perms_to_user(backupfile, self.info)

                break

    def restore_database(self, backupfile, drop=True):

        config = self.config

        source = config.sources()['system']

        helper = get_adv_func_helper(source['db-driver'])

        app = config.appid

        if not exists(backupfile):

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

        if self.confirm('restore %s database from %s ?' % (app, backupfile)):

            for cmd in helper.restore_commands(source['db-name'], source.get('db-host'),

                                               source.get('db-user'), backupfile,

                                               source['db-encoding'],

                                               keepownership=False, drop=drop):

                while True:

                    print cmd

                    if os.system(cmd):

                        print 'error while restoring the base'

                        answer = self.confirm('continue anyway?',

                                              shell=False, abort=False, retry=True)

                        if not answer:

                            raise SystemExit(1)

                        if answer == 1: # 1: continue, 2: retry

                            break

                    else:

                        break

            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)

            return self._cnx

    @property

    def session(self):

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

    @property

    @cached

    def rqlcursor(self):

        """lazy rql cursor"""

        # should not give session as cnx.cursor(), else we may try to execute

        # some query while no pool is set on the session (eg on entity attribute

        # access for instance)

        return self.cnx.cursor()

    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_function(self.cmd_sync_schema_props_perms),

                        'synchronize_eschema': deprecated_function(self.cmd_sync_schema_props_perms),

                        'synchronize_rschema': deprecated_function(self.cmd_sync_schema_props_perms),

                        })

        return context

    @cached

    def group_mapping(self):

        """cached group mapping"""

        return ss.group_mapping(self.rqlcursor)

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

                          *args, **kwargs):

        if cubepath:

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

        else:

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

        if 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 ('eid', 'has_text', 'identity'):

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

                    continue

                self._synchronize_rdef_schema(subj, rschema, obj)

        if syncperms:

            self._synchronize_permissions(rtype)

    def _synchronize_eschema(self, etype, syncperms=True):

        """synchronize properties of the persistent entity schema against

        its current definition:

        * description

        * internationalizable, fulltextindexed, indexed, meta

        * relations from/to this entity

        * permissions if `syncperms`

        """

        etype = str(etype)

        if etype in self._synchronized:

            return

        self._synchronized.add(etype)

        repoeschema = self.repo.schema.eschema(etype)

        try: