# 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

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(session, 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 session.describe(term.eval(args))[0]

def check_read_access(session, 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 = session.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(session, rel.children[0],

                                                    solution, args))

                rdef = eschema.rdef(rschema)

            else:

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

                                               solution, args),

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

                                               solution, args))

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

                if DBG:

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

                           (rdef, session.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, session.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 session.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, session.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, 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.cnxset.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 make_temp_table_name(self, table):

        """

        return a temp table name according to db backend

        """

        return self.syssource.make_temp_table_name(table)

    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

        """

        cached = None

        if security and self.session.read_security:

            # ensure security is turned of when security is inserted,

            # else we may loop for ever...

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

                key = self.cache_key

            else:

                key = None

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

                cachedunion, args = self.session.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.session.security_enabled(read=False):

                    noinvariant = self._insert_security(union)

                if key is not None:

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

        """

        session = self.session

        msgs = []

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

        # which have a known eid

        varkwargs = {}

        if not session.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(session, rqlst, solution, self.args)

            except Unauthorized as ex:

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

                msg %= (solution, session.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 session.added_in_transaction(eid):

                        continue

                    for rqlexpr in rqlexprs:

                        if rqlexpr.check(session, 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, 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 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)