unit tests

Author: miguelgrinberg

elasticsearch/dsl/document_base.py

@@ -123,6 +123,13 @@ def __mod__(self, value: Any) -> "InstrumentedExpression":
     def __rmod__(self, value: Any) -> "InstrumentedExpression":
         return InstrumentedExpression(f"{json.dumps(value)} % {self._expr}")
 
+    def where(self, expr: "InstrumentedExpression") -> "InstrumentedExpression":
+        """Add a condition to be met for the row to be included.
+
+        Use only in expressions given in the ``STATS`` command.
+        """
+        return InstrumentedExpression(f"{self._expr} WHERE {json.dumps(expr)}")
+
 
 class InstrumentedField(InstrumentedExpression):
     """Proxy object for a mapped document field.
@@ -170,15 +177,31 @@ def __neg__(self) -> str:  # type: ignore[override]
         return f"-{self._expr}"
 
     def asc(self) -> "InstrumentedField":
+        """Return the field name representation for ascending sort order.
+
+        For use in ES|QL queries only.
+        """
         return InstrumentedField(f"{self._expr} ASC", None)
 
     def desc(self) -> "InstrumentedField":
+        """Return the field name representation for descending sort order.
+
+        For use in ES|QL queries only.
+        """
         return InstrumentedField(f"{self._expr} DESC", None)
 
     def nulls_first(self) -> "InstrumentedField":
+        """Return the field name representation for nulls first sort order.
+
+        For use in ES|QL queries only.
+        """
         return InstrumentedField(f"{self._expr} NULLS FIRST", None)
 
     def nulls_last(self) -> "InstrumentedField":
+        """Return the field name representation for nulls last sort order.
+
+        For use in ES|QL queries only.
+        """
         return InstrumentedField(f"{self._expr} NULLS LAST", None)
 
     def __str__(self) -> str:

elasticsearch/esql/esql.py

@@ -17,12 +17,12 @@
 
 import json
 from abc import ABC, abstractmethod
-from typing import Any, Dict, Optional, Tuple, Type
+from typing import Any, Dict, Optional, Tuple, Type, Union
 
-from ..dsl.document_base import DocumentBase, InstrumentedField
+from ..dsl.document_base import DocumentBase, InstrumentedExpression, InstrumentedField
 
-FieldType = InstrumentedField | str
-IndexType = Type[DocumentBase] | str
+FieldType = Union[InstrumentedField, str]
+IndexType = Union[Type[DocumentBase], str]
 ExpressionType = Any
 
 
@@ -537,12 +537,13 @@ class Row(ESQLBase):
 
     def __init__(self, **params: ExpressionType):
         super().__init__()
-        self._params = params
+        self._params = {
+            k: json.dumps(v) if not isinstance(v, InstrumentedExpression) else v
+            for k, v in params.items()
+        }
 
     def _render_internal(self) -> str:
-        return "ROW " + ", ".join(
-            [f"{k} = {json.dumps(v)}" for k, v in self._params.items()]
-        )
+        return "ROW " + ", ".join([f"{k} = {v}" for k, v in self._params.items()])
 
 
 class Show(ESQLBase):