"""Helper classes to execute RQL queries on a set of sources, performing

security checking and data aggregation.

:organization: Logilab

:copyright: 2001-2009 LOGILAB S.A. (Paris, FRANCE), license is LGPL v2.

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

:license: GNU Lesser General Public License, v2.1 - http://www.gnu.org/licenses

"""

__docformat__ = "restructuredtext en"

from itertools import repeat

from logilab.common.cache import Cache

from logilab.common.compat import any

from rql import RQLHelper, RQLSyntaxError

from rql.stmts import Union, Select

from rql.nodes import (Relation, VariableRef, Constant, SubQuery)

from cubicweb import Unauthorized, QueryError, UnknownEid, typed_eid

from cubicweb import server

from cubicweb.rset import ResultSet

from cubicweb.server.utils import cleanup_solutions

from cubicweb.server.rqlannotation import SQLGenAnnotator, set_qdata

from cubicweb.server.ssplanner import add_types_restriction

READ_ONLY_RTYPES = set(('eid', 'has_text', 'is', 'is_instance_of', 'identity'))

def empty_rset(session, rql, args, rqlst=None):

    """build an empty result set object"""

    return ResultSet([], rql, args, rqlst=rqlst)

def update_varmap(varmap, selected, table):

    """return a sql schema to store RQL query result"""

    for i, term in enumerate(selected):

        key = term.as_string()

        value = '%s.C%s' % (table, i)

        if varmap.get(key, value) != value:

            raise Exception('variable name conflict on %s' % key)

        varmap[key] = value

# permission utilities ########################################################

def var_kwargs(restriction, args):

    varkwargs = {}

    for rel in restriction.iget_nodes(Relation):

        cmp = rel.children[1]

        if rel.r_type == 'eid' and cmp.operator == '=' and \

               not rel.neged(strict=True) and \

               isinstance(cmp.children[0], Constant) and \

               cmp.children[0].type == 'Substitute':

            varkwargs[rel.children[0].name] = typed_eid(cmp.children[0].eval(args))

    return varkwargs

def check_no_password_selected(rqlst):

    """check that Password entities are not selected"""

    for solution in rqlst.solutions:

        if 'Password' in solution.itervalues():

            raise Unauthorized('Password selection is not allowed')

def check_read_access(schema, user, rqlst, solution):

    """check that the given user has credentials to access data read the

    query

    return a dict defining necessary local checks (due to use of rql expression

    in the schema), keys are variable names and values associated rql expression

    for the associated variable with the given solution

    """

    if rqlst.where is not None:

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

            # XXX has_text may have specific perm ?

            if rel.r_type in READ_ONLY_RTYPES:

                continue

            if not schema.rschema(rel.r_type).has_access(user, 'read'):

                raise Unauthorized('read', rel.r_type)

    localchecks = {}

    # iterate on defined_vars and not on solutions to ignore column aliases

    for varname in rqlst.defined_vars:

        etype = solution[varname]

        eschema = schema.eschema(etype)

        if not eschema.has_access(user, 'read'):

            erqlexprs = eschema.get_rqlexprs('read')

            if not erqlexprs:

                ex = Unauthorized('read', etype)

                ex.var = varname

                raise ex

            #assert len(erqlexprs) == 1

            localchecks[varname] = tuple(erqlexprs)

    return localchecks

def noinvariant_vars(restricted, select, nbtrees):

    # a variable can actually be invariant if it has not been restricted for

    # security reason or if security assertion hasn't modified the possible

    # solutions for the query

    if nbtrees != 1:

        for vname in restricted:

            try:

                yield select.defined_vars[vname]

            except KeyError:

                # this is an alias

                continue

    else:

        for vname in restricted:

            try:

                var = select.defined_vars[vname]

            except KeyError:

                # this is an alias

                continue

            if len(var.stinfo['possibletypes']) != 1:

                yield var

def _expand_selection(terms, selected, aliases, select, newselect):

    for term in terms:

        for vref in term.iget_nodes(VariableRef):

            if not vref.name in selected:

                select.append_selected(vref)

                colalias = newselect.get_variable(vref.name, len(aliases))

                aliases.append(VariableRef(colalias))

                selected.add(vref.name)

# Plans #######################################################################

class ExecutionPlan(object):

    """the execution model of a rql query, composed of querier steps"""

    def __init__(self, querier, rqlst, args, session):

        # original rql syntax tree

        self.rqlst = rqlst

        self.args = args or {}

        # session executing the query

        self.session = session

        # quick reference to the system source

        self.syssource = session.pool.source('system')

        # execution steps

        self.steps = []

        # index of temporary tables created during execution

        self.temp_tables = {}

        # various resource accesors

        self.querier = querier

        self.schema = querier.schema

        self.sqlannotate = querier.sqlgen_annotate

        self.rqlhelper = session.vreg.rqlhelper

    def annotate_rqlst(self):

        if not self.rqlst.annotated:

            self.rqlhelper.annotate(self.rqlst)

    def add_step(self, step):

        """add a step to the plan"""

        self.steps.append(step)

    def clean(self):

        """remove temporary tables"""

        self.syssource.clean_temp_data(self.session, self.temp_tables)

    def sqlexec(self, sql, args=None):

        return self.syssource.sqlexec(self.session, sql, args)

    def execute(self):

        """execute a plan and return resulting rows"""

        try:

            for step in self.steps:

                result = step.execute()

            # the latest executed step contains the full query result

            return result

        finally:

            self.clean()

    def init_temp_table(self, table, selected, sol):

        """initialize sql schema and variable map for a temporary table which

        will be used to store result for the given rqlst

        """

        try:

            outputmap, sqlschema, _ = self.temp_tables[table]

            update_varmap(outputmap, selected, table)

        except KeyError:

            sqlschema, outputmap = self.syssource.temp_table_def(selected, sol,

                                                                 table)

            self.temp_tables[table] = [outputmap, sqlschema, False]

        return outputmap

    def create_temp_table(self, table):

        """create a temporary table to store result for the given rqlst"""

        if not self.temp_tables[table][-1]:

            sqlschema = self.temp_tables[table][1]

            self.syssource.create_temp_table(self.session, table, sqlschema)

            self.temp_tables[table][-1] = True

    def preprocess(self, union, security=True):

        """insert security when necessary then annotate rql st for sql generation

        return rqlst to actually execute

        """

        noinvariant = set()

        if security and not self.session.is_super_session:

            self._insert_security(union, noinvariant)

        self.rqlhelper.simplify(union)

        self.sqlannotate(union)

        set_qdata(self.schema.rschema, union, noinvariant)

        if union.has_text_query:

            self.cache_key = None

    def _insert_security(self, union, noinvariant):

        rh = self.rqlhelper

        for select in union.children[:]:

            for subquery in select.with_:

                self._insert_security(subquery.query, noinvariant)

            localchecks, restricted = self._check_permissions(select)

            if any(localchecks):

                rewrite = self.session.rql_rewriter.rewrite

                nbtrees = len(localchecks)

                myunion = union

                # transform in subquery when len(localchecks)>1 and groups

                if nbtrees > 1 and (select.orderby or select.groupby or

                                    select.having or select.has_aggregat or

                                    select.limit or select.offset):

                    newselect = Select()

                    # only select variables in subqueries

                    origselection = select.selection

                    select.select_only_variables()

                    select.has_aggregat = False

                    # create subquery first so correct node are used on copy

                    # (eg ColumnAlias instead of Variable)

                    aliases = [VariableRef(newselect.get_variable(vref.name, i))

                               for i, vref in enumerate(select.selection)]

                    selected = set(vref.name for vref in aliases)

                    # now copy original selection and groups

                    for term in origselection:

                        newselect.append_selected(term.copy(newselect))

                    if select.orderby:

                        newselect.set_orderby([s.copy(newselect) for s in select.orderby])

                        _expand_selection(select.orderby, selected, aliases, select, newselect)

                        select.orderby = () # XXX dereference?

                    if select.groupby:

                        newselect.set_groupby([g.copy(newselect) for g in select.groupby])

                        _expand_selection(select.groupby, selected, aliases, select, newselect)

                        select.groupby = () # XXX dereference?

                    if select.having:

                        newselect.set_having([g.copy(newselect) for g in select.having])

                        _expand_selection(select.having, selected, aliases, select, newselect)

                        select.having = () # XXX dereference?

                    if select.limit:

                        newselect.limit = select.limit

                        select.limit = None

                    if select.offset:

                        newselect.offset = select.offset

                        select.offset = 0

                    myunion = Union()

                    newselect.set_with([SubQuery(aliases, myunion)], check=False)

                    solutions = [sol.copy() for sol in select.solutions]

                    cleanup_solutions(newselect, solutions)

                    newselect.set_possible_types(solutions)

                    # if some solutions doesn't need rewriting, insert original

                    # select as first union subquery

                    if () in localchecks:

                        myunion.append(select)

                    # we're done, replace original select by the new select with

                    # subqueries (more added in the loop below)

                    union.replace(select, newselect)

                elif not () in localchecks:

                    union.remove(select)

                for lcheckdef, lchecksolutions in localchecks.iteritems():

                    if not lcheckdef:

                        continue

                    myrqlst = select.copy(solutions=lchecksolutions)

                    myunion.append(myrqlst)

                    # in-place rewrite + annotation / simplification

                    lcheckdef = [((varmap, 'X'), rqlexprs)

                                 for varmap, rqlexprs in lcheckdef]

                    rewrite(myrqlst, lcheckdef, lchecksolutions, self.args)

                    noinvariant.update(noinvariant_vars(restricted, myrqlst, nbtrees))

                if () in localchecks:

                    select.set_possible_types(localchecks[()])

                    add_types_restriction(self.schema, select)

                    noinvariant.update(noinvariant_vars(restricted, select, nbtrees))

    def _check_permissions(self, rqlst):

        """return a dict defining "local checks", e.g. RQLExpression defined in

        the schema that should be inserted in the original query

        solutions where a variable has a type which the user can't definitly read

        are removed, else if the user may read it (eg if an rql expression is

        defined for the "read" permission of the related type), the local checks

        dict for the solution is updated

        return a dict with entries for each different local check necessary,

        with associated solutions as value. A local check is defined by a list

        of 2-uple, with variable name as first item and the necessary rql

        expression as second item for each variable which has to be checked.

        So solutions which don't require local checks will be associated to

        the empty tuple key.

        note: rqlst should not have been simplified at this point

        """

        assert not self.session.is_super_session

        user = self.session.user

        schema = self.schema

        msgs = []

        # dictionnary of variables restricted for security reason

        localchecks = {}

        if rqlst.where is not None:

            varkwargs = var_kwargs(rqlst.where, self.args)

            neweids = self.session.transaction_data.get('neweids', ())

        else:

            varkwargs = None

        restricted_vars = set()

        newsolutions = []

        for solution in rqlst.solutions:

            try:

                localcheck = check_read_access(schema, user, rqlst, solution)

            except Unauthorized, ex:

                msg = 'remove %s from solutions since %s has no %s access to %s'

                msg %= (solution, user.login, ex.args[0], ex.args[1])

                msgs.append(msg)

                LOGGER.info(msg)

            else:

                newsolutions.append(solution)

                if varkwargs:

                    # try to benefit of rqlexpr.check cache for entities which

                    # are specified by eid in query'args

                    for varname, eid in varkwargs.iteritems():

                        try:

                            rqlexprs = localcheck.pop(varname)

                        except KeyError:

                            continue

                        if eid in neweids:

                            continue

                        for rqlexpr in rqlexprs:

                            if rqlexpr.check(self.session, eid):

                                break

                        else:

                            raise Unauthorized()

                restricted_vars.update(localcheck)

                localchecks.setdefault(tuple(localcheck.iteritems()), []).append(solution)

        # raise Unautorized exception if the user can't access to any solution

        if not newsolutions:

            raise Unauthorized('\n'.join(msgs))

        rqlst.set_possible_types(newsolutions)

        return localchecks, restricted_vars

    def finalize(self, select, solutions, insertedvars):

        rqlst = Union()

        rqlst.append(select)

        for mainvarname, rschema, newvarname in insertedvars:

            nvartype = str(rschema.objects(solutions[0][mainvarname])[0])

            for sol in solutions:

                sol[newvarname] = nvartype

        select.clean_solutions(solutions)

        self.rqlhelper.annotate(rqlst)

        self.preprocess(rqlst, security=False)

        return rqlst

class InsertPlan(ExecutionPlan):

    """an execution model specific to the INSERT rql query

    """

    def __init__(self, querier, rqlst, args, session):

        ExecutionPlan.__init__(self, querier, rqlst, args, session)

        # save originaly selected variable, we may modify this

        # dictionary for substitution (query parameters)

        self.selected = rqlst.selection

        # list of new or updated entities definition (utils.Entity)

        self.e_defs = [[]]

        # list of new relation definition (3-uple (from_eid, r_type, to_eid)

        self.r_defs = []

        # indexes to track entity definitions bound to relation definitions

        self._r_subj_index = {}

        self._r_obj_index = {}

        self._expanded_r_defs = {}

    def relation_definitions(self, rqlst, to_build):

        """add constant values to entity def, mark variables to be selected

        """

        to_select = {}

        for relation in rqlst.main_relations:

            lhs, rhs = relation.get_variable_parts()

            rtype = relation.r_type

            if rtype in READ_ONLY_RTYPES:

                raise QueryError("can't assign to %s" % rtype)

            try:

                edef = to_build[str(lhs)]

            except KeyError:

                # lhs var is not to build, should be selected and added as an

                # object relation

                edef = to_build[str(rhs)]

                to_select.setdefault(edef, []).append((rtype, lhs, 1))

            else: