elasticsearch 9.1.0__py3-none-any.whl → 9.1.2__py3-none-any.whl

elasticsearch/esql/esql.py

@@ -16,6 +16,7 @@
 #  under the License.
 
 import json
+import re
 from abc import ABC, abstractmethod
 from typing import Any, Dict, Optional, Tuple, Type, Union
 
@@ -111,6 +112,29 @@ class ESQLBase(ABC):
     def _render_internal(self) -> str:
         pass
 
+    @staticmethod
+    def _format_index(index: IndexType) -> str:
+        return index._index._name if hasattr(index, "_index") else str(index)
+
+    @staticmethod
+    def _format_id(id: FieldType, allow_patterns: bool = False) -> str:
+        s = str(id)  # in case it is an InstrumentedField
+        if allow_patterns and "*" in s:
+            return s  # patterns cannot be escaped
+        if re.fullmatch(r"[a-zA-Z_@][a-zA-Z0-9_\.]*", s):
+            return s
+        # this identifier needs to be escaped
+        s.replace("`", "``")
+        return f"`{s}`"
+
+    @staticmethod
+    def _format_expr(expr: ExpressionType) -> str:
+        return (
+            json.dumps(expr)
+            if not isinstance(expr, (str, InstrumentedExpression))
+            else str(expr)
+        )
+
     def _is_forked(self) -> bool:
         if self.__class__.__name__ == "Fork":
             return True
@@ -427,7 +451,7 @@ class ESQLBase(ABC):
         """
         return Sample(self, probability)
 
-    def sort(self, *columns: FieldType) -> "Sort":
+    def sort(self, *columns: ExpressionType) -> "Sort":
         """The ``SORT`` processing command sorts a table on one or more columns.
 
         :param columns: The columns to sort on.
@@ -570,15 +594,12 @@ class From(ESQLBase):
         return self
 
     def _render_internal(self) -> str:
-        indices = [
-            index if isinstance(index, str) else index._index._name
-            for index in self._indices
-        ]
+        indices = [self._format_index(index) for index in self._indices]
         s = f'{self.__class__.__name__.upper()} {", ".join(indices)}'
         if self._metadata_fields:
             s = (
                 s
-                + f' METADATA {", ".join([str(field) for field in self._metadata_fields])}'
+                + f' METADATA {", ".join([self._format_id(field) for field in self._metadata_fields])}'
             )
         return s
 
@@ -594,7 +615,11 @@ class Row(ESQLBase):
     def __init__(self, **params: ExpressionType):
         super().__init__()
         self._params = {
-            k: json.dumps(v) if not isinstance(v, InstrumentedExpression) else v
+            self._format_id(k): (
+                json.dumps(v)
+                if not isinstance(v, InstrumentedExpression)
+                else self._format_expr(v)
+            )
             for k, v in params.items()
         }
 
@@ -615,7 +640,7 @@ class Show(ESQLBase):
         self._item = item
 
     def _render_internal(self) -> str:
-        return f"SHOW {self._item}"
+        return f"SHOW {self._format_id(self._item)}"
 
 
 class Branch(ESQLBase):
@@ -667,11 +692,11 @@ class ChangePoint(ESQLBase):
         return self
 
     def _render_internal(self) -> str:
-        key = "" if not self._key else f" ON {self._key}"
+        key = "" if not self._key else f" ON {self._format_id(self._key)}"
         names = (
             ""
             if not self._type_name and not self._pvalue_name
-            else f' AS {self._type_name or "type"}, {self._pvalue_name or "pvalue"}'
+            else f' AS {self._format_id(self._type_name or "type")}, {self._format_id(self._pvalue_name or "pvalue")}'
         )
         return f"CHANGE_POINT {self._value}{key}{names}"
 
@@ -709,12 +734,13 @@ class Completion(ESQLBase):
     def _render_internal(self) -> str:
         if self._inference_id is None:
             raise ValueError("The completion command requires an inference ID")
+        with_ = {"inference_id": self._inference_id}
         if self._named_prompt:
             column = list(self._named_prompt.keys())[0]
             prompt = list(self._named_prompt.values())[0]
-            return f"COMPLETION {column} = {prompt} WITH {self._inference_id}"
+            return f"COMPLETION {self._format_id(column)} = {self._format_id(prompt)} WITH {json.dumps(with_)}"
         else:
-            return f"COMPLETION {self._prompt[0]} WITH {self._inference_id}"
+            return f"COMPLETION {self._format_id(self._prompt[0])} WITH {json.dumps(with_)}"
 
 
 class Dissect(ESQLBase):
@@ -742,9 +768,13 @@ class Dissect(ESQLBase):
 
     def _render_internal(self) -> str:
         sep = (
-            "" if self._separator is None else f' APPEND_SEPARATOR="{self._separator}"'
+            ""
+            if self._separator is None
+            else f" APPEND_SEPARATOR={json.dumps(self._separator)}"
+        )
+        return (
+            f"DISSECT {self._format_id(self._input)} {json.dumps(self._pattern)}{sep}"
         )
-        return f"DISSECT {self._input} {json.dumps(self._pattern)}{sep}"
 
 
 class Drop(ESQLBase):
@@ -760,7 +790,7 @@ class Drop(ESQLBase):
         self._columns = columns
 
     def _render_internal(self) -> str:
-        return f'DROP {", ".join([str(col) for col in self._columns])}'
+        return f'DROP {", ".join([self._format_id(col, allow_patterns=True) for col in self._columns])}'
 
 
 class Enrich(ESQLBase):
@@ -814,12 +844,18 @@ class Enrich(ESQLBase):
         return self
 
     def _render_internal(self) -> str:
-        on = "" if self._match_field is None else f" ON {self._match_field}"
+        on = (
+            ""
+            if self._match_field is None
+            else f" ON {self._format_id(self._match_field)}"
+        )
         with_ = ""
         if self._named_fields:
-            with_ = f' WITH {", ".join([f"{name} = {field}" for name, field in self._named_fields.items()])}'
+            with_ = f' WITH {", ".join([f"{self._format_id(name)} = {self._format_id(field)}" for name, field in self._named_fields.items()])}'
         elif self._fields is not None:
-            with_ = f' WITH {", ".join([str(field) for field in self._fields])}'
+            with_ = (
+                f' WITH {", ".join([self._format_id(field) for field in self._fields])}'
+            )
         return f"ENRICH {self._policy}{on}{with_}"
 
 
@@ -832,7 +868,10 @@ class Eval(ESQLBase):
     """
 
     def __init__(
-        self, parent: ESQLBase, *columns: FieldType, **named_columns: FieldType
+        self,
+        parent: ESQLBase,
+        *columns: ExpressionType,
+        **named_columns: ExpressionType,
     ):
         if columns and named_columns:
             raise ValueError(
@@ -844,10 +883,13 @@ class Eval(ESQLBase):
     def _render_internal(self) -> str:
         if isinstance(self._columns, dict):
             cols = ", ".join(
-                [f"{name} = {value}" for name, value in self._columns.items()]
+                [
+                    f"{self._format_id(name)} = {self._format_expr(value)}"
+                    for name, value in self._columns.items()
+                ]
             )
         else:
-            cols = ", ".join([f"{col}" for col in self._columns])
+            cols = ", ".join([f"{self._format_expr(col)}" for col in self._columns])
         return f"EVAL {cols}"
 
 
@@ -900,7 +942,7 @@ class Grok(ESQLBase):
         self._pattern = pattern
 
     def _render_internal(self) -> str:
-        return f"GROK {self._input} {json.dumps(self._pattern)}"
+        return f"GROK {self._format_id(self._input)} {json.dumps(self._pattern)}"
 
 
 class Keep(ESQLBase):
@@ -916,7 +958,7 @@ class Keep(ESQLBase):
         self._columns = columns
 
     def _render_internal(self) -> str:
-        return f'KEEP {", ".join([f"{col}" for col in self._columns])}'
+        return f'KEEP {", ".join([f"{self._format_id(col, allow_patterns=True)}" for col in self._columns])}'
 
 
 class Limit(ESQLBase):
@@ -932,7 +974,7 @@ class Limit(ESQLBase):
         self._max_number_of_rows = max_number_of_rows
 
     def _render_internal(self) -> str:
-        return f"LIMIT {self._max_number_of_rows}"
+        return f"LIMIT {json.dumps(self._max_number_of_rows)}"
 
 
 class LookupJoin(ESQLBase):
@@ -967,7 +1009,9 @@ class LookupJoin(ESQLBase):
             if isinstance(self._lookup_index, str)
             else self._lookup_index._index._name
         )
-        return f"LOOKUP JOIN {index} ON {self._field}"
+        return (
+            f"LOOKUP JOIN {self._format_index(index)} ON {self._format_id(self._field)}"
+        )
 
 
 class MvExpand(ESQLBase):
@@ -983,7 +1027,7 @@ class MvExpand(ESQLBase):
         self._column = column
 
     def _render_internal(self) -> str:
-        return f"MV_EXPAND {self._column}"
+        return f"MV_EXPAND {self._format_id(self._column)}"
 
 
 class Rename(ESQLBase):
@@ -999,7 +1043,7 @@ class Rename(ESQLBase):
         self._columns = columns
 
     def _render_internal(self) -> str:
-        return f'RENAME {", ".join([f"{old_name} AS {new_name}" for old_name, new_name in self._columns.items()])}'
+        return f'RENAME {", ".join([f"{self._format_id(old_name)} AS {self._format_id(new_name)}" for old_name, new_name in self._columns.items()])}'
 
 
 class Sample(ESQLBase):
@@ -1015,7 +1059,7 @@ class Sample(ESQLBase):
         self._probability = probability
 
     def _render_internal(self) -> str:
-        return f"SAMPLE {self._probability}"
+        return f"SAMPLE {json.dumps(self._probability)}"
 
 
 class Sort(ESQLBase):
@@ -1026,12 +1070,16 @@ class Sort(ESQLBase):
     in a single expression.
     """
 
-    def __init__(self, parent: ESQLBase, *columns: FieldType):
+    def __init__(self, parent: ESQLBase, *columns: ExpressionType):
         super().__init__(parent)
         self._columns = columns
 
     def _render_internal(self) -> str:
-        return f'SORT {", ".join([f"{col}" for col in self._columns])}'
+        sorts = [
+            " ".join([self._format_id(term) for term in str(col).split(" ")])
+            for col in self._columns
+        ]
+        return f'SORT {", ".join([f"{sort}" for sort in sorts])}'
 
 
 class Stats(ESQLBase):
@@ -1062,14 +1110,17 @@ class Stats(ESQLBase):
 
     def _render_internal(self) -> str:
         if isinstance(self._expressions, dict):
-            exprs = [f"{key} = {value}" for key, value in self._expressions.items()]
+            exprs = [
+                f"{self._format_id(key)} = {self._format_expr(value)}"
+                for key, value in self._expressions.items()
+            ]
         else:
-            exprs = [f"{expr}" for expr in self._expressions]
+            exprs = [f"{self._format_expr(expr)}" for expr in self._expressions]
         expression_separator = ",\n        "
         by = (
             ""
             if self._grouping_expressions is None
-            else f'\n        BY {", ".join([f"{expr}" for expr in self._grouping_expressions])}'
+            else f'\n        BY {", ".join([f"{self._format_expr(expr)}" for expr in self._grouping_expressions])}'
         )
         return f'STATS {expression_separator.join([f"{expr}" for expr in exprs])}{by}'
 
@@ -1087,7 +1138,7 @@ class Where(ESQLBase):
         self._expressions = expressions
 
     def _render_internal(self) -> str:
-        return f'WHERE {" AND ".join([f"{expr}" for expr in self._expressions])}'
+        return f'WHERE {" AND ".join([f"{self._format_expr(expr)}" for expr in self._expressions])}'
 
 
 def and_(*expressions: InstrumentedExpression) -> "InstrumentedExpression":

elasticsearch/esql/functions.py

@@ -19,11 +19,15 @@ import json
 from typing import Any
 
 from elasticsearch.dsl.document_base import InstrumentedExpression
-from elasticsearch.esql.esql import ExpressionType
+from elasticsearch.esql.esql import ESQLBase, ExpressionType
 
 
 def _render(v: Any) -> str:
-    return json.dumps(v) if not isinstance(v, InstrumentedExpression) else str(v)
+    return (
+        json.dumps(v)
+        if not isinstance(v, InstrumentedExpression)
+        else ESQLBase._format_expr(v)
+    )
 
 
 def abs(number: ExpressionType) -> InstrumentedExpression:
@@ -69,7 +73,9 @@ def atan2(
     :param y_coordinate: y coordinate. If `null`, the function returns `null`.
     :param x_coordinate: x coordinate. If `null`, the function returns `null`.
     """
-    return InstrumentedExpression(f"ATAN2({y_coordinate}, {x_coordinate})")
+    return InstrumentedExpression(
+        f"ATAN2({_render(y_coordinate)}, {_render(x_coordinate)})"
+    )
 
 
 def avg(number: ExpressionType) -> InstrumentedExpression:
@@ -114,7 +120,7 @@ def bucket(
     :param to: End of the range. Can be a number, a date or a date expressed as a string.
     """
     return InstrumentedExpression(
-        f"BUCKET({_render(field)}, {_render(buckets)}, {from_}, {_render(to)})"
+        f"BUCKET({_render(field)}, {_render(buckets)}, {_render(from_)}, {_render(to)})"
     )
 
 
@@ -169,7 +175,7 @@ def cidr_match(ip: ExpressionType, block_x: ExpressionType) -> InstrumentedExpre
     :param ip: IP address of type `ip` (both IPv4 and IPv6 are supported).
     :param block_x: CIDR block to test the IP against.
     """
-    return InstrumentedExpression(f"CIDR_MATCH({_render(ip)}, {block_x})")
+    return InstrumentedExpression(f"CIDR_MATCH({_render(ip)}, {_render(block_x)})")
 
 
 def coalesce(first: ExpressionType, rest: ExpressionType) -> InstrumentedExpression:
@@ -264,7 +270,7 @@ def date_diff(
     :param end_timestamp: A string representing an end timestamp
     """
     return InstrumentedExpression(
-        f"DATE_DIFF({_render(unit)}, {start_timestamp}, {end_timestamp})"
+        f"DATE_DIFF({_render(unit)}, {_render(start_timestamp)}, {_render(end_timestamp)})"
     )
 
 
@@ -285,7 +291,9 @@ def date_extract(
         the function returns `null`.
     :param date: Date expression. If `null`, the function returns `null`.
     """
-    return InstrumentedExpression(f"DATE_EXTRACT({date_part}, {_render(date)})")
+    return InstrumentedExpression(
+        f"DATE_EXTRACT({_render(date_part)}, {_render(date)})"
+    )
 
 
 def date_format(
@@ -301,7 +309,7 @@ def date_format(
     """
     if date_format is not None:
         return InstrumentedExpression(
-            f"DATE_FORMAT({json.dumps(date_format)}, {_render(date)})"
+            f"DATE_FORMAT({_render(date_format)}, {_render(date)})"
         )
     else:
         return InstrumentedExpression(f"DATE_FORMAT({_render(date)})")
@@ -317,7 +325,9 @@ def date_parse(
     :param date_string: Date expression as a string. If `null` or an empty
                         string, the function returns `null`.
     """
-    return InstrumentedExpression(f"DATE_PARSE({date_pattern}, {date_string})")
+    return InstrumentedExpression(
+        f"DATE_PARSE({_render(date_pattern)}, {_render(date_string)})"
+    )
 
 
 def date_trunc(
@@ -639,7 +649,7 @@ def min_over_time(field: ExpressionType) -> InstrumentedExpression:
 
 
 def multi_match(
-    query: ExpressionType, fields: ExpressionType, options: ExpressionType = None
+    query: ExpressionType, *fields: ExpressionType, options: ExpressionType = None
 ) -> InstrumentedExpression:
     """Use `MULTI_MATCH` to perform a multi-match query on the specified field.
     The multi_match query builds on the match query to allow multi-field queries.
@@ -651,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])})'
         )
 
 
@@ -929,7 +939,7 @@ def replace(
     :param new_string: Replacement string.
     """
     return InstrumentedExpression(
-        f"REPLACE({_render(string)}, {_render(regex)}, {new_string})"
+        f"REPLACE({_render(string)}, {_render(regex)}, {_render(new_string)})"
     )
 
 
@@ -1004,7 +1014,7 @@ def scalb(d: ExpressionType, scale_factor: ExpressionType) -> InstrumentedExpres
     :param scale_factor: Numeric expression for the scale factor. If `null`, the
                          function returns `null`.
     """
-    return InstrumentedExpression(f"SCALB({_render(d)}, {scale_factor})")
+    return InstrumentedExpression(f"SCALB({_render(d)}, {_render(scale_factor)})")
 
 
 def sha1(input: ExpressionType) -> InstrumentedExpression:
@@ -1116,7 +1126,7 @@ def st_contains(
                    first. This means it is not possible to combine `geo_*` and
                    `cartesian_*` parameters.
     """
-    return InstrumentedExpression(f"ST_CONTAINS({geom_a}, {geom_b})")
+    return InstrumentedExpression(f"ST_CONTAINS({_render(geom_a)}, {_render(geom_b)})")
 
 
 def st_disjoint(
@@ -1135,7 +1145,7 @@ def st_disjoint(
                    first. This means it is not possible to combine `geo_*` and
                    `cartesian_*` parameters.
     """
-    return InstrumentedExpression(f"ST_DISJOINT({geom_a}, {geom_b})")
+    return InstrumentedExpression(f"ST_DISJOINT({_render(geom_a)}, {_render(geom_b)})")
 
 
 def st_distance(
@@ -1153,7 +1163,7 @@ def st_distance(
                    also have the same coordinate system as the first. This means it
                    is not possible to combine `geo_point` and `cartesian_point` parameters.
     """
-    return InstrumentedExpression(f"ST_DISTANCE({geom_a}, {geom_b})")
+    return InstrumentedExpression(f"ST_DISTANCE({_render(geom_a)}, {_render(geom_b)})")
 
 
 def st_envelope(geometry: ExpressionType) -> InstrumentedExpression:
@@ -1208,7 +1218,7 @@ def st_geohash_to_long(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input geohash grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOHASH_TO_LONG({grid_id})")
+    return InstrumentedExpression(f"ST_GEOHASH_TO_LONG({_render(grid_id)})")
 
 
 def st_geohash_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
@@ -1218,7 +1228,7 @@ def st_geohash_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input geohash grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOHASH_TO_STRING({grid_id})")
+    return InstrumentedExpression(f"ST_GEOHASH_TO_STRING({_render(grid_id)})")
 
 
 def st_geohex(
@@ -1254,7 +1264,7 @@ def st_geohex_to_long(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input geohex grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOHEX_TO_LONG({grid_id})")
+    return InstrumentedExpression(f"ST_GEOHEX_TO_LONG({_render(grid_id)})")
 
 
 def st_geohex_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
@@ -1264,7 +1274,7 @@ def st_geohex_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input Geohex grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOHEX_TO_STRING({grid_id})")
+    return InstrumentedExpression(f"ST_GEOHEX_TO_STRING({_render(grid_id)})")
 
 
 def st_geotile(
@@ -1300,7 +1310,7 @@ def st_geotile_to_long(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input geotile grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOTILE_TO_LONG({grid_id})")
+    return InstrumentedExpression(f"ST_GEOTILE_TO_LONG({_render(grid_id)})")
 
 
 def st_geotile_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
@@ -1310,7 +1320,7 @@ def st_geotile_to_string(grid_id: ExpressionType) -> InstrumentedExpression:
     :param grid_id: Input geotile grid-id. The input can be a single- or
                     multi-valued column or an expression.
     """
-    return InstrumentedExpression(f"ST_GEOTILE_TO_STRING({grid_id})")
+    return InstrumentedExpression(f"ST_GEOTILE_TO_STRING({_render(grid_id)})")
 
 
 def st_intersects(
@@ -1330,7 +1340,9 @@ def st_intersects(
                    first. This means it is not possible to combine `geo_*` and
                    `cartesian_*` parameters.
     """
-    return InstrumentedExpression(f"ST_INTERSECTS({geom_a}, {geom_b})")
+    return InstrumentedExpression(
+        f"ST_INTERSECTS({_render(geom_a)}, {_render(geom_b)})"
+    )
 
 
 def st_within(geom_a: ExpressionType, geom_b: ExpressionType) -> InstrumentedExpression:
@@ -1346,7 +1358,7 @@ def st_within(geom_a: ExpressionType, geom_b: ExpressionType) -> InstrumentedExp
                    first. This means it is not possible to combine `geo_*` and
                    `cartesian_*` parameters.
     """
-    return InstrumentedExpression(f"ST_WITHIN({geom_a}, {geom_b})")
+    return InstrumentedExpression(f"ST_WITHIN({_render(geom_a)}, {_render(geom_b)})")
 
 
 def st_x(point: ExpressionType) -> InstrumentedExpression:

elasticsearch/helpers/__init__.py

@@ -19,12 +19,21 @@ from .._async.helpers import async_bulk, async_reindex, async_scan, async_stream
 from .._utils import fixup_module_metadata
 from .actions import _chunk_actions  # noqa: F401
 from .actions import _process_bulk_chunk  # noqa: F401
-from .actions import bulk, expand_action, parallel_bulk, reindex, scan, streaming_bulk
+from .actions import (
+    BULK_FLUSH,
+    bulk,
+    expand_action,
+    parallel_bulk,
+    reindex,
+    scan,
+    streaming_bulk,
+)
 from .errors import BulkIndexError, ScanError
 
 __all__ = [
     "BulkIndexError",
     "ScanError",
+    "BULK_FLUSH",
     "expand_action",
     "streaming_bulk",
     "bulk",