"""RQL to SQL generator for native sources.

SQL queries optimization

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

1. CWUser X WHERE X in_group G, G name 'users':

   CWUser is the only subject entity type for the in_group relation,

   which allow us to do ::

     SELECT eid_from FROM in_group, CWGroup

     WHERE in_group.eid_to = CWGroup.eid_from

     AND CWGroup.name = 'users'

2. Any X WHERE X nonfinal1 Y, Y nonfinal2 Z

   -> direct join between nonfinal1 and nonfinal2, whatever X,Y, Z (unless

      inlined...)

      NOT IMPLEMENTED (and quite hard to implement)

Potential optimization information is collected by the querier, sql generation

is done according to this information

cross RDMS note : read `Comparison of different SQL implementations`_

by Troels Arvin. Features SQL ISO Standard, PG, mysql, Oracle, MS SQL, DB2

and Informix.

.. _Comparison of different SQL implementations: http://www.troels.arvin.dk/db/rdbms 

:organization: Logilab

:copyright: 2001-2010 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"

import threading

from logilab.database import FunctionDescr, SQL_FUNCTIONS_REGISTRY

from rql import BadRQLQuery, CoercionError

from rql.stmts import Union, Select

from rql.nodes import (SortTerm, VariableRef, Constant, Function, Not,

                       Variable, ColumnAlias, Relation, SubQuery, Exists)

from cubicweb import QueryError

from cubicweb.server.sqlutils import SQL_PREFIX

from cubicweb.server.utils import cleanup_solutions

ColumnAlias._q_invariant = False # avoid to check for ColumnAlias / Variable

FunctionDescr.source_execute = None

def default_update_cb_stack(self, stack):

    stack.append(self.source_execute)

FunctionDescr.update_cb_stack = default_update_cb_stack

LENGTH = SQL_FUNCTIONS_REGISTRY.get_function('LENGTH')

def length_source_execute(source, value):

    return len(value.getvalue())

LENGTH.source_execute = length_source_execute

def _new_var(select, varname):

    newvar = select.get_variable(varname)

    if not 'relations' in newvar.stinfo:

        # not yet initialized

        newvar.prepare_annotation()

        newvar.stinfo['scope'] = select

        newvar._q_invariant = False

    return newvar

def _fill_to_wrap_rel(var, newselect, towrap, schema):

    for rel in var.stinfo['relations'] - var.stinfo['rhsrelations']:

        rschema = schema.rschema(rel.r_type)

        if rschema.inlined:

            towrap.add( (var, rel) )

            for vref in rel.children[1].iget_nodes(VariableRef):

                newivar = _new_var(newselect, vref.name)

                newselect.selection.append(VariableRef(newivar))

                _fill_to_wrap_rel(vref.variable, newselect, towrap, schema)

        elif rschema.final:

            towrap.add( (var, rel) )

def rewrite_unstable_outer_join(select, solutions, unstable, schema):

    """if some optional variables are unstable, they should be selected in a

    subquery. This function check this and rewrite the rql syntax tree if

    necessary (in place). Return a boolean telling if the tree has been modified

    """

    torewrite = set()

    modified = False

    for varname in tuple(unstable):

        var = select.defined_vars[varname]

        if not var.stinfo.get('optrelations'):

            continue

        modified = True

        unstable.remove(varname)

        torewrite.add(var)

        newselect = Select()

        newselect.need_distinct = newselect.need_intersect = False

        myunion = Union()

        myunion.append(newselect)

        # extract aliases / selection

        newvar = _new_var(newselect, var.name)

        newselect.selection = [VariableRef(newvar)]

        for avar in select.defined_vars.itervalues():

            if avar.stinfo['attrvar'] is var:

                newavar = _new_var(newselect, avar.name)

                newavar.stinfo['attrvar'] = newvar

                newselect.selection.append(VariableRef(newavar))

        towrap_rels = set()

        _fill_to_wrap_rel(var, newselect, towrap_rels, schema)

        # extract relations

        for var, rel in towrap_rels:

            newrel = rel.copy(newselect)

            newselect.add_restriction(newrel)

            select.remove_node(rel)

            var.stinfo['relations'].remove(rel)

            newvar.stinfo['relations'].add(newrel)

            if rel.optional in ('left', 'both'):

                newvar.add_optional_relation(newrel)

            for vref in newrel.children[1].iget_nodes(VariableRef):

                var = vref.variable

                var.stinfo['relations'].add(newrel)

                var.stinfo['rhsrelations'].add(newrel)

                if rel.optional in ('right', 'both'):

                    var.add_optional_relation(newrel)

        # extract subquery solutions

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

        cleanup_solutions(newselect, mysolutions)

        newselect.set_possible_types(solutions)

        # full sub-query

        aliases = [VariableRef(select.get_variable(avar.name, i))

                   for i, avar in enumerate(newselect.selection)]

        select.add_subquery(SubQuery(aliases, myunion), check=False)

    return modified

def _new_solutions(rqlst, solutions):

    """first filter out subqueries variables from solutions"""

    newsolutions = []

    for origsol in solutions:

        asol = {}

        for vname in rqlst.defined_vars:

            asol[vname] = origsol[vname]

        if not asol in newsolutions:

            newsolutions.append(asol)

    return newsolutions

def remove_unused_solutions(rqlst, solutions, varmap, schema):

    """cleanup solutions: remove solutions where invariant variables are taking

    different types

    """

    newsols = _new_solutions(rqlst, solutions)

    existssols = {}

    unstable = set()

    invariants = {}

    for vname, var in rqlst.defined_vars.iteritems():

        vtype = newsols[0][vname]

        if var._q_invariant or vname in varmap:

            # remove invariant variable from solutions to remove duplicates

            # later, then reinserting a type for the variable even later

            for sol in newsols:

                invariants.setdefault(id(sol), {})[vname] = sol.pop(vname)

        elif var.scope is not rqlst:

            # move appart variables which are in a EXISTS scope and are variating

            try:

                thisexistssols, thisexistsvars = existssols[var.scope]

            except KeyError:

                thisexistssols = [newsols[0]]

                thisexistsvars = set()

                existssols[var.scope] = thisexistssols, thisexistsvars

            for i in xrange(len(newsols)-1, 0, -1):

                if vtype != newsols[i][vname]:

                    thisexistssols.append(newsols.pop(i))

                    thisexistsvars.add(vname)

        else:

            # remember unstable variables

            for i in xrange(1, len(newsols)):

                if vtype != newsols[i][vname]:

                    unstable.add(vname)

    if invariants:

        # filter out duplicates

        newsols_ = []

        for sol in newsols:

            if not sol in newsols_:

                newsols_.append(sol)

        newsols = newsols_

        # reinsert solutions for invariants

        for sol in newsols:

            for invvar, vartype in invariants[id(sol)].iteritems():

                sol[invvar] = vartype

        for sol in existssols:

            try:

                for invvar, vartype in invariants[id(sol)].iteritems():

                    sol[invvar] = vartype

            except KeyError:

                continue

    if len(newsols) > 1:

        if rewrite_unstable_outer_join(rqlst, newsols, unstable, schema):

            # remove variables extracted to subqueries from solutions

            newsols = _new_solutions(rqlst, newsols)

    return newsols, existssols, unstable

def relation_info(relation):

    lhs, rhs = relation.get_variable_parts()

    try:

        lhs = lhs.variable

        lhsconst = lhs.stinfo['constnode']

    except AttributeError:

        lhsconst = lhs

        lhs = None

    except KeyError:

        lhsconst = None # ColumnAlias

    try:

        rhs = rhs.variable

        rhsconst = rhs.stinfo['constnode']

    except AttributeError:

        rhsconst = rhs

        rhs = None

    except KeyError:

        rhsconst = None # ColumnAlias

    return lhs, lhsconst, rhs, rhsconst

def switch_relation_field(sql, table=''):

    switchedsql = sql.replace(table + '.eid_from', '__eid_from__')

    switchedsql = switchedsql.replace(table + '.eid_to',

                                      table + '.eid_from')

    return switchedsql.replace('__eid_from__', table + '.eid_to')

def sort_term_selection(sorts, selectedidx, rqlst, groups):

    # XXX beurk

    if isinstance(rqlst, list):

        def append(term):

            rqlst.append(term)

    else:

        def append(term):

            rqlst.selection.append(term.copy(rqlst))

    for sortterm in sorts:

        term = sortterm.term

        if not isinstance(term, Constant) and not str(term) in selectedidx:

            selectedidx.append(str(term))

            append(term)

            if groups:

                for vref in term.iget_nodes(VariableRef):

                    if not vref in groups:

                        groups.append(vref)

def fix_selection_and_group(rqlst, selectedidx, needwrap, selectsortterms,

                            sorts, groups, having):

    if selectsortterms and sorts:

        sort_term_selection(sorts, selectedidx, rqlst, not needwrap and groups)

    if sorts and groups:

        # when a query is grouped, ensure sort terms are grouped as well

        for sortterm in sorts:

            term = sortterm.term

            if not isinstance(term, Constant):

                for vref in term.iget_nodes(VariableRef):

                    if not vref in groups:

                        groups.append(vref)

    if needwrap:

        if groups:

            for vref in groups:

                if not vref.name in selectedidx:

                    selectedidx.append(vref.name)

                    rqlst.selection.append(vref)

        if having:

            for term in having:

                for vref in term.iget_nodes(VariableRef):

                    if not vref.name in selectedidx:

                        selectedidx.append(vref.name)

                        rqlst.selection.append(vref)

def iter_mapped_var_sels(stmt, variable):

    # variable is a Variable or ColumnAlias node mapped to a source side

    # callback

    if not (len(variable.stinfo['rhsrelations']) <= 1 and # < 1 on column alias

            variable.stinfo['selected']):

        raise QueryError("can't use %s as a restriction variable"

                         % variable.name)

    for selectidx in variable.stinfo['selected']:

        vrefs = stmt.selection[selectidx].get_nodes(VariableRef)

        if len(vrefs) != 1:

            raise QueryError()

        yield selectidx, vrefs[0]

def update_source_cb_stack(state, stmt, node, stack):

    while True:

        node = node.parent

        if node is stmt:

            break

        if not isinstance(node, Function):

            raise QueryError()

        func = SQL_FUNCTIONS_REGISTRY.get_function(node.name)

        if func.source_execute is None:

            raise QueryError('%s can not be called on mapped attribute'

                             % node.name)

        state.source_cb_funcs.add(node)

        func.update_cb_stack(stack)

# IGenerator implementation for RQL->SQL #######################################

class StateInfo(object):

    def __init__(self, existssols, unstablevars):

        self.existssols = existssols

        self.unstablevars = unstablevars

        self.subtables = {}

        self.needs_source_cb = None

        self.subquery_source_cb = None

        self.source_cb_funcs = set()

    def reset(self, solution):

        """reset some visit variables"""

        self.solution = solution

        self.count = 0

        self.done = set()

        self.tables = self.subtables.copy()

        self.actual_tables = [[]]

        for _, tsql in self.tables.itervalues():

            self.actual_tables[-1].append(tsql)

        self.outer_tables = {}

        self.duplicate_switches = []

        self.aliases = {}

        self.restrictions = []

        self._restr_stack = []

        self.ignore_varmap = False

        self._needs_source_cb = {}

    def merge_source_cbs(self, needs_source_cb):

        if self.needs_source_cb is None:

            self.needs_source_cb = needs_source_cb

        elif needs_source_cb != self.needs_source_cb:

            raise QueryError('query fetch some source mapped attribute, some not')

    def finalize_source_cbs(self):

        if self.subquery_source_cb is not None:

            self.needs_source_cb.update(self.subquery_source_cb)

    def add_restriction(self, restr):

        if restr:

            self.restrictions.append(restr)

    def iter_exists_sols(self, exists):

        if not exists in self.existssols:

            yield 1

            return

        thisexistssols, thisexistsvars = self.existssols[exists]

        origsol = self.solution

        origtables = self.tables

        done = self.done

        for thisexistssol in thisexistssols:

            for vname in self.unstablevars:

                if thisexistssol[vname] != origsol[vname] and vname in thisexistsvars:

                    break

            else:

                self.tables = origtables.copy()

                self.solution = thisexistssol

                yield 1

                # cleanup self.done from stuff specific to exists

                for var in thisexistsvars:

                    if var in done:

                        done.remove(var)

                for rel in exists.iget_nodes(Relation):

                    if rel in done:

                        done.remove(rel)

        self.solution = origsol

        self.tables = origtables

    def push_scope(self):

        self.actual_tables.append([])

        self._restr_stack.append(self.restrictions)

        self.restrictions = []

    def pop_scope(self):

        restrictions = self.restrictions

        self.restrictions = self._restr_stack.pop()

        return restrictions, self.actual_tables.pop()

class SQLGenerator(object):