# copyright 2003-2014 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

"""

__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.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 import repoapi

from cubicweb.migration import MigrationHelper, yes

from cubicweb.server import hook, schemaserial as ss

from cubicweb.server.schema2sql import eschema2sql, rschema2sql, unique_index_name

from cubicweb.server.utils import manager_userpasswd

from cubicweb.server.sqlutils import sqlexec, SQL_PREFIX

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,

    providing 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

            self.session = cnx._session

        elif connect:

            self.repo_connect()

        else:

            self.session = None

        # no config on shell to a remote instance

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

            repo = self.repo

            # 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 config.quick_start:

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

        try:

            login = self.repo.config.default_admin_config['login']

            pwd = self.repo.config.default_admin_config['password']

        except KeyError:

            login, pwd = manager_userpasswd()

        while True:

            try:

                self.cnx = repoapi.connect(self.repo, login, password=pwd)

                if not 'managers' in self.cnx.user.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 = self.repo._get_session(self.cnx.sessionid)

    # overriden from base MigrationHelper ######################################

    @cached

    def repo_connect(self):

        self.repo = repoapi.get_repository(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.cnx.allow_all_hooks_but('notification'):

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

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

        with self.cnx._cnx.ensure_cnx_set:

            try:

                return super(ServerMigrationHelper, self).cmd_process_script(

                      migrscript, funcname, *args, **kwargs)

            except ExecutionError as 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

        source = repo.system_source

        tmpdir = tempfile.mkdtemp()

        try:

            failed = False

            try:

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

            except Exception as 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, 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()

        source = repo.system_source

        try:

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

        except Exception as 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.'

    def commit(self):

        if hasattr(self, 'cnx'):

            self.cnx.commit(free_cnxset=False)

    def rollback(self):

        if hasattr(self, 'cnx'):

            self.cnx.rollback(free_cnxset=False)

    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.cnx._cnx,

                        'repo' : self.repo,

                        })

        return context

    @cached

    def group_mapping(self):

        """cached group mapping"""

        return ss.group_mapping(self.cnx)

    def cstrtype_mapping(self):

        """cached constraint types mapping"""

        return ss.cstrtype_mapping(self.cnx)

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

            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:

                if event == 'postcreate':

                    with self.cnx.allow_all_hooks_but():

                        return self.cmd_process_script(apc, funcname, *args, **kwargs)

                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 cmd_install_custom_sql_scripts(self, cube=None):

        """install a cube custom sql scripts `schema/*.<driver>.sql` where

        <driver> depends on the instance main database backend (eg 'postgres',

        'mysql'...)

        """

        driver = self.repo.system_source.dbdriver

        if cube is None:

            directory = osp.join(CW_SOFTWARE_ROOT, 'schemas')

        else: