1#  Licensed to Elasticsearch B.V. under one or more contributor
2#  license agreements. See the NOTICE file distributed with
3#  this work for additional information regarding copyright
4#  ownership. Elasticsearch B.V. licenses this file to you under
5#  the Apache License, Version 2.0 (the "License"); you may
6#  not use this file except in compliance with the License.
7#  You may obtain a copy of the License at
8#
9# 	http://www.apache.org/licenses/LICENSE-2.0
10#
11#  Unless required by applicable law or agreed to in writing,
12#  software distributed under the License is distributed on an
13#  "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14#  KIND, either express or implied.  See the License for the
15#  specific language governing permissions and limitations
16#  under the License.
17
18from .connections import get_connection
19from .query import Bool, Q
20from .response import UpdateByQueryResponse
21from .search import ProxyDescriptor, QueryProxy, Request
22
23
24class UpdateByQuery(Request):
25
26    query = ProxyDescriptor("query")
27
28    def __init__(self, **kwargs):
29        """
30        Update by query request to elasticsearch.
31
32        :arg using: `Elasticsearch` instance to use
33        :arg index: limit the search to index
34        :arg doc_type: only query this type.
35
36        All the parameters supplied (or omitted) at creation type can be later
37        overriden by methods (`using`, `index` and `doc_type` respectively).
38
39        """
40        super(UpdateByQuery, self).__init__(**kwargs)
41        self._response_class = UpdateByQueryResponse
42        self._script = {}
43        self._query_proxy = QueryProxy(self, "query")
44
45    def filter(self, *args, **kwargs):
46        return self.query(Bool(filter=[Q(*args, **kwargs)]))
47
48    def exclude(self, *args, **kwargs):
49        return self.query(Bool(filter=[~Q(*args, **kwargs)]))
50
51    @classmethod
52    def from_dict(cls, d):
53        """
54        Construct a new `UpdateByQuery` instance from a raw dict containing the search
55        body. Useful when migrating from raw dictionaries.
56
57        Example::
58
59            ubq = UpdateByQuery.from_dict({
60                "query": {
61                    "bool": {
62                        "must": [...]
63                    }
64                },
65                "script": {...}
66            })
67            ubq = ubq.filter('term', published=True)
68        """
69        u = cls()
70        u.update_from_dict(d)
71        return u
72
73    def _clone(self):
74        """
75        Return a clone of the current search request. Performs a shallow copy
76        of all the underlying objects. Used internally by most state modifying
77        APIs.
78        """
79        ubq = super(UpdateByQuery, self)._clone()
80
81        ubq._response_class = self._response_class
82        ubq._script = self._script.copy()
83        ubq.query._proxied = self.query._proxied
84        return ubq
85
86    def response_class(self, cls):
87        """
88        Override the default wrapper used for the response.
89        """
90        ubq = self._clone()
91        ubq._response_class = cls
92        return ubq
93
94    def update_from_dict(self, d):
95        """
96        Apply options from a serialized body to the current instance. Modifies
97        the object in-place. Used mostly by ``from_dict``.
98        """
99        d = d.copy()
100        if "query" in d:
101            self.query._proxied = Q(d.pop("query"))
102        if "script" in d:
103            self._script = d.pop("script")
104        self._extra.update(d)
105        return self
106
107    def script(self, **kwargs):
108        """
109        Define update action to take:
110        https://www.elastic.co/guide/en/elasticsearch/reference/current/modules-scripting-using.html
111        for more details.
112
113        Note: the API only accepts a single script, so
114        calling the script multiple times will overwrite.
115
116        Example::
117
118            ubq = Search()
119            ubq = ubq.script(source="ctx._source.likes++"")
120            ubq = ubq.script(source="ctx._source.likes += params.f"",
121                         lang="expression",
122                         params={'f': 3})
123        """
124        ubq = self._clone()
125        if ubq._script:
126            ubq._script = {}
127        ubq._script.update(kwargs)
128        return ubq
129
130    def to_dict(self, **kwargs):
131        """
132        Serialize the search into the dictionary that will be sent over as the
133        request'ubq body.
134
135        All additional keyword arguments will be included into the dictionary.
136        """
137        d = {}
138        if self.query:
139            d["query"] = self.query.to_dict()
140
141        if self._script:
142            d["script"] = self._script
143
144        d.update(self._extra)
145
146        d.update(kwargs)
147        return d
148
149    def execute(self):
150        """
151        Execute the search and return an instance of ``Response`` wrapping all
152        the data.
153        """
154        es = get_connection(self._using)
155
156        self._response = self._response_class(
157            self,
158            es.update_by_query(index=self._index, body=self.to_dict(), **self._params),
159        )
160        return self._response
161