"""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_msplanner.py

What you need to know

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

1. The system source is expected  to support every entity and relation types

2. Given "X relation Y":

   * if relation, X and Y types are supported by the external source, we suppose

     by default that X and Y should both come from the same source as the

     relation. You can specify otherwise by adding relation into the

     "cross_relations" set in the source's mapping file and it that case, we'll

     consider that we can also find in the system source some relation between

     X and Y coming from different sources.

   * if "relation" isn't supported by the external source but X or Y

     types (or both) are, we suppose by default that can find in the system

     source some relation where X and/or Y come from the external source. You

     can specify otherwise by adding relation into the "dont_cross_relations"

     set in the source's mapping file and it that case, we'll consider that we

     can only find in the system source some relation between X and Y coming

     the system source.

Implementation

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

XXX explain algorithm

Exemples of multi-sources query execution

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

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

:copyright: 2003-2009 LOGILAB S.A. (Paris, FRANCE), all rights reserved.

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

"""

__docformat__ = "restructuredtext en"

from copy import deepcopy

from itertools import imap, ifilterfalse

from logilab.common.compat import any

from logilab.common.decorators import cached

from rql.stmts import Union, Select

from rql.nodes import VariableRef, Comparison, Relation, Constant, Exists, Variable

from cubicweb import server

from cubicweb.common.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 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 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))

def allequals(solutions):

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

    sol = solutions.next()

    noconstsol = None

    for sol_ in solutions:

        if sol_ != sol:

            return False

    return True

# XXX move functions below to rql ##############################################

def is_ancestor(n1, n2):

    p = n1.parent

    while p is not None:

        if p is n2:

            return True

        p = p.parent

    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

################################################################################

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). Actually

    a PPI is created for each subquery and for each query in a union.

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

    multi-sources repository, modifying its internal structure during the

    process.

    :attr plan:

      the execution plan

    :attr rqlst:

      the original rql syntax tree handled by this part

    :attr needsplit:

      bool telling if the query has to be split into multiple steps for

      execution or if it can be executed at once

    :attr temptable:

      a SQL temporary table name or None, if necessary to handle aggregate /

      sorting for this part of the query

    :attr finaltable:

      a SQL table name or None, if results for this part of the query should be

      written into a temporary table (usually shared by multiple PPI)

    :attr sourcesterms:

      a dictionary {source : {term: set([solution index, ])}} telling for each

      source which terms are supported for which solutions. A "term" may be

      either a rql Variable, Constant or Relation node.

    """

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

        self.plan = plan

        self.rqlst = rqlst

        self.needsplit = False

        self.temptable = None

        self.finaltable = None

        self._schema = plan.schema

        self._session = plan.session

        self._repo = self._session.repo

        self._solutions = rqlst.solutions

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

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

        self.sourcesterms = self._sourcesterms = {}

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

        self._crossrelations = {}

        # dictionary of variables and constants which are linked to each other

        # using a non final relation supported by multiple sources (crossed or

        # not).

        self._linkedterms = {}

        # processing

        termssources = self._compute_sourcesterms()

        self._remove_invalid_sources(termssources)

        self._compute_needsplit()

        # after initialisation, .sourcesterms contains the same thing as

        # ._sourcesterms though during plan construction, ._sourcesterms will

        # be modified while .sourcesterms will be kept unmodified

        self.sourcesterms = {}

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

            self.sourcesterms[k] = {}

            for k2, v2 in v.iteritems():

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

        # cleanup linked var

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

            self._linkedterms[var] = frozenset(x[0] for x in linkedrelsinfo)

        # map output of a step to input of a following step

        self._inputmaps = {}

        # record input map conflicts to resolve them on final step generation

        self._conflicts = []

        if rqlhelper is not None: # else test

            self._insert_identity_variable = rqlhelper._annotator.rewrite_shared_optional

        if server.DEBUG:

            print 'sourcesterms:'

            for source, terms in self.sourcesterms.items():

                print source, terms

    def copy_solutions(self, solindices):

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

    @property

    @cached

    def part_sources(self):

        if self._sourcesterms:

            return tuple(sorted(self._sourcesterms))

        return (self._repo.system_source,)

    @property

    @cached

    def _sys_source_set(self):

        return frozenset((self._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._repo.sources

                         for solidx in self._solindices

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

                                 or relation.r_type in source.dont_cross_relations))

    def _compute_sourcesterms(self):

        """compute for each term (variable, rewritten constant, relation) and

        for each solution in the rqlst which sources support them

        """

        repo = self._repo

        eschema = self._schema.eschema

        sourcesterms = self._sourcesterms

        # 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_term(repo.system_source, varobj)

                        else:

                            self._set_source_for_term(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_term(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'))):

                            sourcesterms.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               

        # add source for rewritten constants to sourcesterms

        for vconsts in self.rqlst.stinfo['rewritten'].itervalues():

            const = vconsts[0]

            source = self._session.source_from_eid(const.eval(self.plan.args))

            if source is self._repo.system_source:

                for const in vconsts:

                    self._set_source_for_term(source, const)

            elif source in self._sourcesterms:

                source_scopes = frozenset(t.scope for t in self._sourcesterms[source])

                for const in vconsts:

                    if const.scope in source_scopes:

                        self._set_source_for_term(source, const)

        # add source for relations

        rschema = self._schema.rschema

        termssources = {}

        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)

                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 sourcesterms

                        self._sourcesterms.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 dictionary of sources supporting lhs and rhs vars

                if not lhsv in termssources:

                    termssources[lhsv] = self._term_sources(lhs)

                if not rhsv in termssources:

                    termssources[rhsv] = self._term_sources(rhs)

                self._handle_cross_relation(rel, relsources, termssources)

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

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

        return termssources

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

        for source in relsources:

            if rel.r_type in source.cross_relations:

                ssource = self._repo.system_source

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

                for const in rel.get_nodes(Constant):