# copyright 2003-2013 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/>.

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

security checking and data aggregation.

"""

__docformat__ = "restructuredtext en"

from itertools import repeat

from logilab.common.compat import any

from rql import RQLSyntaxError, CoercionError

from rql.stmts import Union

from rql.nodes import ETYPE_PYOBJ_MAP, etype_from_pyobj, Relation, Exists, Not

from yams import BASE_TYPES

from cubicweb import ValidationError, Unauthorized, UnknownEid

from cubicweb import Binary, server

from cubicweb.rset import ResultSet

from cubicweb.utils import QueryCache, RepeatList

from cubicweb.server.rqlannotation import SQLGenAnnotator, set_qdata

from cubicweb.server.ssplanner import READ_ONLY_RTYPES, add_types_restriction

from cubicweb.server.edition import EditedEntity

from cubicweb.server.ssplanner import SSPlanner

ETYPE_PYOBJ_MAP[Binary] = 'Bytes'

def empty_rset(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: got %s / %s'

                            % (key, value, varmap))

        varmap[key] = value

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

def check_no_password_selected(rqlst):

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

    for solution in rqlst.solutions:

        for var, etype in solution.iteritems():

            if etype == 'Password':

                raise Unauthorized('Password selection is not allowed (%s)' % var)

def term_etype(cnx, term, solution, args):

    """return the entity type for the given term (a VariableRef or a Constant

    node)

    """

    try:

        return solution[term.name]

    except AttributeError:

        return cnx.entity_metas(term.eval(args))['type']

def check_read_access(cnx, rqlst, solution, args):

    """Check that the given user has credentials to access data read by the

    query and return a dict defining necessary "local checks" (i.e. rql

    expression in read permission defined in the schema) where no group grants

    him the permission.

    Returned dictionary's keys are variable names and values the rql expressions

    for this variable (with the given solution).

    """

    # use `term_etype` since we've to deal with rewritten constants here,

    # when used as an external source by another repository.

    # XXX what about local read security w/ those rewritten constants...

    DBG = (server.DEBUG & server.DBG_SEC) and 'read' in server._SECURITY_CAPS

    schema = cnx.repo.schema

    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

            rschema = schema.rschema(rel.r_type)

            if rschema.final:

                eschema = schema.eschema(term_etype(cnx, rel.children[0],

                                                    solution, args))

                rdef = eschema.rdef(rschema)

            else:

                rdef = rschema.rdef(term_etype(cnx, rel.children[0],

                                               solution, args),

                                    term_etype(cnx, rel.children[1].children[0],

                                               solution, args))

            if not cnx.user.matching_groups(rdef.get_groups('read')):

                if DBG:

                    print ('check_read_access: %s %s does not match %s' %

                           (rdef, cnx.user.groups, rdef.get_groups('read')))

                # XXX rqlexpr not allowed

                raise Unauthorized('read', rel.r_type)

            if DBG:

                print ('check_read_access: %s %s matches %s' %

                       (rdef, cnx.user.groups, rdef.get_groups('read')))

    localchecks = {}

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

    for varname in rqlst.defined_vars:

        eschema = schema.eschema(solution[varname])

        if eschema.final:

            continue

        if not cnx.user.matching_groups(eschema.get_groups('read')):

            erqlexprs = eschema.get_rqlexprs('read')

            if not erqlexprs:

                ex = Unauthorized('read', solution[varname])

                ex.var = varname

                if DBG:

                    print ('check_read_access: %s %s %s %s' %

                           (varname, eschema, cnx.user.groups, eschema.get_groups('read')))

                raise ex

            # don't insert security on variable only referenced by 'NOT X relation Y' or

            # 'NOT EXISTS(X relation Y)'

            varinfo = rqlst.defined_vars[varname].stinfo

            if varinfo['selected'] or (

                len([r for r in varinfo['relations']

                     if (not schema.rschema(r.r_type).final

                         and ((isinstance(r.parent, Exists) and r.parent.neged(strict=True))

                              or isinstance(r.parent, Not)))])

                != len(varinfo['relations'])):

                localchecks[varname] = erqlexprs

    return localchecks

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

class ExecutionPlan(object):

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

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

        # original rql syntax tree

        self.rqlst = rqlst

        self.args = args or {}

        # cnx executing the query

        self.cnx = cnx

        # quick reference to the system source

        self.syssource = cnx.repo.system_source

        # execution steps

        self.steps = []

        # various resource accesors

        self.querier = querier

        self.schema = querier.schema

        self.sqlannotate = querier.sqlgen_annotate

        self.rqlhelper = cnx.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 sqlexec(self, sql, args=None):

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

    def execute(self):

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

        for step in self.steps:

            result = step.execute()

        # the latest executed step contains the full query result

        return result

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

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

        return rqlst to actually execute

        """

        cached = None

        if security and self.cnx.read_security:

            # ensure security is turned of when security is inserted,

            # else we may loop for ever...

            if self.cnx.transaction_data.get('security-rqlst-cache'):

                key = self.cache_key

            else:

                key = None

            if key is not None and key in self.cnx.transaction_data:

                cachedunion, args = self.cnx.transaction_data[key]

                union.children[:] = []

                for select in cachedunion.children:

                    union.append(select)

                union.has_text_query = cachedunion.has_text_query

                args.update(self.args)

                self.args = args

                cached = True

            else:

                with self.cnx.security_enabled(read=False):

                    noinvariant = self._insert_security(union)

                if key is not None:

                    self.cnx.transaction_data[key] = (union, self.args)

        else:

            noinvariant = ()

        if cached is None:

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

        for select in union.children[:]:

            for subquery in select.with_:

                self._insert_security(subquery.query)

            localchecks, restricted = self._check_permissions(select)

            if any(localchecks):

                self.cnx.rql_rewriter.insert_local_checks(

                    select, self.args, localchecks, restricted, noinvariant)

        return noinvariant

    def _check_permissions(self, rqlst):

        """Return a dict defining "local checks", i.e. RQLExpression defined in

        the schema that should be inserted in the original query, together with

        a set of variable names which requires some security to be inserted.

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

        read are removed, else if the user *may* read it (i.e. if an rql

        expression is defined for the "read" permission of the related type),

        the local checks dict is updated.

        The local checks dict has entries for each different local check

        necessary, with associated solutions as value, a local check being

        defined by a list of 2-uple (variable name, rql expressions) for each

        variable which has to be checked. 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.

        """

        cnx = self.cnx

        msgs = []

        # dict(varname: eid), allowing to check rql expression for variables

        # which have a known eid

        varkwargs = {}

        if not cnx.transaction_data.get('security-rqlst-cache'):

            for var in rqlst.defined_vars.itervalues():

                if var.stinfo['constnode'] is not None:

                    eid = var.stinfo['constnode'].eval(self.args)

                    varkwargs[var.name] = int(eid)

        # dictionary of variables restricted for security reason

        localchecks = {}

        restricted_vars = set()

        newsolutions = []

        for solution in rqlst.solutions:

            try:

                localcheck = check_read_access(cnx, rqlst, solution, self.args)

            except Unauthorized as ex:

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

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

                msgs.append(msg)

                LOGGER.info(msg)

            else:

                newsolutions.append(solution)

                # 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 entity has been added in the current transaction, the

                    # user can read it whatever rql expressions are associated

                    # to its type

                    if cnx.added_in_transaction(eid):

                        continue

                    for rqlexpr in rqlexprs:

                        if rqlexpr.check(cnx, eid):

                            break

                    else:

                        raise Unauthorized('No read acces on %r with eid %i.' % (var, eid))

                # mark variables protected by an rql expression

                restricted_vars.update(localcheck)

                # turn local check into a dict key

                localcheck = tuple(sorted(localcheck.iteritems()))

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

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

        if not newsolutions:

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

        # if there is some message, solutions have been modified and must be

        # reconsidered by the syntax treee

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

        add_types_restriction(self.schema, select)

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

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

        # save originaly selected variable, we may modify this

        # dictionary for substitution (query parameters)

        self.selected = rqlst.selection

        # list of rows of entities definition (ssplanner.EditedEntity)

        self.e_defs = [[]]

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

        self.r_defs = set()

        # indexes to track entity definitions bound to relation definitions

        self._r_subj_index = {}

        self._r_obj_index = {}

        self._expanded_r_defs = {}

    def add_entity_def(self, edef):

        """add an entity definition to build"""

        self.e_defs[-1].append(edef)

    def add_relation_def(self, rdef):

        """add an relation definition to build"""

        self.r_defs.add(rdef)

        if not isinstance(rdef[0], int):

            self._r_subj_index.setdefault(rdef[0], []).append(rdef)

        if not isinstance(rdef[2], int):

            self._r_obj_index.setdefault(rdef[2], []).append(rdef)

    def substitute_entity_def(self, edef, edefs):

        """substitute an incomplete entity definition by a list of complete

        equivalents

        e.g. on queries such as ::

          INSERT Personne X, Societe Y: X nom N, Y nom 'toto', X travaille Y

          WHERE U login 'admin', U login N

        X will be inserted as many times as U exists, and so the X travaille Y

        relations as to be added as many time as X is inserted

        """

        if not edefs or not self.e_defs:

            # no result, no entity will be created

            self.e_defs = ()

            return

        # first remove the incomplete entity definition

        colidx = self.e_defs[0].index(edef)

        for i, row in enumerate(self.e_defs[:]):

            self.e_defs[i][colidx] = edefs[0]

            samplerow = self.e_defs[i]

            for edef_ in edefs[1:]:

                row = [ed.clone() for i, ed in enumerate(samplerow)

                       if i != colidx]

                row.insert(colidx, edef_)

                self.e_defs.append(row)

        # now, see if this entity def is referenced as subject in some relation

        # definition

        if edef in self._r_subj_index:

            for rdef in self._r_subj_index[edef]:

                expanded = self._expanded(rdef)

                result = []

                for exp_rdef in expanded:

                    for edef_ in edefs: