"""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

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

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

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 : {varname: [solution index, ]}

        self._sourcesvars = {}

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

        # final relation which is supported by multiple sources

        self._linkedvars = {}

        # processing

        self._compute_sourcesvars()

        self._remove_invalid_sources()

        #if server.DEBUG:

        #    print 'planner sources vars', self._sourcesvars

        self._compute_needsplit()

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

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

                                 or rtype 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 _remove_invalid_sources(self):

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

        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

                relsources = [source for source in repo.sources

                              if source.support_relation(rel.r_type)

                              or rel.r_type in source.dont_cross_relations]

                if len(relsources) < 2:

                    if relsources:# and not relsources[0] in self._sourcesvars:

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

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

                self._linkedvars.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

                    relsources = [source for source in repo.sources

                                  if source.support_relation(rel.r_type)

                                  or rel.r_type in source.dont_cross_relations]

                    if len(relsources) < 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.r_type)

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

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

        # if there are other sources than the system source, consider simplified

        # variables'source

        if self._sourcesvars and self._sourcesvars.keys() != [self._session.repo.system_source]:

            # add source for rewritten constants to sourcesvars

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

                const = vconsts[0]

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

                source = self._session.source_from_eid(eid)

                if source is self._session.repo.system_source:

                    for const in vconsts:

                        self._set_source_for_var(source, const)

                elif source in self._sourcesvars:

                    source_scopes = frozenset(v.scope for v in self._sourcesvars[source])

                    for const in vconsts:

                        if const.scope in source_scopes:

                            self._set_source_for_var(source, const)

    def _extern_term(self, term, vsources, inserted):

        var = term.variable

        if var.stinfo['constnode']:

            termv = var.stinfo['constnode']

            vsources[termv] = self._term_sources(termv)

        elif var in inserted:

            termv = var

            source = self._session.repo.locate_etype_source(inserted[var])

            vsources[termv] = set((source, solindex) for solindex in self._solindices)

        else:

            termv = self.rqlst.defined_vars[var.name]

            if not termv in vsources:

                vsources[termv] = self._term_sources(termv)

        return termv

    def _remove_sources_until_stable(self, var, vsources):

        for ovar, rel in self._linkedvars.get(var, ()):

            if not var.scope is ovar.scope and rel.scope.neged(strict=True):

                # can't get information from relation inside a NOT exists

                # where variables don't belong to the same scope

                continue

            if rel.neged(strict=True):

                # neged relation doesn't allow to infer variable sources

                continue

            norelsup = self._norel_support_set(rel.r_type)

            # compute invalid sources for variables and remove them

            self._remove_var_sources(var, norelsup, ovar, vsources)

            self._remove_var_sources(ovar, norelsup, var, vsources)

    def _remove_var_sources(self, var, norelsup, ovar, vsources):

        """remove invalid sources for var according to ovar's sources and the

        relation between those two variables. 

        """

        varsources = vsources[var]

        invalid_sources = varsources - (vsources[ovar] | norelsup)

        if invalid_sources:

            self._remove_sources(var, invalid_sources)

            varsources -= invalid_sources

            self._remove_sources_until_stable(var, vsources)

    def _compute_needsplit(self):

        """tell according to sourcesvars if the rqlst has to be splitted for

        execution among multiple sources

        the execution has to be split if

        * a source support an entity (non invariant) but doesn't support a

          relation on it

        * a source support an entity which is accessed by an optional relation

        * there is more than one sources and either all sources'supported        

          variable/solutions are not equivalent or multiple variables have to

          be fetched from some source

        """

        # NOTE: < 2 since may be 0 on queries such as Any X WHERE X eid 2