Skip to content

Commit b4b6da4

Browse files
miguelgrinbergmiguelgrinberg
committed
more operators
1 parent 3d88ee5 commit b4b6da4

File tree

5 files changed

+256
-81
lines changed

5 files changed

+256
-81
lines changed

‎elasticsearch/dsl/__init__.py‎

Lines changed: 1 addition & 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,

‎elasticsearch/dsl/document_base.py‎

Lines changed: 89 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):
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,127 @@ 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):
132+
"""Compare the expression against NULL."""
133+
return InstrumentedExpression(f"{self._expr} IS NULL")
134+
135+
def is_not_null(self):
136+
"""Compare the expression against NOT NULL."""
137+
return InstrumentedExpression(f"{self._expr} IS NOT NULL")
138+
139+
def in_(self, *values: Any):
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):
145+
"""Filter the expression using a string pattern."""
146+
if len(patterns) == 1:
147+
return InstrumentedExpression(f'{self._expr} LIKE {self._render_value(patterns[0])}')
148+
else:
149+
return InstrumentedExpression(f'{self._expr} LIKE ({", ".join([self._render_value(p) for p in patterns])})')
150+
151+
def rlike(self, *patterns: str):
152+
"""Filter the expression using a regular expression."""
153+
if len(patterns) == 1:
154+
return InstrumentedExpression(f'{self._expr} RLIKE {self._render_value(patterns[0])}')
155+
else:
156+
return InstrumentedExpression(f'{self._expr} RLIKE ({", ".join([self._render_value(p) for p in patterns])})')
157+
158+
def match(self, query):
159+
"""Perform a match query on the field."""
160+
return InstrumentedExpression(f"{self._expr}:{self._render_value(query)}")
161+
162+
def asc(self) -> "InstrumentedExpression":
163+
"""Return the field name representation for ascending sort order.
164+
165+
For use in ES|QL queries only.
166+
"""
167+
return InstrumentedExpression(f"{self._expr} ASC")
168+
169+
def desc(self) -> "InstrumentedExpression":
170+
"""Return the field name representation for descending sort order.
171+
172+
For use in ES|QL queries only.
173+
"""
174+
return InstrumentedExpression(f"{self._expr} DESC")
175+
176+
def nulls_first(self) -> "InstrumentedExpression":
177+
"""Return the field name representation for nulls first sort order.
178+
179+
For use in ES|QL queries only.
180+
"""
181+
return InstrumentedExpression(f"{self._expr} NULLS FIRST")
182+
183+
def nulls_last(self) -> "InstrumentedExpression":
184+
"""Return the field name representation for nulls last sort order.
185+
186+
For use in ES|QL queries only.
187+
"""
188+
return InstrumentedExpression(f"{self._expr} NULLS LAST")
125189

126190
def where(
127-
self, expr: Union[str, "InstrumentedExpression"]
191+
self, *expressions: Union[str, "InstrumentedExpression"]
128192
) -> "InstrumentedExpression":
129193
"""Add a condition to be met for the row to be included.
130194
131195
Use only in expressions given in the ``STATS`` command.
132196
"""
133-
return InstrumentedExpression(f"{self._expr} WHERE {expr}")
197+
if len(expressions) == 1:
198+
return InstrumentedExpression(f"{self._expr} WHERE {expressions[0]}")
199+
else:
200+
return InstrumentedExpression(
201+
f'{self._expr} WHERE {" AND ".join([f"({expr})" for expr in expressions])}')
202+
203+
204+
E = InstrumentedExpression
134205

135206

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

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-
209252
def __str__(self) -> str:
210253
return self._expr
211254

‎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_, or_, not_ # 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):
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):
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):
1009+
"""Negate an expression."""
1010+
return InstrumentedExpression(f"NOT ({expression})")

0 commit comments

Comments
 (0)