# copyright 2003-2010 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

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

#

# This file is part of CubicWeb.

#

# CubicWeb is free software: you can redistribute it and/or modify it under the

# terms of the GNU Lesser General Public License as published by the Free

# Software Foundation, either version 2.1 of the License, or (at your option)

# any later version.

#

# CubicWeb is distributed in the hope that it will be useful, but WITHOUT

# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS

# FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more

# details.

#

# You should have received a copy of the GNU Lesser General Public License along

# with CubicWeb.  If not, see <http://www.gnu.org/licenses/>.

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

"""

from __future__ import with_statement

__docformat__ = "restructuredtext en"

import sys

import os

import tarfile

import tempfile

import shutil

import os.path as osp

from datetime import datetime

from glob import glob

from copy import copy

from warnings import warn

from contextlib import contextmanager

from logilab.common.deprecation import deprecated

from logilab.common.decorators import cached, clear_cache

from logilab.common.testlib import mock_object

from yams.constraints import SizeConstraint

from yams.schema2sql import eschema2sql, rschema2sql

from yams.schema import RelationDefinitionSchema

from cubicweb import AuthenticationError, ExecutionError

from cubicweb.selectors import is_instance

from cubicweb.schema import (ETYPE_NAME_MAP, META_RTYPES, VIRTUAL_RTYPES,

                             PURE_VIRTUAL_RTYPES,

                             CubicWebRelationSchema, order_eschemas)

from cubicweb.cwvreg import CW_EVENT_MANAGER

from cubicweb.dbapi import get_repository, repo_connect

from cubicweb.migration import MigrationHelper, yes

from cubicweb.server.session import hooks_control

from cubicweb.server import hook

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 ClearGroupMap(hook.Hook):

    __regid__ = 'cw.migration.clear_group_mapping'

    __select__ = hook.Hook.__select__ & is_instance('CWGroup')

    events = ('after_add_entity', 'after_update_entity',)

    def __call__(self):

        clear_cache(self.mih, 'group_mapping')

        self.mih._synchronized.clear()

    @classmethod

    def mih_register(cls, repo):

        # may be already registered in tests (e.g. unittest_migractions at

        # least)

        if not cls.__regid__ in repo.vreg['after_add_entity_hooks']:

            repo.vreg.register(ClearGroupMap)

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

        # no config on shell to a remote instance

        if config is not None and (cnx or connect):

            repo = self.repo

            self.session.data['rebuild-infered'] = False

            # register a hook to clear our group_mapping cache and the

            # self._synchronized set when some group is added or updated

            ClearGroupMap.mih = self

            ClearGroupMap.mih_register(repo)

            CW_EVENT_MANAGER.bind('after-registry-reload',

                                  ClearGroupMap.mih_register, repo)

            # notify we're starting maintenance (called instead of server_start

            # which is called on regular start

            repo.hm.call_hooks('server_maintenance', repo=repo)

        if not schema and not getattr(config, 'quick_start', False):

            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)

        # disable notification during migration

        with hooks_control(self.session, self.session.HOOKS_ALLOW_ALL, 'notification'):

            super(ServerMigrationHelper, self).migrate(vcconf, toupgrade, options)

    def cmd_process_script(self, migrscript, funcname=None, *args, **kwargs):

        try:

            return super(ServerMigrationHelper, self).cmd_process_script(

                  migrscript, funcname, *args, **kwargs)

        except ExecutionError, err:

            print >> sys.stderr, "-> %s" % err

        except:

            self.rollback()

            raise

    # Adjust docstring

    cmd_process_script.__doc__ = MigrationHelper.cmd_process_script.__doc__

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

        try:

            for source in repo.sources:

                try:

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

                except Exception, ex:

                    print '-> error trying to backup %s [%s]' % (source.uri, ex)

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

                        raise SystemExit(1)

                    else:

                        break

            else:

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

                for filename in os.listdir(tmpdir):

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

                bkup.close()

                # call hooks

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

                # done

                print '-> backup file',  backupfile

        finally:

            shutil.rmtree(tmpdir)

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

                         askconfirm=True):

        # check

        if not osp.exists(backupfile):

            raise ExecutionError("Backup file %s doesn't exist" % backupfile)

        if askconfirm and not self.confirm('Restore %s database from %s ?'

                                           % (self.config.appid, backupfile)):

            return

        # unpack backup

        tmpdir = tempfile.mkdtemp()

        try:

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

        except tarfile.ReadError:

            # assume restoring old backup

            shutil.copy(backupfile, osp.join(tmpdir, 'system'))

        else:

            for name in bkup.getnames():

                if name[0] in '/.':

                    raise ExecutionError('Security check failed, path starts with "/" or "."')

            bkup.close() # XXX seek error if not close+open !?!

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

            bkup.extractall(path=tmpdir)

            bkup.close()

        self.config.open_connections_pools = False

        repo = self.repo_connect()

        for source in repo.sources:

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

                continue

            try:

                source.restore(osp.join(tmpdir, source.uri), self.confirm, drop)

            except Exception, exc:

                print '-> error trying to restore %s [%s]' % (source.uri, exc)

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

                    raise SystemExit(1)

        shutil.rmtree(tmpdir)

        # call hooks

        repo.open_connections_pools()

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

            self.session.data['rebuild-infered'] = False

            return self._cnx

    @property

    def session(self):

        if self.config is not None:

            session = self.repo._get_session(self.cnx.sessionid)

            if session.pool is None:

                session.set_read_security(False)

                session.set_write_security(False)

            session.set_pool()

            return session

        # no access to session on remote instance

        return None

    def commit(self):

        if hasattr(self, '_cnx'):

            self._cnx.commit()

        if self.session:

            self.session.set_pool()

    def rollback(self):

        if hasattr(self, '_cnx'):

            self._cnx.rollback()

        if self.session:

            self.session.set_pool()

    def rqlexecall(self, rqliter, ask_confirm=False):

        for rql, kwargs in rqliter:

            self.rqlexec(rql, kwargs, ask_confirm=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({'commit': self.checkpoint,

                        'rollback': self.rollback,

                        'checkpoint': deprecated('[3.6] use commit')(self.checkpoint),

                        'sql': self.sqlexec,

                        'rql': self.rqlexec,

                        'rqliter': self.rqliter,

                        'schema': self.repo.get_schema(),

                        'cnx': self.cnx,

                        'fsschema': self.fs_schema,

                        'session' : self.session,

                        'repo' : self.repo,

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

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

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

                        })

        return context

    @cached

    def group_mapping(self):

        """cached group mapping"""

        return ss.group_mapping(self._cw)

    @cached

    def cstrtype_mapping(self):

        """cached constraint types mapping"""

        return ss.cstrtype_mapping(self._cw)

    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:

                self.cmd_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.cmd_process_script(apc, funcname, *args, **kwargs)

            finally:

                self.confirm = confirm

                self.execscript_confirm = execscript_confirm

                if self.config.free_wheel:

                    self.cmd_reactivate_verification_hooks()

    def install_custom_sql_scripts(self, directory, driver):

        for fpath in glob(osp.join(directory, '*.sql.%s' % driver)):

            newname = osp.basename(fpath).replace('.sql.%s' % driver,

                                                  '.%s.sql' % driver)

            warn('[3.5.6] rename %s into %s' % (fpath, newname),

                 DeprecationWarning)

            print '-> installing', fpath

            sqlexec(open(fpath).read(), self.session.system_sql, False,

                    delimiter=';;')

        for fpath in glob(osp.join(directory, '*.%s.sql' % driver)):

            print '-> installing', fpath

            sqlexec(open(fpath).read(), self.session.system_sql, False,