# coding: utf-8

# copyright 2015-2016 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/>.

from collections import defaultdict

from itertools import chain

import logging

from uuid import uuid4

from six import text_type

from six.moves import range

from cubicweb.dataimport import stores, pgstore

from cubicweb.server.schema2sql import eschema_sql_def

class MassiveObjectStore(stores.RQLObjectStore):

    """Store for massive import of data, with delayed insertion of meta data.

    WARNINGS:

    - This store may only be used with PostgreSQL for now, as it relies

      on the COPY FROM method, and on specific PostgreSQL tables to get all

      the indexes.

    - This store can only insert relations that are not inlined (i.e.,

      which do *not* have inlined=True in their definition in the schema),

      unless they are specified as entity attributes.

    It should be used as follows:

       store = MassiveObjectStore(cnx)

       eid_p = store.prepare_insert_entity('Person',

                                           cwuri=u'http://dbpedia.org/toto',

                                           name=u'Toto')

       eid_loc = store.prepare_insert_entity('Location',

                                             cwuri=u'http://geonames.org/11111',

                                             name=u'Somewhere')

       store.prepare_insert_relation(eid_p, 'lives_in', eid_loc)

       store.flush()

       ...

       store.commit()

       store.finish()

    Full-text indexation is not handled, you'll have to reindex the proper entity types by yourself

    if desired.

    """

    def __init__(self, cnx, slave_mode=False, eids_seq_range=10000, metagen=None):

        """Create a MassiveObject store, with the following arguments:

        - `cnx`, a connection to the repository

        - `metagen`, optional :class:`MetadataGenerator` instance

        - `eids_seq_range`: size of eid range reserved by the store for each batch

        """

        super(MassiveObjectStore, self).__init__(cnx)

        self.uuid = text_type(uuid4()).replace('-', '')

        self.slave_mode = slave_mode

        if metagen is None:

            metagen = stores.MetadataGenerator(cnx)

        self.metagen = metagen

        self.logger = logging.getLogger('dataimport.massive_store')

        self.sql = cnx.system_sql

        self.schema = cnx.vreg.schema

        self.default_values = get_default_values(self.schema)

        self.get_next_eid = lambda g=self._get_eid_gen(eids_seq_range): next(g)

        self._source_dbhelper = cnx.repo.system_source.dbhelper

        self._dbh = PGHelper(cnx)

        self._data_entities = defaultdict(list)

        self._data_relations = defaultdict(list)

        self._initialized = {}

    def _get_eid_gen(self, eids_seq_range):

        """ Function getting the next eid. This is done by preselecting

        a given number of eids from the 'entities_id_seq', and then

        storing them"""

        while True:

            last_eid = self._cnx.repo.system_source.create_eid(self._cnx, eids_seq_range)

            for eid in range(last_eid - eids_seq_range + 1, last_eid + 1):

                yield eid

    # master/slaves specific API

    def master_init(self, commit=True):

        """Initialize database for massive insertion.

        This is expected to be called once, by the master store in master/slaves configuration.

        """

        assert not self.slave_mode

        if self not in self._initialized:

            self.sql('DROP TABLE IF EXISTS cwmassive_initialized')

            self.sql('CREATE TABLE cwmassive_initialized'

                     '(retype text, type varchar(128), uuid varchar(32))')

            self._initialized[self] = None

            if commit:

                self.commit()

    # SQL utilities #########################################################

    def _drop_metadata_constraints(self):

        """Drop constraints and indexes for the metadata tables.

        They will be recreated by the `finish` method.

        """

        rtypes = [rtype for rtype in self.metagen.meta_relations

                  if not self.schema.rschema(rtype).final]

        rtypes += ('is_instance_of', 'is', 'cw_source')

        for rtype in rtypes:

            self._dbh.drop_constraints(rtype + '_relation')

            self._dbh.drop_indexes(rtype + '_relation')

        # don't drop constraints for the entities table, the only one is the primary key's index on

        # eid and we want to keep it

        self._dbh.drop_indexes('entities')

    def restart_eid_sequence(self, start_eid):

        self.sql(self._cnx.repo.system_source.dbhelper.sql_restart_numrange(

            'entities_id_seq', initial_value=start_eid))

        self._cnx.commit()

    # store api ################################################################

    def prepare_insert_entity(self, etype, **data):

        """Given an entity type, attributes and inlined relations, returns the inserted entity's

        eid.

        """

        if etype not in self._initialized:

            if not self.slave_mode:

                self.master_init(commit=False)

            tablename = 'cw_%s' % etype.lower()

            tmp_tablename = '%s_%s' % (tablename, self.uuid)

            self.sql("INSERT INTO cwmassive_initialized VALUES (%(e)s, 'etype', %(uuid)s)",

                     {'e': etype, 'uuid': self.uuid})

            attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])

            self.sql('CREATE TABLE %s(%s);' % (tmp_tablename,

                                               ', '.join('cw_%s %s' % (column, sqltype)

                                                         for column, sqltype in attr_defs)))

            self._initialized[etype] = [attr for attr, _ in attr_defs]

        if 'eid' not in data:

            # If eid is not given and the eids sequence is set, use the value from the sequence

            eid = self.get_next_eid()

            data['eid'] = eid

        self._data_entities[etype].append(data)

        return data['eid']

    def prepare_insert_relation(self, eid_from, rtype, eid_to, **kwargs):

        """Insert into the database a  relation ``rtype`` between entities with eids ``eid_from``

        and ``eid_to``.

        Relation must not be inlined.

        """

        if rtype not in self._initialized:

            if not self.slave_mode:

                self.master_init(commit=False)

            assert not self._cnx.vreg.schema.rschema(rtype).inlined

            self._initialized[rtype] = None

            tablename = '%s_relation' % rtype.lower()

            tmp_tablename = '%s_%s' % (tablename, self.uuid)

            self.sql("INSERT INTO cwmassive_initialized VALUES (%(r)s, 'rtype', %(uuid)s)",

                     {'r': rtype, 'uuid': self.uuid})

            self.sql('CREATE TABLE %s(eid_from integer, eid_to integer)' % tmp_tablename)

        self._data_relations[rtype].append({'eid_from': eid_from, 'eid_to': eid_to})

    def flush(self):

        """Flush the data"""

        self.flush_entities()

        self.flush_relations()

    def finish(self):

        """Remove temporary tables and columns."""

        try:

            self._finish()

            self._cnx.commit()

        except Exception:

            self._cnx.rollback()

            raise

        finally:

            # delete the meta data table

            self.sql('DROP TABLE IF EXISTS cwmassive_initialized')

            self.commit()

    def _finish(self):

        """Remove temporary tables and columns."""

        assert not self.slave_mode, 'finish method should only be called by the master store'

        self.logger.info("Start cleaning")

        # Get all the initialized etypes/rtypes

        if self._dbh.table_exists('cwmassive_initialized'):

            cu = self.sql('SELECT retype, type, uuid FROM cwmassive_initialized')

            entities = defaultdict(list)

            relations = defaultdict(list)

            for retype, _type, uuid in cu.fetchall():

                if _type == 'rtype':

                    relations[retype].append(uuid)

                else:  # _type = 'etype'

                    entities[retype].append(uuid)

            # if there is some entities to insert, delete constraint on metadata tables once for all

            if entities:

                self._drop_metadata_constraints()

            # get back entity data from the temporary tables

            for etype, uuids in entities.items():

                tablename = 'cw_%s' % etype.lower()

                attr_defs = eschema_sql_def(self._source_dbhelper, self.schema[etype])

                columns = ','.join('cw_%s' % attr for attr, _ in attr_defs)

                self._dbh.drop_constraints(tablename)

                self._dbh.drop_indexes(tablename)

                for uuid in uuids:

                    tmp_tablename = '%s_%s' % (tablename, uuid)

                    self.sql('INSERT INTO %(table)s(%(columns)s) '

                             'SELECT %(columns)s FROM %(tmp_table)s'

                             % {'table': tablename, 'tmp_table': tmp_tablename,

                                'columns': columns})

                    self._insert_etype_metadata(etype, tmp_tablename)

                    self._tmp_data_cleanup(tmp_tablename, etype, uuid)

            # get back relation data from the temporary tables

            for rtype, uuids in relations.items():

                tablename = '%s_relation' % rtype.lower()

                self._dbh.drop_constraints(tablename)

                self._dbh.drop_indexes(tablename)

                for uuid in uuids:

                    tmp_tablename = '%s_%s' % (tablename, uuid)

                    self.fill_relation_table(tablename, tmp_tablename)

                    self._tmp_data_cleanup(tmp_tablename, rtype, uuid)

        # restore all deleted indexes and constraints

        self._dbh.restore_indexes_and_constraints()

    def _insert_etype_metadata(self, etype, tmp_tablename):

        """Massive insertion of meta data for `etype`, with new entities in `tmp_tablename`.

        """

        # insert standard metadata relations

        for rtype, eid in self.metagen.base_etype_rels(etype).items():

            self.fill_meta_relation_table(tmp_tablename, rtype, eid)

        # insert cw_source, is and is_instance_of relations (normally handled by the system source)

        self.fill_meta_relation_table(tmp_tablename, 'cw_source', self.metagen.source.eid)

        eschema = self.schema[etype]

        self.fill_meta_relation_table(tmp_tablename, 'is', eschema.eid)

        for parent_eschema in chain(eschema.ancestors(), [eschema]):

            self.fill_meta_relation_table(tmp_tablename, 'is_instance_of', parent_eschema.eid)

        self.fill_entities_table(etype, tmp_tablename)

    def fill_entities_table(self, etype, tmp_tablename):

        # finally insert records into the entities table

        self.sql("INSERT INTO entities(eid, type) "

                 "SELECT cw_eid, '%s' FROM %s "

                 "WHERE NOT EXISTS (SELECT 1 FROM entities WHERE eid=cw_eid)"

                 % (etype, tmp_tablename))

    def fill_relation_table(self, tablename, tmp_tablename):

        # XXX no index on the original relation table, EXISTS subquery may be sloooow

        self.sql('INSERT INTO %(table)s(eid_from, eid_to) SELECT DISTINCT '

                 'T.eid_from, T.eid_to FROM %(tmp_table)s AS T '

                 'WHERE NOT EXISTS (SELECT 1 FROM %(table)s AS TT WHERE '

                 'TT.eid_from=T.eid_from AND TT.eid_to=T.eid_to);'

                 % {'table': tablename, 'tmp_table': tmp_tablename})

    def fill_meta_relation_table(self, tmp_tablename, rtype, eid_to):

        self.sql("INSERT INTO %s_relation(eid_from, eid_to) SELECT cw_eid, %s FROM %s "

                 "WHERE NOT EXISTS (SELECT 1 FROM entities WHERE eid=cw_eid)"

                 % (rtype, eid_to, tmp_tablename))

    def _tmp_data_cleanup(self, tmp_tablename, ertype, uuid):

        """Drop temporary relation table and record from cwmassive_initialized."""

        self.sql('DROP TABLE %(tmp_table)s' % {'tmp_table': tmp_tablename})

        self.sql('DELETE FROM cwmassive_initialized '

                 'WHERE retype = %(rtype)s AND uuid = %(uuid)s',

                 {'rtype': ertype, 'uuid': uuid})

    # FLUSH #################################################################

    def flush_relations(self):

        """Flush the relations data from in-memory structures to a temporary table."""

        for rtype, data in self._data_relations.items():

            if not data:

                # There is no data for these etype for this flush round.

                continue

            buf = pgstore._create_copyfrom_buffer(data, ('eid_from', 'eid_to'))

            cursor = self._cnx.cnxset.cu

            tablename = '%s_relation' % rtype.lower()

            tmp_tablename = '%s_%s' % (tablename, self.uuid)

            cursor.copy_from(buf, tmp_tablename, null='NULL', columns=('eid_from', 'eid_to'))

            # Clear data cache

            self._data_relations[rtype] = []

    def flush_entities(self):

        """Flush the entities data from in-memory structures to a temporary table."""

        metagen = self.metagen

        for etype, data in self._data_entities.items():

            if not data:

                # There is no data for these etype for this flush round.

                continue

            attrs = self._initialized[etype]

            _base_data = dict.fromkeys(attrs)

            _base_data.update(self.default_values[etype])

            _base_data.update(metagen.base_etype_attrs(etype))

            _data = []

            for d in data:

                # do this first on `d`, because it won't fill keys associated to None as provided by

                # `_base_data`

                metagen.init_entity_attrs(etype, d['eid'], d)

                # XXX warn/raise if there is some key not in attrs?

                _d = _base_data.copy()

                _d.update(d)

                _data.append(_d)

            buf = pgstore._create_copyfrom_buffer(_data, attrs)

            tablename = 'cw_%s' % etype.lower()

            tmp_tablename = '%s_%s' % (tablename, self.uuid)

            columns = ['cw_%s' % attr for attr in attrs]

            cursor = self._cnx.cnxset.cu

            cursor.copy_from(buf, tmp_tablename, null='NULL', columns=columns)

            # Clear data cache

            self._data_entities[etype] = []

def get_default_values(schema):

    """analyzes yams ``schema`` and returns the list of default values.

    The returned value is a dictionary mapping entity types to a

    sub-dictionnaries mapping attribute names -> default values.

    """

    default_values = {}

    # iterates on all entity types

    for eschema in schema.entities():

        # for each entity type, iterates on attribute definitions

        default_values[eschema.type] = eschema_constraints = {}

        for rschema, _ in eschema.attribute_definitions():

            # for each attribute, if a size constraint is found,

            # append it to the size constraint list

            if eschema.default(rschema.type) is not None:

                eschema_constraints[rschema.type] = eschema.default(rschema.type)

    return default_values

class PGHelper(object):

    """This class provides some helper methods to manipulate a postgres database metadata (index and

    constraints).

    """

    def __init__(self, cnx):

        self.sql = cnx.system_sql

        # Deals with pg schema, see #3216686

        pg_schema = cnx.repo.config.system_source_config.get('db-namespace') or 'public'

        self.pg_schema = pg_schema

    def drop_indexes(self, tablename):

        """Drop indexes and constraints, storing them in a table for later restore."""

        # Create a table to save the constraints, it allows reloading even after crash

        self.sql('CREATE TABLE IF NOT EXISTS cwmassive_constraints(sql TEXT, insert_order SERIAL)')

        indexes = self.table_indexes(tablename)

        for name, query in indexes.items():

            self.sql('INSERT INTO cwmassive_constraints(sql) VALUES (%(sql)s)', {'sql': query})

            self.sql('DROP INDEX %s' % name)

    def drop_constraints(self, tablename):

        self.sql('CREATE TABLE IF NOT EXISTS cwmassive_constraints(sql TEXT, insert_order SERIAL)')

        constraints = self.table_constraints(tablename)

        for name, query in constraints.items():