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

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

def set_sql_prefix(prefix):

    """3.1.5 migration function: allow to unset/reset SQL_PREFIX"""

    for module in ('checkintegrity', 'migractions', 'schemahooks',

                   'sources.rql2sql', 'sources.native'):

        try:

            sys.modules['cubicweb.server.%s' % module].SQL_PREFIX = prefix

            print 'changed SQL_PREFIX for %s' % module

        except KeyError:

            pass

def update_database(repo):

    """3.1.3 migration function: update database schema by adding SQL_PREFIX to

    entity type tables and columns

    """

    pool = repo._get_pool()

    source = repo.system_source

    sqlcu = pool['system']

    for etype in repo.schema.entities():

        if etype.is_final():

            continue

        try:

            sqlcu.execute('ALTER TABLE %s RENAME TO cw_%s' % (etype, etype))

            print 'renamed %s table for source %s' % (etype, uri)

        except:

            pass

        for rschema in etype.subject_relations():

            if rschema == 'has_text':

                continue

            if rschema.is_final() or rschema.inlined:

                sqlcu.execute('ALTER TABLE cw_%s RENAME %s TO cw_%s'

                              % (etype, rschema, rschema))

                print 'renamed %s.%s column for source %s' % (

                    etype, rschema, uri)

    pool.commit()

    repo._free_pool(pool)

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

        try:

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

        except:

            import traceback

            traceback.print_exc()

            print '3.1.5 migration'

            # XXX 3.1.5 migration

            set_sql_prefix('')

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

            update_database(self.repo)

            set_sql_prefix('cw_')

        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,

                        # XXX deprecate

                        'newschema': self.fs_schema,

                        'fsschema': self.fs_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):

        self.cmd_add_cubes( (cube,), update_database)

    def cmd_add_cubes(self, cubes, 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_cubes(cubes)

        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

        newcubes_schema = self.config.load_schema(construction_mode='non-strict')

        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 newcubes_schema.relations():

            if not rschema in self.repo.schema:

                self.cmd_add_relation_type(rschema.type)

                new.add(rschema.type)

        for eschema in newcubes_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 newcubes_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

        fsschema = self.fs_schema

        removedcubes_schema = self.config.load_schema(construction_mode='non-strict')

        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 fsschema.relations():

            if not rschema in removedcubes_schema and rschema in reposchema:

                self.cmd_drop_relation_type(rschema.type)

        for eschema in fsschema.entities():

            if not eschema in removedcubes_schema and eschema in reposchema:

                self.cmd_drop_entity_type(eschema.type)

        for rschema in fsschema.relations():

            if rschema in removedcubes_schema and rschema in reposchema: 

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