"""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-2008 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

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

"""

__docformat__ = "restructuredtext en"

import sys

import os

from os.path import join, exists

from mx.DateTime import now

from logilab.common.decorators import cached

from logilab.common.adbh import get_adv_func_helper

from yams.constraints import SizeConstraint

from yams.schema2sql import eschema2sql, rschema2sql

from cubicweb import AuthenticationError

from cubicweb.dbapi import get_repository, repo_connect

from cubicweb.common.migration import MigrationHelper, yes

try:

    from cubicweb.server import schemaserial as ss

    from cubicweb.server.utils import manager_userpasswd

    from cubicweb.server.sqlutils import sqlexec

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.new_schema = schema

        self._synchronized = set()

    @cached

    def repo_connect(self):

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

        return self.repo

    def shutdown(self):

        if self.repo is not None:

            self.repo.shutdown()

    def rewrite_vcconfiguration(self):

        """write current installed versions (of cubicweb software

        and of each used cube) into the database

        """

        self.cmd_set_property('system.version.cubicweb', self.config.cubicweb_version())

        for pkg in self.config.cubes():

            pkgversion = self.config.cube_version(pkg)

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

        self.commit()

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

        config = self.config

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

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

        date = 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:

                print 'database backup:', backupfile

                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'

    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)

        self.rewrite_configuration()

    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

        """

        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)

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

        return self.cnx.cursor(self.session)    

    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,

                        'newschema': self.new_schema,

                        'cnx': self.cnx,

                        'session' : self.session,

                        'repo' : self.repo,

                        })

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

    # base actions ############################################################

    def checkpoint(self):

        """checkpoint action"""

        if self.confirm('commit now ?', shell=False):

            self.commit()

    def cmd_add_cube(self, cube, update_database=True):

        """update_database is telling if the database schema should be updated

        or if only the relevant eproperty should be inserted (for the case where

        a cube has been extracted from an existing application, so the

        cube schema is already in there)

        """

        newcubes = super(ServerMigrationHelper, self).cmd_add_cube(

            cube)

        if not newcubes:

            return

        for pack in newcubes:

            self.cmd_set_property('system.version.'+pack,

                                  self.config.cube_version(pack))

        if not update_database:

            self.commit()

            return

        self.new_schema = self.config.load_schema()

        new = set()

        # execute pre-create files

        for pack in reversed(newcubes):

            self.exec_event_script('precreate', self.config.cube_dir(pack))

        # add new entity and relation types

        for rschema in self.new_schema.relations():

            if not rschema in self.repo.schema:

                self.cmd_add_relation_type(rschema.type)

                new.add(rschema.type)

        for eschema in self.new_schema.entities():

            if not eschema in self.repo.schema:

                self.cmd_add_entity_type(eschema.type)

                new.add(eschema.type)

        # check if attributes has been added to existing entities

        for rschema in self.new_schema.relations():

            existingschema = self.repo.schema.rschema(rschema.type)

            for (fromtype, totype) in rschema.iter_rdefs():

                if existingschema.has_rdef(fromtype, totype):

                    continue

                # check we should actually add the relation definition

                if not (fromtype in new or totype in new or rschema in new):

                    continue

                self.cmd_add_relation_definition(str(fromtype), rschema.type, 

                                                 str(totype))

        # execute post-create files

        for pack in reversed(newcubes):

            self.exec_event_script('postcreate', self.config.cube_dir(pack))

            self.commit()        

    def cmd_remove_cube(self, cube):

        removedcubes = super(ServerMigrationHelper, self).cmd_remove_cube(cube)

        if not removedcubes:

            return

        oldschema = self.new_schema

        self.new_schema = newschema = self.config.load_schema()

        reposchema = self.repo.schema

        # execute pre-remove files

        for pack in reversed(removedcubes):

            self.exec_event_script('preremove', self.config.cube_dir(pack))

        # remove cubes'entity and relation types

        for rschema in oldschema.relations():

            if not rschema in newschema and rschema in reposchema:

                self.cmd_drop_relation_type(rschema.type)

        for eschema in oldschema.entities():

            if not eschema in newschema and eschema in reposchema:

                self.cmd_drop_entity_type(eschema.type)

        for rschema in oldschema.relations():

            if rschema in newschema and rschema in reposchema: 

                # check if attributes/relations has been added to entities from 

                # other cubes

                for fromtype, totype in rschema.iter_rdefs():

                    if not newschema[rschema.type].has_rdef(fromtype, totype) and \

                           reposchema[rschema.type].has_rdef(fromtype, totype):

                        self.cmd_drop_relation_definition(

                            str(fromtype), rschema.type, str(totype))

        # execute post-remove files

        for pack in reversed(removedcubes):

            self.exec_event_script('postremove', self.config.cube_dir(pack))

            self.rqlexec('DELETE EProperty X WHERE X pkey %(pk)s',

                         {'pk': u'system.version.'+pack}, ask_confirm=False)

            self.commit()

    # schema migration actions ################################################

    def cmd_add_attribute(self, etype, attrname, attrtype=None, commit=True):

        """add a new attribute on the given entity type"""

        if attrtype is None:

            rschema = self.new_schema.rschema(attrname)

            attrtype = rschema.objects(etype)[0]

        self.cmd_add_relation_definition(etype, attrname, attrtype, commit=commit)

    def cmd_drop_attribute(self, etype, attrname, commit=True):

        """drop an existing attribute from the given entity type

        `attrname` is a string giving the name of the attribute to drop

        """

        rschema = self.repo.schema.rschema(attrname)

        attrtype = rschema.objects(etype)[0]

        self.cmd_drop_relation_definition(etype, attrname, attrtype, commit=commit)

    def cmd_rename_attribute(self, etype, oldname, newname, commit=True):

        """rename an existing attribute of the given entity type

        `oldname` is a string giving the name of the existing attribute

        `newname` is a string giving the name of the renamed attribute

        """

        eschema = self.new_schema.eschema(etype)

        attrtype = eschema.destination(newname)

        # have to commit this first step anyway to get the definition

        # actually in the schema

        self.cmd_add_attribute(etype, newname, attrtype, commit=True)

        # skipp NULL values if the attribute is required

        rql = 'SET X %s VAL WHERE X is %s, X %s VAL' % (newname, etype, oldname)

        card = eschema.rproperty(newname, 'cardinality')[0]

        if card == '1':

            rql += ', NOT X %s NULL' % oldname

        self.rqlexec(rql, ask_confirm=self.verbosity>=2)

        self.cmd_drop_attribute(etype, oldname, commit=commit)

    def cmd_add_entity_type(self, etype, auto=True, commit=True):

        """register a new entity type

        in auto mode, automatically register entity's relation where the

        targeted type is known

        """

        applschema = self.repo.schema