Skip to content

Commit 27d8977

Browse files
miguelgrinbergmiguelgrinberg
committed
more operators
1 parent 3d88ee5 commit 27d8977

File tree

5 files changed

+278
-81
lines changed

5 files changed

+278
-81
lines changed

‎elasticsearch/dsl/__init__.py‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@
1919
from .aggs import A, Agg
2020
from .analysis import analyzer, char_filter, normalizer, token_filter, tokenizer
2121
from .document import AsyncDocument, Document
22-
from .document_base import InnerDoc, M, MetaField, mapped_field
22+
from .document_base import E, InnerDoc, M, MetaField, mapped_field
2323
from .exceptions import (
2424
ElasticsearchDslException,
2525
IllegalOperation,
@@ -135,6 +135,7 @@
135135
"Double",
136136
"DoubleRange",
137137
"DslBase",
138+
"E",
138139
"ElasticsearchDslException",
139140
"EmptySearch",
140141
"Facet",

‎elasticsearch/dsl/document_base.py‎

Lines changed: 98 additions & 46 deletions
Original file line numberDiff line numberDiff line change
@@ -63,6 +63,11 @@ class InstrumentedExpression:
6363
def __init__(self, expr: str):
6464
self._expr = expr
6565

66+
def _render_value(self, value: Any) -> str:
67+
if isinstance(value, InstrumentedExpression):
68+
return str(value)
69+
return json.dumps(value)
70+
6671
def __str__(self) -> str:
6772
return self._expr
6873

@@ -76,61 +81,136 @@ def __neg__(self) -> "InstrumentedExpression":
7681
return InstrumentedExpression(f"-({self._expr})")
7782

7883
def __eq__(self, value: Any) -> "InstrumentedExpression": # type: ignore[override]
79-
return InstrumentedExpression(f"{self._expr} == {json.dumps(value)}")
84+
return InstrumentedExpression(f"{self._expr} == {self._render_value(value)}")
8085

8186
def __ne__(self, value: Any) -> "InstrumentedExpression": # type: ignore[override]
82-
return InstrumentedExpression(f"{self._expr} != {json.dumps(value)}")
87+
return InstrumentedExpression(f"{self._expr} != {self._render_value(value)}")
8388

8489
def __lt__(self, value: Any) -> "InstrumentedExpression":
85-
return InstrumentedExpression(f"{self._expr} < {json.dumps(value)}")
90+
return InstrumentedExpression(f"{self._expr} < {self._render_value(value)}")
8691

8792
def __gt__(self, value: Any) -> "InstrumentedExpression":
88-
return InstrumentedExpression(f"{self._expr} > {json.dumps(value)}")
93+
return InstrumentedExpression(f"{self._expr} > {self._render_value(value)}")
8994

9095
def __le__(self, value: Any) -> "InstrumentedExpression":
91-
return InstrumentedExpression(f"{self._expr} <= {json.dumps(value)}")
96+
return InstrumentedExpression(f"{self._expr} <= {self._render_value(value)}")
9297

9398
def __ge__(self, value: Any) -> "InstrumentedExpression":
94-
return InstrumentedExpression(f"{self._expr} >= {json.dumps(value)}")
99+
return InstrumentedExpression(f"{self._expr} >= {self._render_value(value)}")
95100

96101
def __add__(self, value: Any) -> "InstrumentedExpression":
97-
return InstrumentedExpression(f"{self._expr} + {json.dumps(value)}")
102+
return InstrumentedExpression(f"{self._expr} + {self._render_value(value)}")
98103

99104
def __radd__(self, value: Any) -> "InstrumentedExpression":
100-
return InstrumentedExpression(f"{json.dumps(value)} + {self._expr}")
105+
return InstrumentedExpression(f"{self._render_value(value)} + {self._expr}")
101106

102107
def __sub__(self, value: Any) -> "InstrumentedExpression":
103-
return InstrumentedExpression(f"{self._expr} - {json.dumps(value)}")
108+
return InstrumentedExpression(f"{self._expr} - {self._render_value(value)}")
104109

105110
def __rsub__(self, value: Any) -> "InstrumentedExpression":
106-
return InstrumentedExpression(f"{json.dumps(value)} - {self._expr}")
111+
return InstrumentedExpression(f"{self._render_value(value)} - {self._expr}")
107112

108113
def __mul__(self, value: Any) -> "InstrumentedExpression":
109-
return InstrumentedExpression(f"{self._expr} * {json.dumps(value)}")
114+
return InstrumentedExpression(f"{self._expr} * {self._render_value(value)}")
110115

111116
def __rmul__(self, value: Any) -> "InstrumentedExpression":
112-
return InstrumentedExpression(f"{json.dumps(value)} * {self._expr}")
117+
return InstrumentedExpression(f"{self._render_value(value)} * {self._expr}")
113118

114119
def __truediv__(self, value: Any) -> "InstrumentedExpression":
115-
return InstrumentedExpression(f"{self._expr} / {json.dumps(value)}")
120+
return InstrumentedExpression(f"{self._expr} / {self._render_value(value)}")
116121

117122
def __rtruediv__(self, value: Any) -> "InstrumentedExpression":
118-
return InstrumentedExpression(f"{json.dumps(value)} / {self._expr}")
123+
return InstrumentedExpression(f"{self._render_value(value)} / {self._expr}")
119124

120125
def __mod__(self, value: Any) -> "InstrumentedExpression":
121-
return InstrumentedExpression(f"{self._expr} % {json.dumps(value)}")
126+
return InstrumentedExpression(f"{self._expr} % {self._render_value(value)}")
122127

123128
def __rmod__(self, value: Any) -> "InstrumentedExpression":
124-
return InstrumentedExpression(f"{json.dumps(value)} % {self._expr}")
129+
return InstrumentedExpression(f"{self._render_value(value)} % {self._expr}")
130+
131+
def is_null(self) -> "InstrumentedExpression":
132+
"""Compare the expression against NULL."""
133+
return InstrumentedExpression(f"{self._expr} IS NULL")
134+
135+
def is_not_null(self) -> "InstrumentedExpression":
136+
"""Compare the expression against NOT NULL."""
137+
return InstrumentedExpression(f"{self._expr} IS NOT NULL")
138+
139+
def in_(self, *values: Any) -> "InstrumentedExpression":
140+
"""Test if the expression equals one of the given values."""
141+
rendered_values = ", ".join([f"{value}" for value in values])
142+
return InstrumentedExpression(f"{self._expr} IN ({rendered_values})")
143+
144+
def like(self, *patterns: str) -> "InstrumentedExpression":
145+
"""Filter the expression using a string pattern."""
146+
if len(patterns) == 1:
147+
return InstrumentedExpression(
148+
f"{self._expr} LIKE {self._render_value(patterns[0])}"
149+
)
150+
else:
151+
return InstrumentedExpression(
152+
f'{self._expr} LIKE ({", ".join([self._render_value(p) for p in patterns])})'
153+
)
154+
155+
def rlike(self, *patterns: str) -> "InstrumentedExpression":
156+
"""Filter the expression using a regular expression."""
157+
if len(patterns) == 1:
158+
return InstrumentedExpression(
159+
f"{self._expr} RLIKE {self._render_value(patterns[0])}"
160+
)
161+
else:
162+
return InstrumentedExpression(
163+
f'{self._expr} RLIKE ({", ".join([self._render_value(p) for p in patterns])})'
164+
)
165+
166+
def match(self, query: str) -> "InstrumentedExpression":
167+
"""Perform a match query on the field."""
168+
return InstrumentedExpression(f"{self._expr}:{self._render_value(query)}")
169+
170+
def asc(self) -> "InstrumentedExpression":
171+
"""Return the field name representation for ascending sort order.
172+
173+
For use in ES|QL queries only.
174+
"""
175+
return InstrumentedExpression(f"{self._expr} ASC")
176+
177+
def desc(self) -> "InstrumentedExpression":
178+
"""Return the field name representation for descending sort order.
179+
180+
For use in ES|QL queries only.
181+
"""
182+
return InstrumentedExpression(f"{self._expr} DESC")
183+
184+
def nulls_first(self) -> "InstrumentedExpression":
185+
"""Return the field name representation for nulls first sort order.
186+
187+
For use in ES|QL queries only.
188+
"""
189+
return InstrumentedExpression(f"{self._expr} NULLS FIRST")
190+
191+
def nulls_last(self) -> "InstrumentedExpression":
192+
"""Return the field name representation for nulls last sort order.
193+
194+
For use in ES|QL queries only.
195+
"""
196+
return InstrumentedExpression(f"{self._expr} NULLS LAST")
125197

126198
def where(
127-
self, expr: Union[str, "InstrumentedExpression"]
199+
self, *expressions: Union[str, "InstrumentedExpression"]
128200
) -> "InstrumentedExpression":
129201
"""Add a condition to be met for the row to be included.
130202
131203
Use only in expressions given in the ``STATS`` command.
132204
"""
133-
return InstrumentedExpression(f"{self._expr} WHERE {expr}")
205+
if len(expressions) == 1:
206+
return InstrumentedExpression(f"{self._expr} WHERE {expressions[0]}")
207+
else:
208+
return InstrumentedExpression(
209+
f'{self._expr} WHERE {" AND ".join([f"({expr})" for expr in expressions])}'
210+
)
211+
212+
213+
E = InstrumentedExpression
134214

135215

136216
class InstrumentedField(InstrumentedExpression):
@@ -178,34 +258,6 @@ def __neg__(self) -> str: # type: ignore[override]
178258
"""Return the field name representation for descending sort order"""
179259
return f"-{self._expr}"
180260

181-
def asc(self) -> "InstrumentedField":
182-
"""Return the field name representation for ascending sort order.
183-
184-
For use in ES|QL queries only.
185-
"""
186-
return InstrumentedField(f"{self._expr} ASC", None)
187-
188-
def desc(self) -> "InstrumentedField":
189-
"""Return the field name representation for descending sort order.
190-
191-
For use in ES|QL queries only.
192-
"""
193-
return InstrumentedField(f"{self._expr} DESC", None)
194-
195-
def nulls_first(self) -> "InstrumentedField":
196-
"""Return the field name representation for nulls first sort order.
197-
198-
For use in ES|QL queries only.
199-
"""
200-
return InstrumentedField(f"{self._expr} NULLS FIRST", None)
201-
202-
def nulls_last(self) -> "InstrumentedField":
203-
"""Return the field name representation for nulls last sort order.
204-
205-
For use in ES|QL queries only.
206-
"""
207-
return InstrumentedField(f"{self._expr} NULLS LAST", None)
208-
209261
def __str__(self) -> str:
210262
return self._expr
211263

‎elasticsearch/esql/__init__.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,4 +15,4 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from .esql import ESQL # noqa: F401
18+
from .esql import ESQL, and_, not_, or_ # noqa: F401

‎elasticsearch/esql/esql.py‎

Lines changed: 31 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -61,7 +61,7 @@ def row(**params: ExpressionType) -> "Row":
6161
6262
query1 = ESQL.row(a=1, b="two", c=None)
6363
query2 = ESQL.row(a=[1, 2])
64-
query3 = ESQL.row(a="ROUND(1.23, 0)")
64+
query3 = ESQL.row(a=functions.round(1.23, 0))
6565
"""
6666
return Row(**params)
6767

@@ -126,8 +126,9 @@ def change_point(self, value: FieldType) -> "ChangePoint":
126126
query = (
127127
ESQL.row(key=list(range(1, 26)))
128128
.mv_expand("key")
129-
.eval(value="CASE(key<13, 0, 42)")
130-
.change_point("value").on("key")
129+
.eval(value=functions.case("key<13", 0, 42))
130+
.change_point("value")
131+
.on("key")
131132
.where("type IS NOT NULL")
132133
)
133134
"""
@@ -206,7 +207,7 @@ def eval(self, *columns: ExpressionType, **named_columns: ExpressionType) -> "Ev
206207
query2 = (
207208
ESQL.from_("employees")
208209
.eval("height * 3.281")
209-
.stats(avg_height_feet="AVG(`height * 3.281`)")
210+
.stats(avg_height_feet=functions.avg("`height * 3.281`"))
210211
)
211212
"""
212213
return Eval(self, *columns, **named_columns)
@@ -261,6 +262,15 @@ def grok(self, input: FieldType, pattern: str) -> "Grok":
261262
.keep("date", "ip", "email", "num")
262263
)
263264
query2 = (
265+
ESQL.row(a="2023-01-23T12:15:00.000Z 127.0.0.1 [email protected] 42")
266+
.grok(
267+
"a",
268+
"%{TIMESTAMP_ISO8601:date} %{IP:ip} %{EMAILADDRESS:email} %{NUMBER:num:int}",
269+
)
270+
.keep("date", "ip", "email", "num")
271+
.eval(date=functions.to_datetime("date"))
272+
)
273+
query3 = (
264274
ESQL.from_("addresses")
265275
.keep("city.name", "zip_code")
266276
.grok("zip_code", "%{WORD:zip_parts} %{WORD:zip_parts}")
@@ -291,7 +301,8 @@ def limit(self, max_number_of_rows: int) -> "Limit":
291301
292302
Examples::
293303
294-
query = ESQL.from_("employees").sort("emp_no ASC").limit(5)
304+
query1 = ESQL.from_("employees").sort("emp_no ASC").limit(5)
305+
query2 = ESQL.from_("index").stats(functions.avg("field1")).by("field2").limit(20000)
295306
"""
296307
return Limit(self, max_number_of_rows)
297308

@@ -982,3 +993,18 @@ def __init__(self, parent: ESQLBase, *expressions: ExpressionType):
982993

983994
def _render_internal(self) -> str:
984995
return f'WHERE {" AND ".join([f"{expr}" for expr in self._expressions])}'
996+
997+
998+
def and_(*expressions: InstrumentedExpression) -> "InstrumentedExpression":
999+
"""Combine two or more expressions with the AND operator."""
1000+
return InstrumentedExpression(" AND ".join([f"({expr})" for expr in expressions]))
1001+
1002+
1003+
def or_(*expressions: InstrumentedExpression) -> "InstrumentedExpression":
1004+
"""Combine two or more expressions with the OR operator."""
1005+
return InstrumentedExpression(" OR ".join([f"({expr})" for expr in expressions]))
1006+
1007+
1008+
def not_(expression: InstrumentedExpression) -> "InstrumentedExpression":
1009+
"""Negate an expression."""
1010+
return InstrumentedExpression(f"NOT ({expression})")

0 commit comments

Comments
 (0)