stmts.py

# copyright 2004-2021 LOGILAB S.A. (Paris, FRANCE), all rights reserved.
# contact http://www.logilab.fr/ -- mailto:contact@logilab.fr
#
# This file is part of rql.
#
# rql 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.
#
# rql 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 rql. If not, see <http://www.gnu.org/licenses/>.
"""Construction and manipulation of RQL syntax trees.

This module defines only first level nodes (i.e. statements). Child nodes are
defined in the nodes module
"""
from copy import deepcopy
from warnings import warn

from logilab.common.deprecation import callable_deprecated

from rql import BadRQLQuery, CoercionError, nodes
from rql.base import BaseNode, Node
from rql.utils import rqlvar_maker
from rql import rqltypes as rt

from typing import (
    TYPE_CHECKING,
    Dict,
    List,
    Union as Union_,
    Any,
    Iterable,
    Optional,
    cast,
    Set as Set_,
    Tuple,
    Iterator,
)

_MARKER: object = object()

__docformat__: str = "restructuredtext en"

if TYPE_CHECKING:
    import rql

Solution = Dict[str, str]
SolutionsList = List[Solution]


def _check_references(
    defined: Dict[str, Union_["rql.nodes.Variable", "rql.nodes.ColumnAlias"]],
    varrefs: Iterable[Union_["rql.nodes.VariableRef", "rql.base.BaseNode"]],
) -> bool:
    refs = {}
    for var in defined.values():
        for vref in var.references():
            # be careful, Variable and VariableRef define __cmp__
            if not [v for v in varrefs if v is vref]:
                raise AssertionError("vref %r is not in the tree" % vref)
            refs[id(vref)] = 1
    for vref in varrefs:
        if id(vref) not in refs:
            raise AssertionError("vref %r is not referenced (%r)" % (vref, vref.stmt))
    return True


class undo_modification:
    def __init__(self, select):
        self.select = select

    def __enter__(self):
        self.select.save_state()

    def __exit__(self):
        self.select.recover()


class ScopeNode(BaseNode):
    def __init__(self):
        # dictionnary of defined variables in the original RQL syntax tree
        self.defined_vars: Dict[str, "rql.nodes.Variable"] = {}
        self.with_: List["rql.nodes.SubQuery"] = []
        # list of possibles solutions for used variables
        self.solutions: SolutionsList = []
        self._varmaker = None  # variable names generator, built when necessary
        self.where: Optional["rql.base.Node"] = None  # where clause node
        self.having: Iterable["rql.base.Node"] = ()  # XXX now a single node
        # "ScopeNode" has no attribute "schema"
        self.schema: Optional[Any] = None
        # "ScopeNode" has no attribute "aliases"
        self.aliases: Dict[str, "rql.nodes.ColumnAlias"] = {}

    # "ScopeNode" has no attribute "undo_manager"
    @property
    def undo_manager(self):
        try:
            return self._undo_manager
        except AttributeError:
            from rql.undo import SelectionManager

            self._undo_manager = SelectionManager(self)
            return self._undo_manager

    @property
    def should_register_op(self):
        return None

    def get_selected_variables(self):
        return self.selected_terms()

    def set_where(self, node: "rql.base.Node") -> None:
        self.where = node
        node.parent = self

    def set_having(self, terms: Iterable["rql.base.Node"]) -> None:
        if self.should_register_op:
            from rql.undo import SetHavingOperation

            # "ScopeNode" has no attribute "undo_manager"  [attr-defined]
            # either we ignore, either we redefine
            self.undo_manager.add_operation(SetHavingOperation(self, self.having))
        self.having = terms
        for node in terms:
            node.parent = self

    # Signature of "copy" incompatible with supertype "BaseNode"
    # either add a parameter that won't be used, either ignore
    def copy(
        self,
        stmt: Optional["rql.stmts.Statement"] = None,
        copy_solutions: bool = True,
        solutions: SolutionsList = None,
    ) -> "rql.base.BaseNode":
        new = self.__class__()
        if self.schema is not None:
            new.schema = self.schema
        if solutions is not None:
            new.solutions = solutions
        elif copy_solutions and self.solutions:
            new.solutions = deepcopy(self.solutions)
        return new

    # construction helper methods #############################################
    def get_etype(self, name: str) -> "rql.nodes.Constant":
        """return the type object for the given entity's type name

        raise BadRQLQuery on unknown type
        """
        return nodes.Constant(name, "etype")

    def get_variable(
        self, name: str
    ) -> Union_["rql.nodes.Variable", "rql.nodes.ColumnAlias"]:
        """get a variable instance from its name

        the variable is created if it doesn't exist yet
        """
        try:
            return self.defined_vars[name]
        except Exception:
            self.defined_vars[name] = var = nodes.Variable(name)
            var.stmt = self
            return var

    def allocate_varname(self) -> str:
        """return an yet undefined variable name"""
        if self._varmaker is None:
            self._varmaker = rqlvar_maker(
                defined=self.defined_vars,
                # XXX only on Select node
                aliases=getattr(self, "aliases", None),
            )
        return next(self._varmaker)

    def make_variable(self) -> "rql.nodes.Variable":
        """create a new variable with an unique name for this tree"""
        var = self.get_variable(self.allocate_varname())
        if self.should_register_op:
            from rql.undo import MakeVarOperation

            self.undo_manager.add_operation(MakeVarOperation(var))
        return cast("rql.nodes.Variable", var)

    def set_possible_types(
        self,
        solutions: SolutionsList,
        kwargs: Optional[Union_[object, Dict[str, str]]] = _MARKER,
        key: str = "possibletypes",
    ) -> None:
        if key == "possibletypes":
            self.solutions = solutions
        defined = self.defined_vars
        for var in defined.values():
            var.stinfo[key] = set()
            for solution in solutions:
                var.stinfo[key].add(solution[var.name])
        # for debugging
        # for sol in solutions:
        #    for vname in sol:
        #        assert vname in self.defined_vars or vname in self.aliases

    def check_references(self) -> bool:
        """test function"""
        try:
            defined = cast(
                Dict[str, Union_["rql.nodes.ColumnAlias", "rql.nodes.Variable"]],
                self.aliases.copy(),
            )

        except AttributeError:
            defined = cast(
                Dict[str, Union_["rql.nodes.ColumnAlias", "rql.nodes.Variable"]],
                self.defined_vars.copy(),
            )
        else:
            defined.update(self.defined_vars)
            for subq in self.with_:
                subq.query.check_references()
        varrefs = [
            vref for vref in self.get_nodes(nodes.VariableRef) if vref.stmt is self
        ]
        try:
            _check_references(defined, varrefs)
        except Exception:
            print(repr(self))
            raise
        return True


class Statement:
    """base class for statement nodes"""

    # default values for optional instance attributes, set on the instance when
    # used
    schema: Optional["rql.interfaces.ISchema"] = None
    annotated: bool = False  # set by the annotator

    if TYPE_CHECKING:

        def get_variable(self, name, column=None):
            raise NotImplementedError()

    # navigation helper methods #############################################

    @property
    def root(self):
        """return the root node of the tree"""
        return self

    @property
    def stmt(self):
        return self

    @property
    def scope(self):
        return self

    def ored(
        self, traverse_scope: bool = False, _fromnode: Optional["rql.nodes.And"] = None
    ) -> Optional["rql.nodes.Or"]:
        return None

    def neged(
        self, traverse_scope: bool = False, _fromnode: Optional["rql.nodes.Or"] = None
    ) -> Optional["rql.nodes.Not"]:
        return None


class Union(Statement, Node):
    """the select node is the root of the syntax tree for selection statement
    using UNION
    """

    TYPE: str = "select"
    # default values for optional instance attributes, set on the instance when
    # used
    undoing: bool = False  # used to prevent from memorizing when undoing !
    memorizing: int = 0  # recoverable modification attributes
    children: List["rql.stmts.Select"]

    def wrap_selects(self) -> None:
        """return a new rqlst root containing the given union as a subquery"""
        child = Union()
        for select in self.children[:]:
            child.append(select)
            self.remove_select(select)
        newselect: "rql.stmts.Select" = Select()
        aliases: List["rql.nodes.VariableRef"] = []
        for i in range(len(select.selection)):
            aliases.append(nodes.VariableRef(newselect.make_variable()))
        newselect.add_subquery(nodes.SubQuery(aliases, child), check=False)
        for vref in aliases:
            newselect.append_selected(nodes.VariableRef(vref.variable))
        self.append_select(newselect)

    def _get_offset(self) -> int:
        warn("offset is now a Select node attribute", DeprecationWarning, stacklevel=2)
        last_children = self.children[-1]
        return last_children.offset

    def set_offset(self, offset: int) -> None:
        if len(self.children) == 1:
            last_children = cast("rql.stmts.Select", self.children[-1])
            last_children.set_offset(offset)
        # we have to introduce a new root
        # XXX not undoable since a new root has to be introduced
        self.wrap_selects()
        first_child = cast("rql.stmts.Select", self.children[0])
        first_child.set_offset(offset)

    offset = property(_get_offset, set_offset)

    def _get_limit(self):
        warn("limit is now a Select node attribute", DeprecationWarning, stacklevel=2)
        return self.children[-1].limit

    def set_limit(self, limit: int) -> None:
        if len(self.children) == 1:
            self.children[-1].set_limit(limit)
            return None
        self.wrap_selects()
        self.children[0].set_limit(limit)
        return None

    limit = property(_get_limit, set_limit)

    @property
    def root(self):
        """return the root node of the tree"""
        if self.parent is None:
            return self
        return self.parent.root

    def get_description(
        self,
        mainindex: Optional[int] = None,
        tr: Optional[rt.TranslationFunction] = None,
    ) -> List[List[str]]:
        """
        `mainindex`:
          selection index to consider as main column, useful to get smarter
          results
        `tr`:
          optional translation function taking a string as argument and
          returning a string
        """
        if tr is None:

            def tr(msg, context=None):
                return msg

        return [c.get_description(mainindex, tr) for c in self.children]

    # repr / as_string / copy #################################################

    def __repr__(self) -> str:
        return "\nUNION\n".join(repr(select) for select in self.children)

    def as_string(self, kwargs: Optional[Dict] = None) -> str:
        """return the tree as an encoded rql string"""
        strings: List[str] = [
            select.as_string(kwargs=kwargs) for select in self.children
        ]
        if len(strings) == 1:
            return strings[0]
        return " UNION ".join("(%s)" % part for part in strings)

    def copy(
        self, stmt: Optional["rql.stmts.Statement"] = None, copy_children: bool = True
    ) -> "rql.stmts.Union":
        new: "rql.stmts.Union" = Union()
        if self.schema is not None:
            new.schema = self.schema
        if copy_children:
            for child in self.children:
                new.append(child.copy())
                assert new.children[-1].parent is new
        return new

    # union specific methods ##################################################

    # XXX for bw compat, should now use get_variable_indices (cw > 3.8.4)
    def get_variable_variables(self) -> Set_[int]:
        change: Set_[int] = set()
        for idx in self.get_variable_indices():
            first_child = self.children[0]
            vrefs = first_child.selection[idx].iget_nodes(nodes.VariableRef)

            for vref in vrefs:
                change.add(vref.name)
        return change

    def get_variable_indices(self) -> Set_[int]:
        """return the set of selection indexes which take different types
        according to the solutions
        """
        change: Set_[int] = set()
        values: Dict[int, Set_] = {}
        for select in self.children:
            for descr in select.get_selection_solutions():
                for i, etype in enumerate(descr):
                    values.setdefault(i, set()).add(etype)
        for idx, etypes in values.items():
            if len(etypes) > 1:
                change.add(idx)
        return change

    def _locate_subquery(
        self, col: int, etype: str, kwargs: Optional[Dict[Any, Any]] = None
    ) -> Tuple:
        first_child = self.children[0]
        has_children = len(self.children) == 1
        first_child_subqueries = not first_child.with_
        if has_children and first_child_subqueries:
            return self.children[0], col
        for select in self.children:
            term = select.selection[col]
            try:
                if term.name in select.aliases:
                    alias = select.aliases[term.name]
                    return alias.query._locate_subquery(alias.colnum, etype, kwargs)
            except AttributeError:
                # term has no 'name' attribute
                pass
            for i, solution in enumerate(select.solutions):
                if term.get_type(solution, kwargs) == etype:
                    return select, col
        raise Exception(f"internal error, {etype} not found on col {col}")

    def locate_subquery(
        self, col: int, etype: str, kwargs: Optional[Dict] = None
    ) -> Any:
        """return a select node and associated selection index where root
        variable at column `col` is of type `etype`
        """
        try:
            # Cannot determine type of '_subq_cache'  [has-type]
            return self._subq_cache[(col, etype)]  # type: ignore[has-type]
        except AttributeError:
            self._subq_cache = {}
        except KeyError:
            pass
        self._subq_cache[(col, etype)] = self._locate_subquery(col, etype, kwargs)
        return self._subq_cache[(col, etype)]

    def subquery_selection_index(self, subselect: Any, col: int) -> int:
        """given a select sub-query and a column index in the root query, return
        the selection index for this column in the sub-query
        """
        selectpath: List = []
        while subselect.parent.parent is not None:
            subq = subselect.parent.parent
            subselect = subq.parent
            selectpath.insert(0, subselect)
        for select in selectpath:
            col = select.selection[col].variable.colnum
        return col

    # recoverable modification methods ########################################

    # don't use @cached: we want to be able to disable it while this must still
    # be cached
    @property
    def undo_manager(self) -> "rql.undo.SelectionManager":
        from rql.undo import SelectionManager

        undo_manager = getattr(self, "_undo_manager", None)
        if undo_manager:
            return undo_manager
        self._undo_manager = SelectionManager(self)
        return self._undo_manager

    @property
    def should_register_op(self):
        return self.memorizing and not self.undoing

    def undo_modification(self) -> "rql.stmts.undo_modification":
        return undo_modification(self)

    def save_state(self) -> None:
        """save the current tree"""
        self.undo_manager.push_state()
        self.memorizing += 1

    def recover(self) -> None:
        """reverts the tree as it was when save_state() was last called"""
        self.memorizing -= 1
        assert self.memorizing >= 0
        self.undo_manager.recover()

    def check_references(self) -> bool:
        """test function"""
        for select in self.children:
            select.check_references()
        return True

    def append_select(self, select: "rql.stmts.Select") -> None:
        if self.should_register_op:
            from rql.undo import AppendSelectOperation

            self.undo_manager.add_operation(AppendSelectOperation(self, select))
        self.children.append(select)

    def remove_select(self, select: "rql.stmts.Select") -> None:
        idx: int = self.children.index(select)
        if self.should_register_op:
            from rql.undo import RemoveSelectOperation

            self.undo_manager.add_operation(RemoveSelectOperation(self, select, idx))
        self.children.pop(idx)


class Select(Statement, nodes.EditableMixIn, ScopeNode):
    """the select node is the base statement of the syntax tree for selection
    statement, always child of a UNION root.
    """

    vargraph: rt.Graph = {}
    parent = None
    distinct: bool = False
    # limit / offset
    limit: Optional[int] = None
    offset: int = 0
    # already defined inside ScopeNode right?
    # But py3 fails when I change anything here
    # RecursionError: maximum recursion depth excedeed
    # set by the annotator
    has_aggregat: bool = False

    def __init__(self):
        Statement.__init__(self)
        ScopeNode.__init__(self)
        self.selection: List = []
        # subqueries alias
        self.aliases: Dict[str, "rql.nodes.ColumnAlias"] = {}
        # syntax tree meta-information
        self.stinfo: Dict[str, Dict] = {"rewritten": {}}

        # select clauses
        self.groupby: List[Any] = []
        self.orderby: List[Any] = []

    @property
    def root(self):
        """return the root node of the tree"""
        return self.parent

    def get_description(
        self,
        mainindex: Optional[int] = None,
        tr: Optional[rt.TranslationFunction] = None,
    ) -> List[str]:
        """return the list of types or relations (if not found) associated to
        selected variables.
        mainindex is an optional selection index which should be considered has
        'pivot' entity.
        """
        descr: List[str] = []
        for term in self.selection:
            # don't translate Any
            try:
                descr.append(term.get_description(mainindex, tr) or "Any")
            except CoercionError:
                descr.append("Any")
        return descr

    @property
    def children(self):
        children = self.selection[:]
        if self.groupby:
            children += self.groupby
        if self.orderby:
            children += self.orderby
        if self.where:
            children.append(self.where)
        if self.having:
            children += self.having
        if self.with_:
            children += self.with_
        return children

    # repr / as_string / copy #################################################

    def __repr__(self) -> str:
        return self.as_string(userepr=True)

    def as_string(self, kwargs: Optional[Dict] = None, userepr: bool = False) -> str:
        """return the tree as an encoded rql string"""
        if userepr:
            as_string = repr
        else:

            def as_string(x):
                return x.as_string(kwargs=kwargs)

        s = [",".join(as_string(term) for term in self.selection)]
        if self.groupby:
            s.append("GROUPBY " + ",".join(as_string(term) for term in self.groupby))
        if self.orderby:
            s.append("ORDERBY " + ",".join(as_string(term) for term in self.orderby))
        if self.limit is not None:
            s.append("LIMIT %s" % self.limit)
        if self.offset:
            s.append("OFFSET %s" % self.offset)
        if self.where is not None:
            s.append("WHERE " + as_string(self.where))
        if self.having:
            s.append("HAVING " + ",".join(as_string(term) for term in self.having))
        if self.with_:
            s.append("WITH " + ",".join(as_string(term) for term in self.with_))
        if self.distinct:
            return "DISTINCT Any " + " ".join(s)
        return "Any " + " ".join(s)

    def copy(
        self,
        stmt: Optional["rql.stmts.Statement"] = None,
        copy_solutions: bool = True,
        solutions: Optional[SolutionsList] = None,
    ) -> "rql.stmts.Select":
        new = super().copy(self, copy_solutions, solutions)

        # "ScopeNode" has no attribute ....  [attr-defined]
        new = cast("rql.stmts.Select", new)

        if self.with_:
            new.set_with([sq.copy(new) for sq in self.with_], check=False)
        for child in self.selection:
            new.append_selected(child.copy(new))
        if self.groupby:
            new.set_groupby([sq.copy(new) for sq in self.groupby])
        if self.orderby:
            new.set_orderby([sq.copy(new) for sq in self.orderby])
        # Argument 1 to "set_where" of "ScopeNode" has incompatible type "Node"
        # expected "Union[Or, Not, And, Relation]"  [arg-type]
        if self.where:
            new.set_where(self.where.copy(new))
        if self.having:
            new.set_having([sq.copy(new) for sq in self.having])
        new.distinct = self.distinct
        new.limit = self.limit
        new.offset = self.offset
        new.vargraph = self.vargraph
        return new

    # select specific methods #################################################

    def set_possible_types(
        self,
        solutions: SolutionsList,
        kwargs: Optional[Union_[object, Dict[str, str]]] = _MARKER,
        key: str = "possibletypes",
    ) -> None:
        super(Select, self).set_possible_types(solutions, kwargs, key)
        for ca in self.aliases.values():
            ca.stinfo[key] = capt = set()
            for solution in solutions:
                capt.add(solution[ca.name])
            if kwargs is _MARKER:
                continue
            # propagage to subqueries in case we're introducing additional
            # type constraints
            for stmt in ca.query.children[:]:
                # better type for term
                term: Any = stmt.selection[ca.colnum]
                sols: List = [
                    sol for sol in stmt.solutions if term.get_type(sol, kwargs) in capt
                ]
                if not sols:
                    ca.query.remove_select(stmt)
                else:
                    stmt.set_possible_types(sols)

    def set_statement_type(self, etype: str) -> None:
        """set the statement type for this selection
        this method must be called last (i.e. once selected variables has been
        added)
        """
        assert self.selection
        # Person P  ->  Any P where P is Person
        if etype != "Any":
            variables: List["rql.nodes.VariableRef"] = list(
                self.get_selected_variables()
            )
            if not variables:
                raise BadRQLQuery(
                    "Setting type in selection is only allowed "
                    "when some variable is selected"
                )
            for var in variables:
                self.add_type_restriction(var.variable, etype)

    def set_distinct(self, value: bool) -> None:
        """mark DISTINCT query"""
        if self.should_register_op and value != self.distinct:
            from rql.undo import SetDistinctOperation

            self.undo_manager.add_operation(SetDistinctOperation(self.distinct, self))
        self.distinct = value

    def set_limit(self, limit: int) -> None:
        if limit is not None and (not isinstance(limit, int) or limit <= 0):
            raise BadRQLQuery("bad limit %s" % limit)
        if self.should_register_op and limit != self.limit:
            from rql.undo import SetLimitOperation

            self.undo_manager.add_operation(SetLimitOperation(self.limit, self))
        self.limit = limit

    def set_offset(self, offset: int) -> None:
        if offset is not None and (not isinstance(offset, int) or offset < 0):
            raise BadRQLQuery("bad offset %s" % offset)
        if self.should_register_op and offset != self.offset:
            from rql.undo import SetOffsetOperation

            self.undo_manager.add_operation(SetOffsetOperation(self.offset, self))
        self.offset = offset

    def set_orderby(self, terms: List["rql.nodes.SortTerm"]) -> None:
        self.orderby = terms
        for node in terms:
            node.parent = self

    def set_groupby(
        self, terms: List[Union_["rql.nodes.Function", "rql.nodes.VariableRef"]]
    ) -> None:
        self.groupby = terms
        for node in terms:
            node.parent = self

    def set_with(self, terms: List["rql.nodes.SubQuery"], check: bool = True) -> None:
        self.with_ = []
        for node in terms:
            self.add_subquery(node, check)

    def add_subquery(self, node: "rql.nodes.SubQuery", check: bool = True) -> None:

        assert node.query
        node.parent = self
        self.with_.append(node)

        # "BaseNode" has no attribute "selection"
        if check and len(node.aliases) != len(
            cast("rql.stmts.Select", node.query.children[0]).selection
        ):
            raise BadRQLQuery(
                "Should have the same number of aliases than "
                "selected terms in sub-query"
            )
        for i, alias in enumerate(node.aliases):
            if check and alias.name in self.aliases:
                raise BadRQLQuery("Duplicated alias %s" % alias)
            ca = self.get_variable(alias.name, i)
            # "Variable" has no attribute "query"
            # classes with query attribute: ColumnAlias, Exists, SubQuery
            ca = cast("rql.nodes.ColumnAlias", ca)
            ca.query = node.query

    def remove_subquery(self, node: "rql.nodes.SubQuery") -> None:
        self.with_.remove(node)
        node.parent = None
        for i, alias in enumerate(node.aliases):
            del self.aliases[alias.name]

    # Signature of "get_variable" incompatible with supertype "ScopeNode"  [override]
    def get_variable(
        self, name: str, colnum: Optional[int] = None
    ) -> Union_["rql.nodes.Variable", "rql.nodes.ColumnAlias"]:
        """get a variable instance from its name

        the variable is created if it doesn't exist yet
        """
        if name in self.aliases.keys():
            return self.aliases[name]
        if colnum is not None:  # take care, may be 0
            self.aliases[name] = calias = nodes.ColumnAlias(name, colnum)
            calias.stmt = self
            # alias may already have been used as a regular variable, replace it
            if name in self.defined_vars: