Set json/x-ndjson serializer to compatibility mode mimetype too

elastic · Jan 28, 2022 · 0da2ba9 · 0da2ba9
1 parent dc1d7cd
commit 0da2ba9
Show file tree

Hide file tree

Showing 11 changed files with 267 additions and 48 deletions.
diff --git a/elasticsearch/_async/client/__init__.py b/elasticsearch/_async/client/__init__.py
@@ -352,9 +352,22 @@ def __init__(
             if meta_header is not DEFAULT:
                 transport_kwargs["meta_header"] = meta_header
 
-            transport_serializers = DEFAULT_SERIALIZERS
+            transport_serializers = DEFAULT_SERIALIZERS.copy()
             if serializers is not DEFAULT:
                 transport_serializers.update(serializers)
+
+                # Override compatibility serializers from their non-compat mimetypes too.
+                # So we use the same serializer for requests and responses.
+                for mime_subtype in ("json", "x-ndjson"):
+                    if f"application/{mime_subtype}" in serializers:
+                        compat_mimetype = (
+                            f"application/vnd.elasticsearch+{mime_subtype}"
+                        )
+                        if compat_mimetype not in serializers:
+                            transport_serializers[compat_mimetype] = serializers[
+                                f"application/{mime_subtype}"
+                            ]
+
             transport_kwargs["serializers"] = transport_serializers
 
             transport_kwargs["default_mimetype"] = default_mimetype

diff --git a/elasticsearch/_async/client/_base.py b/elasticsearch/_async/client/_base.py
@@ -44,6 +44,7 @@
 )
 from elastic_transport.client_utils import DEFAULT, DefaultType
 
+from ..._version import __versionstr__
 from ...compat import warn_stacklevel
 from ...exceptions import (
     HTTP_EXCEPTIONS,
@@ -56,6 +57,11 @@
 from .utils import _TYPE_ASYNC_SNIFF_CALLBACK, _base64_auth_header, _quote_query
 
 _WARNING_RE = re.compile(r"\"([^\"]*)\"")
+_COMPAT_MIMETYPE_TEMPLATE = "application/vnd.elasticsearch+%s; compatible-with=" + str(
+    __versionstr__.partition(".")[0]
+)
+_COMPAT_MIMETYPE_RE = re.compile(r"application/(json|x-ndjson|vnd\.mapbox-vector-tile)")
+_COMPAT_MIMETYPE_SUB = _COMPAT_MIMETYPE_TEMPLATE % (r"\g<1>",)
 
 
 def resolve_auth_headers(
@@ -166,7 +172,9 @@ async def sniff_callback(
                 meta, node_infos = await transport.perform_request(
                     "GET",
                     "/_nodes/_all/http",
-                    headers={"accept": "application/json"},
+                    headers={
+                        "accept": "application/vnd.elasticsearch+json; compatible-with=8"
+                    },
                     request_timeout=(
                         sniff_options.sniff_timeout
                         if not sniff_options.is_initial_sniff
@@ -257,6 +265,19 @@ async def perform_request(
         else:
             request_headers = self._headers
 
+        def mimetype_header_to_compat(header: str) -> None:
+            # Converts all parts of a Accept/Content-Type headers
+            # from application/X -> application/vnd.elasticsearch+X
+            nonlocal request_headers
+            mimetype = request_headers.get(header, None)
+            if mimetype:
+                request_headers[header] = _COMPAT_MIMETYPE_RE.sub(
+                    _COMPAT_MIMETYPE_SUB, mimetype
+                )
+
+        mimetype_header_to_compat("Accept")
+        mimetype_header_to_compat("Content-Type")
+
         if params:
             target = f"{path}?{_quote_query(params)}"
         else:

diff --git a/elasticsearch/_sync/client/__init__.py b/elasticsearch/_sync/client/__init__.py
@@ -352,9 +352,22 @@ def __init__(
             if meta_header is not DEFAULT:
                 transport_kwargs["meta_header"] = meta_header
 
-            transport_serializers = DEFAULT_SERIALIZERS
+            transport_serializers = DEFAULT_SERIALIZERS.copy()
             if serializers is not DEFAULT:
                 transport_serializers.update(serializers)
+
+                # Override compatibility serializers from their non-compat mimetypes too.
+                # So we use the same serializer for requests and responses.
+                for mime_subtype in ("json", "x-ndjson"):
+                    if f"application/{mime_subtype}" in serializers:
+                        compat_mimetype = (
+                            f"application/vnd.elasticsearch+{mime_subtype}"
+                        )
+                        if compat_mimetype not in serializers:
+                            transport_serializers[compat_mimetype] = serializers[
+                                f"application/{mime_subtype}"
+                            ]
+
             transport_kwargs["serializers"] = transport_serializers
 
             transport_kwargs["default_mimetype"] = default_mimetype

diff --git a/elasticsearch/_sync/client/_base.py b/elasticsearch/_sync/client/_base.py
@@ -44,6 +44,7 @@
 )
 from elastic_transport.client_utils import DEFAULT, DefaultType
 
+from ..._version import __versionstr__
 from ...compat import warn_stacklevel
 from ...exceptions import (
     HTTP_EXCEPTIONS,
@@ -56,6 +57,11 @@
 from .utils import _TYPE_SYNC_SNIFF_CALLBACK, _base64_auth_header, _quote_query
 
 _WARNING_RE = re.compile(r"\"([^\"]*)\"")
+_COMPAT_MIMETYPE_TEMPLATE = "application/vnd.elasticsearch+%s; compatible-with=" + str(
+    __versionstr__.partition(".")[0]
+)
+_COMPAT_MIMETYPE_RE = re.compile(r"application/(json|x-ndjson|vnd\.mapbox-vector-tile)")
+_COMPAT_MIMETYPE_SUB = _COMPAT_MIMETYPE_TEMPLATE % (r"\g<1>",)
 
 
 def resolve_auth_headers(
@@ -166,7 +172,9 @@ def sniff_callback(
                 meta, node_infos = transport.perform_request(
                     "GET",
                     "/_nodes/_all/http",
-                    headers={"accept": "application/json"},
+                    headers={
+                        "accept": "application/vnd.elasticsearch+json; compatible-with=8"
+                    },
                     request_timeout=(
                         sniff_options.sniff_timeout
                         if not sniff_options.is_initial_sniff
@@ -257,6 +265,19 @@ def perform_request(
         else:
             request_headers = self._headers
 
+        def mimetype_header_to_compat(header: str) -> None:
+            # Converts all parts of a Accept/Content-Type headers
+            # from application/X -> application/vnd.elasticsearch+X
+            nonlocal request_headers
+            mimetype = request_headers.get(header, None)
+            if mimetype:
+                request_headers[header] = _COMPAT_MIMETYPE_RE.sub(
+                    _COMPAT_MIMETYPE_SUB, mimetype
+                )
+
+        mimetype_header_to_compat("Accept")
+        mimetype_header_to_compat("Content-Type")
+
         if params:
             target = f"{path}?{_quote_query(params)}"
         else:

diff --git a/elasticsearch/serializer.py b/elasticsearch/serializer.py
@@ -25,7 +25,6 @@
 from elastic_transport import Serializer as Serializer
 from elastic_transport import TextSerializer as TextSerializer
 
-from .compat import to_bytes
 from .exceptions import SerializationError
 
 INTEGER_TYPES = ()
@@ -37,7 +36,8 @@
     "JsonSerializer",
     "TextSerializer",
     "NdjsonSerializer",
-    "CompatibilityModeSerializer",
+    "CompatibilityModeJsonSerializer",
+    "CompatibilityModeNdjsonSerializer",
     "MapboxVectorTileSerializer",
 ]
 
@@ -80,27 +80,12 @@ def default(self, data: Any) -> Any:
         return JsonSerializer.default(self, data)
 
 
-class CompatibilityModeSerializer(JsonSerializer):
+class CompatibilityModeJsonSerializer(JsonSerializer):
     mimetype: ClassVar[str] = "application/vnd.elasticsearch+json"
 
-    def dumps(self, data: Any) -> bytes:
-        if isinstance(data, str):
-            data = data.encode("utf-8", "surrogatepass")
-        if isinstance(data, bytes):
-            return data
-        if isinstance(data, (tuple, list)):
-            return NdjsonSerializer.dumps(self, data)  # type: ignore
-        return JsonSerializer.dumps(self, data)
 
-    def loads(self, data: bytes) -> Any:
-        if isinstance(data, str):
-            data = to_bytes(data, "utf-8")
-        if isinstance(data, bytes) and data.endswith(b"\n"):
-            return NdjsonSerializer.loads(self, data)  # type: ignore
-        try:  # Try as JSON first but if that fails then try NDJSON.
-            return JsonSerializer.loads(self, data)
-        except SerializationError:
-            return NdjsonSerializer.loads(self, data)  # type: ignore
+class CompatibilityModeNdjsonSerializer(NdjsonSerializer):
+    mimetype: ClassVar[str] = "application/vnd.elasticsearch+x-ndjson"
 
 
 class MapboxVectorTileSerializer(Serializer):
@@ -119,7 +104,8 @@ def dumps(self, data: bytes) -> bytes:
     JsonSerializer.mimetype: JsonSerializer(),
     MapboxVectorTileSerializer.mimetype: MapboxVectorTileSerializer(),
     NdjsonSerializer.mimetype: NdjsonSerializer(),
-    CompatibilityModeSerializer.mimetype: CompatibilityModeSerializer(),
+    CompatibilityModeJsonSerializer.mimetype: CompatibilityModeJsonSerializer(),
+    CompatibilityModeNdjsonSerializer.mimetype: CompatibilityModeNdjsonSerializer(),
 }
 
 # Alias for backwards compatibility

diff --git a/test_elasticsearch/test_async/test_transport.py b/test_elasticsearch/test_async/test_transport.py
@@ -243,7 +243,7 @@ async def test_client_meta_header_not_sent(self):
         calls = client.transport.node_pool.get().calls
         assert 1 == len(calls)
         assert calls[0][1]["headers"] == {
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
     async def test_body_surrogates_replaced_encoded_into_bytes(self):
@@ -391,7 +391,9 @@ async def test_sniff_on_start_ignores_sniff_timeout(self):
             ("GET", "/_nodes/_all/http"),
             {
                 "body": None,
-                "headers": {"accept": "application/json"},
+                "headers": {
+                    "accept": "application/vnd.elasticsearch+json; compatible-with=8"
+                },
                 "request_timeout": None,  # <-- Should be None instead of 12
             },
         )
@@ -418,7 +420,7 @@ async def test_sniff_uses_sniff_timeout(self):
             {
                 "body": None,
                 "headers": {
-                    "accept": "application/json",
+                    "accept": "application/vnd.elasticsearch+json; compatible-with=8",
                 },
                 "request_timeout": DEFAULT,
             },
@@ -427,7 +429,9 @@ async def test_sniff_uses_sniff_timeout(self):
             ("GET", "/_nodes/_all/http"),
             {
                 "body": None,
-                "headers": {"accept": "application/json"},
+                "headers": {
+                    "accept": "application/vnd.elasticsearch+json; compatible-with=8"
+                },
                 "request_timeout": 12,
             },
         )
@@ -681,7 +685,7 @@ async def test_unsupported_product_error(headers):
         {
             "body": None,
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "request_timeout": DEFAULT,
         },

diff --git a/test_elasticsearch/test_client/test_deprecated_options.py b/test_elasticsearch/test_client/test_deprecated_options.py
@@ -135,6 +135,7 @@ class CustomSerializer(JsonSerializer):
         "application/json",
         "text/*",
         "application/vnd.elasticsearch+json",
+        "application/vnd.elasticsearch+x-ndjson",
     }
 
     client = Elasticsearch(
@@ -154,5 +155,6 @@ class CustomSerializer(JsonSerializer):
         "application/json",
         "text/*",
         "application/vnd.elasticsearch+json",
+        "application/vnd.elasticsearch+x-ndjson",
         "application/cbor",
     }
diff --git a/test_elasticsearch/test_client/test_options.py b/test_elasticsearch/test_client/test_options.py
@@ -139,7 +139,7 @@ def test_options_passed_to_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
         }
@@ -157,7 +157,7 @@ def test_options_passed_to_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
             "request_timeout": 1,
@@ -182,7 +182,7 @@ def test_options_passed_to_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
             "request_timeout": 1,
@@ -209,7 +209,7 @@ async def test_options_passed_to_async_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
         }
@@ -227,7 +227,7 @@ async def test_options_passed_to_async_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
             "request_timeout": 1,
@@ -252,7 +252,7 @@ async def test_options_passed_to_async_perform_request(self):
         assert call.pop("client_meta") is DEFAULT
         assert call == {
             "headers": {
-                "accept": "application/json",
+                "accept": "application/vnd.elasticsearch+json; compatible-with=8",
             },
             "body": None,
             "request_timeout": 1,
@@ -294,7 +294,7 @@ def test_http_headers_overrides(self):
 
         assert call["headers"] == {
             "key": "val",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client.options(headers={"key1": "val"}).indices.get(index="2")
@@ -303,15 +303,15 @@ def test_http_headers_overrides(self):
         assert call["headers"] == {
             "key": "val",
             "key1": "val",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client.options(headers={"key": "val2"}).indices.get(index="3")
         call = calls[("GET", "/3")][0]
 
         assert call["headers"] == {
             "key": "val2",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client = Elasticsearch(
@@ -338,14 +338,14 @@ def test_user_agent_override(self):
         call = calls[("GET", "/1")][0]
         assert call["headers"] == {
             "user-agent": "custom1",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client.indices.get(index="2", headers={"user-agent": "custom2"})
         call = calls[("GET", "/2")][0]
         assert call["headers"] == {
             "user-agent": "custom2",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client = Elasticsearch(
@@ -359,12 +359,12 @@ def test_user_agent_override(self):
         call = calls[("GET", "/1")][0]
         assert call["headers"] == {
             "user-agent": "custom3",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }
 
         client.indices.get(index="2", headers={"user-agent": "custom4"})
         call = calls[("GET", "/2")][0]
         assert call["headers"] == {
             "user-agent": "custom4",
-            "accept": "application/json",
+            "accept": "application/vnd.elasticsearch+json; compatible-with=8",
         }