"""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), 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 datetime import datetime

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, ETYPE_NAME_MAP

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

    @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 = 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'

    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

        """

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

        # 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,

                        })

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