# copyright 2003-2012 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 yams.constraints import SizeConstraint

from yams.schema2sql import eschema2sql, rschema2sql

from yams.schema import RelationDefinitionSchema

from cubicweb import CW_SOFTWARE_ROOT, AuthenticationError, ExecutionError

from cubicweb.predicates 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 import hook

try:

    from cubicweb.server import SOURCE_TYPES, schemaserial as ss

    from cubicweb.server.utils import manager_userpasswd

    from cubicweb.server.sqlutils import sqlexec, SQL_PREFIX

except ImportError: # LAX

    pass

def mock_object(**params):

    return type('Mock', (), params)()

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

            insert_lperms = self.repo.get_versions()['cubicweb'] < (3, 14, 0) and 'localperms' in config.available_cubes()

            if insert_lperms:

                cubes = config._cubes

                config._cubes += ('localperms',)

            try:

                schema = config.load_schema(expand_cubes=True)

            finally:

                if insert_lperms:

                    config._cubes = cubes

        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 self.session.allow_all_hooks_but('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:

            sys.stderr.write("-> %s\n" % err)

        except BaseException:

            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, format='native'):

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

        try:

            failed = False

            for source in repo.sources:

                try:

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

                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:

                        failed = True

            with open(osp.join(tmpdir, 'format.txt'), 'w') as format_file:

                format_file.write('%s\n' % format)

            with open(osp.join(tmpdir, 'versions.txt'), 'w') as version_file:

                versions = repo.get_versions()

                for cube, version in versions.iteritems():

                    version_file.write('%s %s\n' % (cube, version))

            if not failed:

                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, format='native'):

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

        if osp.isfile(osp.join(tmpdir, 'format.txt')):

            with open(osp.join(tmpdir, 'format.txt')) as format_file:

                written_format = format_file.readline().strip()

                if written_format in ('portable', 'native'):

                    format = written_format

        self.config.init_cnxset_pool = 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, format)

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

        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_cnxset_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.cnxset is None:

                session.set_read_security(False)

                session.set_write_security(False)

            session.set_cnxset()

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

    def rollback(self):

        if hasattr(self, '_cnx'):

            self._cnx.rollback()

        if self.session:

            self.session.set_cnxset()

    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,

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

                        })

        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 cmd_exec_event_script(self, event, cube=None, funcname=None,

                              *args, **kwargs):

        """execute a cube event scripts  `migration/<event>.py` where event

        is one of 'precreate', 'postcreate', 'preremove' and 'postremove'.

        """

        assert event in ('precreate', 'postcreate', 'preremove', 'postremove')

        if cube:

            cubepath = self.config.cube_dir(cube)

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

        elif kwargs.pop('apphome', False):

            apc = osp.join(self.config.apphome, 'migration', '%s.py' % event)

        else:

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

        if osp.exists(apc):