"""plan execution of rql queries on multiple sources

the best way to understand what are we trying to acheive here is to read

the unit-tests in unittest_querier_planner.py

Split and execution specifications

~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

For a system source and a ldap user source (only EUser and its attributes

is supported, no group or such):

:EUser X:

1. fetch EUser X from both sources and return concatenation of results

:EUser X WHERE X in_group G, G name 'users':

* catch 1

  1. fetch EUser X from both sources, store concatenation of results

     into a temporary table

  2. return the result of TMP X WHERE X in_group G, G name 'users' from

     the system source

* catch 2

  1. return the result of EUser X WHERE X in_group G, G name 'users'

     from system source, that's enough (optimization of the sql querier

     will avoid join on EUser, so we will directly get local eids)

:EUser X,L WHERE X in_group G, X login L, G name 'users':

1. fetch Any X,L WHERE X is EUser, X login L from both sources, store

   concatenation of results into a temporary table

2. return the result of Any X, L WHERE X is TMP, X login LX in_group G,

   G name 'users' from the system source

:Any X WHERE X owned_by Y:

* catch 1

  1. fetch EUser X from both sources, store concatenation of results

     into a temporary table

  2. return the result of Any X WHERE X owned_by Y, Y is TMP from

     the system source

* catch 2

  1. return the result of Any X WHERE X owned_by Y

     from system source, that's enough (optimization of the sql querier

     will avoid join on EUser, so we will directly get local eids)

:organization: Logilab

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

"""

__docformat__ = "restructuredtext en"

from itertools import imap, ifilterfalse

from logilab.common.compat import any

from logilab.common.decorators import cached

from rql.stmts import Union, Select

from cubicweb import server

from cubicweb.utils import make_uid

from cubicweb.server.utils import cleanup_solutions

from cubicweb.server.ssplanner import SSPlanner, OneFetchStep, add_types_restriction

from cubicweb.server.mssteps import *

from cubicweb.server.sources import AbstractSource

Variable._ms_table_key = lambda x: x.name

Relation._ms_table_key = lambda x: x.r_type

# str() Constant.value to ensure generated table name won't be unicode

Constant._ms_table_key = lambda x: str(x.value)

AbstractSource.dont_cross_relations = ()

AbstractSource.cross_relations = ()

def allequals(solutions):

    """return true if all solutions are identical"""

    sol = solutions.next()

    for sol_ in solutions:

        if sol_ != sol:

            return False

    return True

def need_aggr_step(select, sources, stepdefs=None):

    """return True if a temporary table is necessary to store some partial

    results to execute the given query

    """

    if len(sources) == 1:

        # can do everything at once with a single source

        return False

    if select.orderby or select.groupby or select.has_aggregat:

        # if more than one source, we need a temp table to deal with sort /

        # groups / aggregat if :

        # * the rqlst won't be splitted (in the other case the last query

        #   using partial temporary table can do sort/groups/aggregat without

        #   the need for a later AggrStep)

        # * the rqlst is splitted in multiple steps and there are more than one

        #   final step

        if stepdefs is None:

            return True

        has_one_final = False

        fstepsolindices = set()

        for stepdef in stepdefs:

            if stepdef[-1]:

                if has_one_final or frozenset(stepdef[2]) != fstepsolindices:

                    return True

                has_one_final = True

            else:

                fstepsolindices.update(stepdef[2])

    return False

def copy_node(newroot, node, subparts=()):

    newnode = node.__class__(*node.initargs(newroot))

    for part in subparts:

        newnode.append(part)

    return newnode

def same_scope(var):

    """return true if the variable is always used in the same scope"""

    try:

        return var.stinfo['samescope']

    except KeyError:

        for rel in var.stinfo['relations']:

            if not rel.scope is var.scope:

                var.stinfo['samescope'] = False

                return False

        var.stinfo['samescope'] = True

        return True

def select_group_sort(select): # XXX something similar done in rql2sql

    # add variables used in groups and sort terms to the selection

    # if necessary

    if select.groupby:

        for vref in select.groupby:

            if not vref in select.selection:

                select.append_selected(vref.copy(select))

    for sortterm in select.orderby:

        for vref in sortterm.iget_nodes(VariableRef):

            if not vref in select.get_selected_variables():

                # we can't directly insert sortterm.term because it references

                # a variable of the select before the copy.

                # XXX if constant term are used to define sort, their value

                # may necessite a decay

                select.append_selected(vref.copy(select))

                if select.groupby and not vref in select.groupby:

                    select.add_group_var(vref.copy(select))

class PartPlanInformation(object):

    """regroups necessary information to execute some part of a "global" rql

    query ("global" means as received by the querier, which may result in

    several internal queries, e.g. parts, due to security insertions)

    it exposes as well some methods helping in executing this part on a

    multi-sources repository, modifying its internal structure during the

    process

    :attr solutions: a list of mappings (varname -> vartype)

    :attr sourcesvars:

      a dictionnary telling for each source which variable/solution are

      supported, of the form {source : {varname: [solution index, ]}}

    """

    def __init__(self, plan, rqlst, rqlhelper=None):

        self.needsplit = False

        self.temptable = None

        self.finaltable = None

        self.plan = plan

        self.rqlst = rqlst

        self._session = plan.session

        self._solutions = rqlst.solutions

        self._solindices = range(len(self._solutions))

        # source : {var: [solution index, ]}

        self.sourcesvars = self._sourcesvars = {}

        # source : {relation: set(child variable and constant)}

        self._crossrelations = {}

        # dictionnary of variables which are linked to each other using a non

        # final relation which is supported by multiple sources

        self._linkedvars = {}

        self._crosslinkedvars = {}

        # processing

        self._compute_sourcesvars()

        self._remove_invalid_sources()

        self._compute_needsplit()

        self.sourcesvars = {}

        for k, v in self._sourcesvars.iteritems():

            self.sourcesvars[k] = {}

            for k2, v2 in v.iteritems():

                self.sourcesvars[k][k2] = v2.copy()

        self._inputmaps = {}

        if rqlhelper is not None: # else test

            self._insert_identity_variable = rqlhelper._annotator.rewrite_shared_optional

    def copy_solutions(self, solindices):

        return [self._solutions[solidx].copy() for solidx in solindices]

    @property

    @cached

    def part_sources(self):

        if self._sourcesvars:

            return tuple(sorted(self._sourcesvars))

        return (self._session.repo.system_source,)

    @property

    @cached

    def _sys_source_set(self):

        return frozenset((self._session.repo.system_source, solindex)

                         for solindex in self._solindices)        

    @cached

    def _norel_support_set(self, relation):

        """return a set of (source, solindex) where source doesn't support the

        relation

        """

        return frozenset((source, solidx) for source in self._session.repo.sources

                         for solidx in self._solindices

                         if not ((source.support_relation(relation.r_type) and

                                  not self.crossed_relation(source, relation))

                                 or relation.r_type in source.dont_cross_relations))

    def _compute_sourcesvars(self):

        """compute for each variable/solution in the rqlst which sources support

        them

        """

        repo = self._session.repo

        eschema = repo.schema.eschema

        sourcesvars = self._sourcesvars

        # find for each source which variable/solution are supported

        for varname, varobj in self.rqlst.defined_vars.items():

            # if variable has an eid specified, we can get its source directly

            # NOTE: use uidrels and not constnode to deal with "X eid IN(1,2,3,4)"

            if varobj.stinfo['uidrels']:

                vrels = varobj.stinfo['relations'] - varobj.stinfo['uidrels']

                for rel in varobj.stinfo['uidrels']:

                    if  rel.neged(strict=True) or rel.operator() != '=':

                        continue

                    for const in rel.children[1].get_nodes(Constant):

                        eid = const.eval(self.plan.args)

                        source = self._session.source_from_eid(eid)

                        if vrels and not any(source.support_relation(r.r_type)

                                             for r in vrels):

                            self._set_source_for_var(repo.system_source, varobj)

                        else:

                            self._set_source_for_var(source, varobj)

                continue

            rels = varobj.stinfo['relations']

            if not rels and not varobj.stinfo['typerels']:

                # (rare) case where the variable has no type specified nor

                # relation accessed ex. "Any MAX(X)"

                self._set_source_for_var(repo.system_source, varobj)

                continue

            for i, sol in enumerate(self._solutions):

                vartype = sol[varname]

                # skip final variable

                if eschema(vartype).is_final():

                    break

                for source in repo.sources:

                    if source.support_entity(vartype):

                        # the source support the entity type, though we will

                        # actually have to fetch from it only if

                        # * the variable isn't invariant

                        # * at least one supported relation specified

                        if not varobj._q_invariant or \

                               any(imap(source.support_relation,

                                        (r.r_type for r in rels if r.r_type != 'eid'))):

                            sourcesvars.setdefault(source, {}).setdefault(varobj, set()).add(i)

                        # if variable is not invariant and is used by a relation

                        # not supported by this source, we'll have to split the

                        # query

                        if not varobj._q_invariant and any(ifilterfalse(

                            source.support_relation, (r.r_type for r in rels))):

                            self.needsplit = True               

    def _handle_cross_relation(self, rel, relsources, vsources):

        crossvars = None

        for source in relsources:

            if rel.r_type in source.cross_relations:

                crossvars = set(x.variable for x in rel.get_nodes(VariableRef))

                crossvars.update(frozenset(x for x in rel.get_nodes(Constant)))

                assert len(crossvars) == 2

                ssource = self._session.repo.system_source

                needsplit = True

                flag = 0

                for v in crossvars:

                    if isinstance(v, Constant):

                        allsols = set(self._solindices)

                        try:

                            self._sourcesvars[ssource][v] = allsols

                        except KeyError:

                            self._sourcesvars[ssource] = {v: allsols}

                    if len(vsources[v]) == 1:

                        if iter(vsources[v]).next()[0].uri == 'system':

                            flag = 1

                            for ov in crossvars:

                                if ov is not v and (isinstance(ov, Constant) or ov._q_invariant):

                                    ssset = frozenset((ssource,))

                                    self._remove_sources(ov, vsources[ov] - ssset)

                        else:

                            for ov in crossvars:

                                if ov is not v and (isinstance(ov, Constant) or ov._q_invariant):

                                    needsplit = False

                                    break

                            else:

                                continue

                        if not rel.neged(strict=True):

                            break

                else:

                    self._crossrelations.setdefault(source, {})[rel] = crossvars

                    if not flag:

                        self._sourcesvars.setdefault(source, {})[rel] = set(self._solindices)

                    self._sourcesvars.setdefault(ssource, {})[rel] = set(self._solindices)

                    if needsplit:

                        self.needsplit = True

        return crossvars is None

    def _remove_invalid_sources(self):

        """removes invalid sources from `sourcesvars` member according to

        traversed relations and their properties (which sources support them,

        can they cross sources, etc...)

        """

        repo = self._session.repo

        rschema = repo.schema.rschema

        vsources = {}

        for rel in self.rqlst.iget_nodes(Relation):

            # process non final relations only

            # note: don't try to get schema for 'is' relation (not available

            # during bootstrap)

            if not rel.is_types_restriction() and not rschema(rel.r_type).is_final():

                # nothing to do if relation is not supported by multiple sources

                # or if some source has it listed in its cross_relations

                # attribute

                #

                # XXX code below don't deal if some source allow relation

                #     crossing but not another one

                relsources = repo.rel_type_sources(rel.r_type)

                crossvars = None

                if len(relsources) < 2:

                    # filter out sources being there because they have this

                    # relation in their dont_cross_relations attribute

                    relsources = [source for source in relsources

                                  if source.support_relation(rel.r_type)]

                    if relsources:

                        # this means the relation is using a variable inlined as

                        # a constant and another unsupported variable, in which

                        # case we put the relation in sourcesvars

                        self._sourcesvars.setdefault(relsources[0], {})[rel] = set(self._solindices)

                    continue

                lhs, rhs = rel.get_variable_parts()

                lhsv, rhsv = getattr(lhs, 'variable', lhs), getattr(rhs, 'variable', rhs)

                # update dictionnary of sources supporting lhs and rhs vars

                if not lhsv in vsources:

                    vsources[lhsv] = self._term_sources(lhs)

                if not rhsv in vsources:

                    vsources[rhsv] = self._term_sources(rhs)

                if self._handle_cross_relation(rel, relsources, vsources):

                    self._linkedvars.setdefault(lhsv, set()).add((rhsv, rel))

                    self._linkedvars.setdefault(rhsv, set()).add((lhsv, rel))

                else:

                    self._crosslinkedvars.setdefault(lhsv, set()).add((rhsv, rel))

                    self._crosslinkedvars.setdefault(rhsv, set()).add((lhsv, rel))

        for term in self._linkedvars:

            self._remove_sources_until_stable(term, vsources)

        if len(self._sourcesvars) > 1 and hasattr(self.plan.rqlst, 'main_relations'):

            # the querier doesn't annotate write queries, need to do it here

            self.plan.annotate_rqlst()

            # insert/update/delete queries, we may get extra information from

            # the main relation (eg relations to the left of the WHERE

            if self.plan.rqlst.TYPE == 'insert':

                inserted = dict((vref.variable, etype)

                                for etype, vref in self.plan.rqlst.main_variables)

            else:

                inserted = {}

            for rel in self.plan.rqlst.main_relations:

                if not rschema(rel.r_type).is_final():

                    # nothing to do if relation is not supported by multiple sources

                    if len(repo.rel_type_sources(rel.r_type)) < 2:

                        continue

                    lhs, rhs = rel.get_variable_parts()

                    try:

                        lhsv = self._extern_term(lhs, vsources, inserted)

                        rhsv = self._extern_term(rhs, vsources, inserted)

                    except KeyError, ex:

                        continue

                    norelsup = self._norel_support_set(rel)

                    self._remove_var_sources(lhsv, norelsup, rhsv, vsources)

                    self._remove_var_sources(rhsv, norelsup, lhsv, vsources)

        # cleanup linked var

        for var, linkedrelsinfo in self._linkedvars.iteritems():