[Backport 8.19] ES|QL query builder integration with the DSL module (elastic#3056)

* ES|QL query builder integration with the DSL module (elastic#3048)

* ES|QL query builder integration with the DSL module

* esql DSL tests

* more esql DSL tests

* documentation

* add esql+dsl example

* review feedback

(cherry picked from commit 228e66c)

* python 3.8 typing fix

---------

Co-authored-by: Miguel Grinberg <[email protected]>

Author: github-actions[bot]

docs/guide/dsl/howto.asciidoc

@@ -4,4 +4,5 @@ include::search_dsl.asciidoc[]
 include::persistence.asciidoc[]
 include::faceted_search.asciidoc[]
 include::update_by_query.asciidoc[]
+include::esql.asciidoc[]
 include::asyncio.asciidoc[]

docs/guide/dsl/tutorials.asciidoc

@@ -83,17 +83,17 @@ system:
 [source,python]
 ----
 from datetime import datetime
-from elasticsearch.dsl import Document, Date, Integer, Keyword, Text, connections
+from elasticsearch.dsl import Document, Date, Integer, Keyword, Text, connections, mapped_field
 
 # Define a default Elasticsearch client
 connections.create_connection(hosts="https://localhost:9200")
 
 class Article(Document):
-    title = Text(analyzer='snowball', fields={'raw': Keyword()})
-    body = Text(analyzer='snowball')
-    tags = Keyword()
-    published_from = Date()
-    lines = Integer()
+    title: str = mapped_field(Text(analyzer='snowball', fields={'raw': Keyword()}))
+    body: str = mapped_field(Text(analyzer='snowball'))
+    tags: list[str] = mapped_field(Keyword())
+    published_from: datetime
+    lines: int
 
     class Index:
         name = 'blog'
@@ -229,13 +229,31 @@ savings offered by the `Search` object, and additionally allows one to
 update the results of the search based on a script assigned in the same
 manner.
 
+==== ES|QL Queries
+
+The DSL module features an integration with the ES|QL query builder, consisting of two methods available in all `Document` sub-classes: `esql_from()` and `esql_execute()`. Using the `Article` document from above, we can search for up to ten articles that include `"world"` in their titles with the following ES|QL query:
+
+[source,python]
+----
+from elasticsearch.esql import functions
+
+query = Article.esql_from().where(functions.match(Article.title, 'world')).limit(10)
+for a in Article.esql_execute(query):
+    print(a.title)
+----
+
+Review the ES|QL Query Builder section to learn more about building ES|QL queries in Python.
+
 ==== Migration from the standard client
 
+<<<<<<< HEAD:docs/guide/dsl/tutorials.asciidoc
 You don't have to port your entire application to get the benefits of
 the DSL module, you can start gradually by creating a `Search` object
 from your existing `dict`, modifying it using the API and serializing it
 back to a `dict`:
 
+==== Migration from the standard client
+
 [source,python]
 ----
 body = {...} # insert complicated query here

docs/guide/esql-query-builder.asciidoc

@@ -21,21 +21,21 @@ You can then see the assembled ES|QL query by printing the resulting query objec
 
 [source, python]
 ----------------------------
->>> query
+>>> print(query)
 FROM employees
 | SORT emp_no
 | KEEP first_name, last_name, height
 | EVAL height_feet = height * 3.281, height_cm = height * 100
 | LIMIT 3
 ----------------------------
 
-To execute this query, you can cast it to a string and pass the string to the `client.esql.query()` endpoint:
+To execute this query, you can pass it to the `client.esql.query()` endpoint:
 
 [source, python]
 ----------------------------
 >>> from elasticsearch import Elasticsearch
 >>> client = Elasticsearch(hosts=[os.environ['ELASTICSEARCH_URL']])
->>> response = client.esql.query(query=str(query))
+>>> response = client.esql.query(query=query)
 ----------------------------
 
 The response body contains a `columns` attribute with the list of columns included in the results, and a `values` attribute with the list of results for the query, each given as a list of column values. Here is a possible response body returned by the example query given above:
@@ -228,7 +228,7 @@ def find_employee_by_name(name):
         .keep("first_name", "last_name", "height")
         .where(E("first_name") == E("?"))
     )
-    return client.esql.query(query=str(query), params=[name])
+    return client.esql.query(query=query, params=[name])
 ----------------------------
 
 Here the part of the query in which the untrusted data needs to be inserted is replaced with a parameter, which in ES|QL is defined by the question mark. When using Python expressions, the parameter must be given as `E("?")` so that it is treated as an expression and not as a literal string.

elasticsearch/dsl/_async/document.py

@@ -20,6 +20,7 @@
     TYPE_CHECKING,
     Any,
     AsyncIterable,
+    AsyncIterator,
     Dict,
     List,
     Optional,
@@ -42,6 +43,7 @@
 
 if TYPE_CHECKING:
     from elasticsearch import AsyncElasticsearch
+    from elasticsearch.esql.esql import ESQLBase
 
 
 class AsyncIndexMeta(DocumentMeta):
@@ -520,3 +522,85 @@ async def __anext__(self) -> Dict[str, Any]:
                 return action
 
         return await async_bulk(es, Generate(actions), **kwargs)
+
+    @classmethod
+    async def esql_execute(
+        cls,
+        query: "ESQLBase",
+        return_additional: bool = False,
+        ignore_missing_fields: bool = False,
+        using: Optional[AsyncUsingType] = None,
+        **kwargs: Any,
+    ) -> AsyncIterator[Union[Self, Tuple[Self, Dict[str, Any]]]]:
+        """
+        Execute the given ES|QL query and return an iterator of 2-element tuples,
+        where the first element is an instance of this ``Document`` and the
+        second a dictionary with any remaining columns requested in the query.
+
+        :arg query: an ES|QL query object created with the ``esql_from()`` method.
+        :arg return_additional: if ``False`` (the default), this method returns
+            document objects. If set to ``True``, the method returns tuples with
+            a document in the first element and a dictionary with any additional
+            columns returned by the query in the second element.
+        :arg ignore_missing_fields: if ``False`` (the default), all the fields of
+            the document must be present in the query, or else an exception is
+            raised. Set to ``True`` to allow missing fields, which will result in
+            partially initialized document objects.
+        :arg using: connection alias to use, defaults to ``'default'``
+        :arg kwargs: additional options for the ``client.esql.query()`` function.
+        """
+        es = cls._get_connection(using)
+        response = await es.esql.query(query=str(query), **kwargs)
+        query_columns = [col["name"] for col in response.body.get("columns", [])]
+
+        # Here we get the list of columns defined in the document, which are the
+        # columns that we will take from each result to assemble the document
+        # object.
+        # When `for_esql=False` is passed below by default, the list will include
+        # nested fields, which ES|QL does not return, causing an error. When passing
+        # `ignore_missing_fields=True` the list will be generated with
+        # `for_esql=True`, so the error will not occur, but the documents will
+        # not have any Nested objects in them.
+        doc_fields = set(cls._get_field_names(for_esql=ignore_missing_fields))
+        if not ignore_missing_fields and not doc_fields.issubset(set(query_columns)):
+            raise ValueError(
+                f"Not all fields of {cls.__name__} were returned by the query. "
+                "Make sure your document does not use Nested fields, which are "
+                "currently not supported in ES|QL. To force the query to be "
+                "evaluated in spite of the missing fields, pass set the "
+                "ignore_missing_fields=True option in the esql_execute() call."
+            )
+        non_doc_fields: set[str] = set(query_columns) - doc_fields - {"_id"}
+        index_id = query_columns.index("_id")
+
+        results = response.body.get("values", [])
+        for column_values in results:
+            # create a dictionary with all the document fields, expanding the
+            # dot notation returned by ES|QL into the recursive dictionaries
+            # used by Document.from_dict()
+            doc_dict: Dict[str, Any] = {}
+            for col, val in zip(query_columns, column_values):
+                if col in doc_fields:
+                    cols = col.split(".")
+                    d = doc_dict
+                    for c in cols[:-1]:
+                        if c not in d:
+                            d[c] = {}
+                        d = d[c]
+                    d[cols[-1]] = val
+
+            # create the document instance
+            obj = cls(meta={"_id": column_values[index_id]})
+            obj._from_dict(doc_dict)
+
+            if return_additional:
+                # build a dict with any other values included in the response
+                other = {
+                    col: val
+                    for col, val in zip(query_columns, column_values)
+                    if col in non_doc_fields
+                }
+
+                yield obj, other
+            else:
+                yield obj

elasticsearch/dsl/_sync/document.py

@@ -21,6 +21,7 @@
     Any,
     Dict,
     Iterable,
+    Iterator,
     List,
     Optional,
     Tuple,
@@ -42,6 +43,7 @@
 
 if TYPE_CHECKING:
     from elasticsearch import Elasticsearch
+    from elasticsearch.esql.esql import ESQLBase
 
 
 class IndexMeta(DocumentMeta):
@@ -512,3 +514,85 @@ def __next__(self) -> Dict[str, Any]:
                 return action
 
         return bulk(es, Generate(actions), **kwargs)
+
+    @classmethod
+    def esql_execute(
+        cls,
+        query: "ESQLBase",
+        return_additional: bool = False,
+        ignore_missing_fields: bool = False,
+        using: Optional[UsingType] = None,
+        **kwargs: Any,
+    ) -> Iterator[Union[Self, Tuple[Self, Dict[str, Any]]]]:
+        """
+        Execute the given ES|QL query and return an iterator of 2-element tuples,
+        where the first element is an instance of this ``Document`` and the
+        second a dictionary with any remaining columns requested in the query.
+
+        :arg query: an ES|QL query object created with the ``esql_from()`` method.
+        :arg return_additional: if ``False`` (the default), this method returns
+            document objects. If set to ``True``, the method returns tuples with
+            a document in the first element and a dictionary with any additional
+            columns returned by the query in the second element.
+        :arg ignore_missing_fields: if ``False`` (the default), all the fields of
+            the document must be present in the query, or else an exception is
+            raised. Set to ``True`` to allow missing fields, which will result in
+            partially initialized document objects.
+        :arg using: connection alias to use, defaults to ``'default'``
+        :arg kwargs: additional options for the ``client.esql.query()`` function.
+        """
+        es = cls._get_connection(using)
+        response = es.esql.query(query=str(query), **kwargs)
+        query_columns = [col["name"] for col in response.body.get("columns", [])]
+
+        # Here we get the list of columns defined in the document, which are the
+        # columns that we will take from each result to assemble the document
+        # object.
+        # When `for_esql=False` is passed below by default, the list will include
+        # nested fields, which ES|QL does not return, causing an error. When passing
+        # `ignore_missing_fields=True` the list will be generated with
+        # `for_esql=True`, so the error will not occur, but the documents will
+        # not have any Nested objects in them.
+        doc_fields = set(cls._get_field_names(for_esql=ignore_missing_fields))
+        if not ignore_missing_fields and not doc_fields.issubset(set(query_columns)):
+            raise ValueError(
+                f"Not all fields of {cls.__name__} were returned by the query. "
+                "Make sure your document does not use Nested fields, which are "
+                "currently not supported in ES|QL. To force the query to be "
+                "evaluated in spite of the missing fields, pass set the "
+                "ignore_missing_fields=True option in the esql_execute() call."
+            )
+        non_doc_fields: set[str] = set(query_columns) - doc_fields - {"_id"}
+        index_id = query_columns.index("_id")
+
+        results = response.body.get("values", [])
+        for column_values in results:
+            # create a dictionary with all the document fields, expanding the
+            # dot notation returned by ES|QL into the recursive dictionaries
+            # used by Document.from_dict()
+            doc_dict: Dict[str, Any] = {}
+            for col, val in zip(query_columns, column_values):
+                if col in doc_fields:
+                    cols = col.split(".")
+                    d = doc_dict
+                    for c in cols[:-1]:
+                        if c not in d:
+                            d[c] = {}
+                        d = d[c]
+                    d[cols[-1]] = val
+
+            # create the document instance
+            obj = cls(meta={"_id": column_values[index_id]})
+            obj._from_dict(doc_dict)
+
+            if return_additional:
+                # build a dict with any other values included in the response
+                other = {
+                    col: val
+                    for col, val in zip(query_columns, column_values)
+                    if col in non_doc_fields
+                }
+
+                yield obj, other
+            else:
+                yield obj

elasticsearch/dsl/document_base.py

@@ -28,6 +28,7 @@
     List,
     Optional,
     Tuple,
+    Type,
     TypeVar,
     Union,
     get_args,
@@ -49,6 +50,7 @@
 if TYPE_CHECKING:
     from elastic_transport import ObjectApiResponse
 
+    from ..esql.esql import ESQLBase
     from .index_base import IndexBase
 
 
@@ -602,3 +604,44 @@ def to_dict(self, include_meta: bool = False, skip_empty: bool = True) -> Dict[s
 
         meta["_source"] = d
         return meta
+
+    @classmethod
+    def _get_field_names(
+        cls, for_esql: bool = False, nested_class: Optional[Type[InnerDoc]] = None
+    ) -> List[str]:
+        """Return the list of field names used by this document.
+        If the document has nested objects, their fields are reported using dot
+        notation. If the ``for_esql`` argument is set to ``True``, the list omits
+        nested fields, which are currently unsupported in ES|QL.
+        """
+        fields = []
+        class_ = nested_class or cls
+        for field_name in class_._doc_type.mapping:
+            field = class_._doc_type.mapping[field_name]
+            if isinstance(field, Object):
+                if for_esql and isinstance(field, Nested):
+                    # ES|QL does not recognize Nested fields at this time
+                    continue
+                sub_fields = cls._get_field_names(
+                    for_esql=for_esql, nested_class=field._doc_class
+                )
+                for sub_field in sub_fields:
+                    fields.append(f"{field_name}.{sub_field}")
+            else:
+                fields.append(field_name)
+        return fields
+
+    @classmethod
+    def esql_from(cls) -> "ESQLBase":
+        """Return a base ES|QL query for instances of this document class.
+
+        The returned query is initialized with ``FROM`` and ``KEEP`` statements,
+        and can be completed as desired.
+        """
+        from ..esql import ESQL  # here to avoid circular imports
+
+        return (
+            ESQL.from_(cls)
+            .metadata("_id")
+            .keep("_id", *tuple(cls._get_field_names(for_esql=True)))
+        )

elasticsearch/esql/__init__.py

@@ -16,4 +16,4 @@
 #  under the License.
 
 from ..dsl import E  # noqa: F401
-from .esql import ESQL, and_, not_, or_  # noqa: F401
+from .esql import ESQL, ESQLBase, and_, not_, or_  # noqa: F401

elasticsearch/esql/functions.py

@@ -661,11 +661,11 @@ def multi_match(
     """
     if options is not None:
         return InstrumentedExpression(
-            f"MULTI_MATCH({_render(query)}, {_render(fields)}, {_render(options)})"
+            f'MULTI_MATCH({_render(query)}, {", ".join([_render(c) for c in fields])}, {_render(options)})'
         )
     else:
         return InstrumentedExpression(
-            f"MULTI_MATCH({_render(query)}, {_render(fields)})"
+            f'MULTI_MATCH({_render(query)}, {", ".join([_render(c) for c in fields])})'
         )