Repository: scoutapp/scout_apm_python

Issue:
Instrument Elasticsearch 8.0
See [changelog](https://github.com/elastic/elasticsearch-py/releases/tag/v8.0.0) for more information.


Discussion:
I'm not sure how to best test this yet. Eventually we'll want to add the following sanitation logic to `_sanitatize_name`:

```python
        if op == "template" and len(ops) >= 2:
            template_name = "{}_{}".format(ops[-2][1:], op)
            template_names = {
                "search_template",
                "msearch_template",
                "render_template",
            }
            if template_name in template_names:
                return template_name.title()
        if op in known_names:
            return op.title()
        if op == "rethrottle":
            rethrottle_name = "{}_{}".format(ops[0][1:], op)
            rethrottle_names = {
                "delete_by_query_rethrottle",
                "reindex_rethrottle",
                "update_by_query_rethrottle",
            }
            if rethrottle_name in rethrottle_names:
                return rethrottle_name.title()
        if ops[0] == "_mvt":
            return "Search_mvt"
        if ops[1] == "_create":
            return "Create"
```

Resolving patch:
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -4,16 +4,25 @@
 import logging
 from collections import namedtuple
 
+import elasticsearch
 import wrapt
 
 from scout_apm.compat import get_pos_args, unwrap_decorators
 from scout_apm.core.tracked_request import TrackedRequest
 
 try:
-    from elasticsearch import Elasticsearch, Transport
+    from elasticsearch import Elasticsearch
 except ImportError:  # pragma: no cover
     Elasticsearch = None
-    Transport = None
+
+try:
+    # Transport was moved to elastic_transport as of v8.0.0
+    from elastic_transport import Transport
+except ImportError:  # pragma: no cover
+    try:
+        from elasticsearch import Transport
+    except ImportError:  # pragma: no cover
+        Transport = None
 
 logger = logging.getLogger(__name__)
 
@@ -21,10 +30,17 @@
 def ensure_installed():
     logger.debug("Instrumenting elasticsearch.")
 
-    if Elasticsearch is None:
-        logger.debug(
-            "Couldn't import elasticsearch.Elasticsearch - probably not installed."
-        )
+    if Elasticsearch is None or Transport is None:
+        if Elasticsearch is None:
+            logger.debug(
+                "Couldn't import elasticsearch.Elasticsearch "
+                "- probably not installed."
+            )
+        if Transport is None:
+            logger.debug(
+                "Couldn't import elasticsearch.Transport or "
+                "elastic_transport.Transport - probably not installed."
+            )
     else:
         ensure_client_instrumented()
         ensure_transport_instrumented()
@@ -32,51 +48,60 @@ def ensure_installed():
 
 ClientMethod = namedtuple("ClientMethod", ["name", "takes_index_argument"])
 
-CLIENT_METHODS = [
-    ClientMethod("bulk", True),
-    ClientMethod("clear_scroll", False),
-    ClientMethod("close", False),
-    ClientMethod("close_point_in_time", False),
-    ClientMethod("count", True),
-    ClientMethod("create", True),
-    ClientMethod("delete", True),
-    ClientMethod("delete_by_query", True),
-    ClientMethod("delete_by_query_rethrottle", False),
-    ClientMethod("delete_script", False),
-    ClientMethod("exists", True),
-    ClientMethod("exists_source", True),
-    ClientMethod("explain", True),
-    ClientMethod("field_caps", True),
-    ClientMethod("get", True),
-    ClientMethod("get_script", False),
-    ClientMethod("get_script_context", False),
-    ClientMethod("get_script_languages", False),
-    ClientMethod("get_source", True),
-    ClientMethod("index", True),
-    ClientMethod("info", False),
-    ClientMethod("mget", True),
-    ClientMethod("msearch", True),
-    ClientMethod("msearch_template", True),
-    ClientMethod("mtermvectors", True),
-    ClientMethod("open_point_in_time", True),
-    ClientMethod("ping", False),
-    ClientMethod("put_script", False),
-    ClientMethod("rank_eval", True),
-    ClientMethod("reindex", False),
-    ClientMethod("reindex_rethrottle", False),
-    ClientMethod("render_search_template", False),
-    ClientMethod("scripts_painless_execute", False),
-    ClientMethod("scroll", False),
-    ClientMethod("search", True),
-    ClientMethod("search_mvt", True),
-    ClientMethod("search_shards", True),
-    ClientMethod("search_template", True),
-    ClientMethod("termvectors", True),
-    ClientMethod("terms_enum", True),
-    ClientMethod("update", True),
-    ClientMethod("update_by_query", True),
-    ClientMethod("update_by_query_rethrottle", False),
-]
+VERSIONED_CLIENT_METHODS = {
+    "v7": [
+        ClientMethod("bulk", True),
+        ClientMethod("clear_scroll", False),
+        ClientMethod("close", False),
+        ClientMethod("close_point_in_time", False),
+        ClientMethod("count", True),
+        ClientMethod("create", True),
+        ClientMethod("delete", True),
+        ClientMethod("delete_by_query", True),
+        ClientMethod("delete_by_query_rethrottle", False),
+        ClientMethod("delete_script", False),
+        ClientMethod("exists", True),
+        ClientMethod("exists_source", True),
+        ClientMethod("explain", True),
+        ClientMethod("field_caps", True),
+        ClientMethod("get", True),
+        ClientMethod("get_script", False),
+        ClientMethod("get_script_context", False),
+        ClientMethod("get_script_languages", False),
+        ClientMethod("get_source", True),
+        ClientMethod("index", True),
+        ClientMethod("info", False),
+        ClientMethod("mget", True),
+        ClientMethod("msearch", True),
+        ClientMethod("msearch_template", True),
+        ClientMethod("mtermvectors", True),
+        ClientMethod("open_point_in_time", True),
+        ClientMethod("ping", False),
+        ClientMethod("put_script", False),
+        ClientMethod("rank_eval", True),
+        ClientMethod("reindex", False),
+        ClientMethod("reindex_rethrottle", False),
+        ClientMethod("render_search_template", False),
+        ClientMethod("scripts_painless_execute", False),
+        ClientMethod("scroll", False),
+        ClientMethod("search", True),
+        ClientMethod("search_mvt", True),
+        ClientMethod("search_shards", True),
+        ClientMethod("search_template", True),
+        ClientMethod("termvectors", True),
+        ClientMethod("terms_enum", True),
+        ClientMethod("update", True),
+        ClientMethod("update_by_query", True),
+        ClientMethod("update_by_query_rethrottle", False),
+    ],
+    "v8": [
+        ClientMethod("knn_search", True),
+    ],
+}
+
+CLIENT_METHODS = VERSIONED_CLIENT_METHODS["v7"][:]
+if elasticsearch.VERSION > (8, 0, 0):
+    CLIENT_METHODS += VERSIONED_CLIENT_METHODS["v8"]
 
 
 have_patched_client = False
@@ -185,8 +210,11 @@ def _sanitize_name(name):
             "bench",
             "bulk",
             "count",
+            "delete_by_query",
+            "execute",
             "exists",
             "explain",
+            "field_caps",
             "field_stats",
             "health",
             "mget",
@@ -195,14 +223,21 @@ def _sanitize_name(name):
             "msearch",
             "mtermvectors",
             "percolate",
+            "pit",
             "query",
+            "rank_eval",
+            "reindex",
+            "script_context",
+            "script_language",
             "scroll",
             "search_shards",
             "source",
             "suggest",
             "template",
             "termvectors",
+            "terms_enum",
             "update",
+            "update_by_query",
             "search",
         )
         if op in known_names:


Test patch:
diff --git a/tests/integration/instruments/test_elasticsearch.py b/tests/integration/instruments/test_elasticsearch.py
--- a/tests/integration/instruments/test_elasticsearch.py
+++ b/tests/integration/instruments/test_elasticsearch.py
@@ -13,6 +13,10 @@
 from tests.compat import mock
 from tests.tools import delete_attributes, skip_if_python_2
 
+skip_if_elasticsearch_v7 = pytest.mark.skipif(
+    elasticsearch.VERSION < (8, 0, 0), reason="Requires ElasticSearch 8"
+)
+
 
 @pytest.fixture(scope="module")
 def elasticsearch_client():
@@ -31,8 +35,9 @@ def test_all_client_attributes_accounted_for():
     public_attributes = {
         m for m in dir(elasticsearch.Elasticsearch) if not m.startswith("_")
     }
-    deliberately_ignored_attributes = set()
+    deliberately_ignored_attributes = {"perform_request", "options", "transport"}
     wrapped_methods = {m.name for m in CLIENT_METHODS}
+
     assert (
         public_attributes - deliberately_ignored_attributes - wrapped_methods
     ) == set()
@@ -71,10 +76,13 @@ def test_ensure_installed_twice(caplog):
 
 
 def test_ensure_installed_fail_no_client(caplog):
-    mock_no_client = mock.patch(
+    mock_no_elasticsearch = mock.patch(
         "scout_apm.instruments.elasticsearch.Elasticsearch", new=None
     )
-    with mock_no_client:
+    mock_no_transport = mock.patch(
+        "scout_apm.instruments.elasticsearch.Transport", new=None
+    )
+    with mock_no_elasticsearch, mock_no_transport:
         ensure_installed()
 
     assert caplog.record_tuples == [
@@ -88,6 +96,12 @@ def test_ensure_installed_fail_no_client(caplog):
             logging.DEBUG,
             "Couldn't import elasticsearch.Elasticsearch - probably not installed.",
         ),
+        (
+            "scout_apm.instruments.elasticsearch",
+            logging.DEBUG,
+            "Couldn't import elasticsearch.Transport or elastic_transport.Transport "
+            "- probably not installed.",
+        ),
     ]
 
 
@@ -168,7 +182,7 @@ def test_ping(elasticsearch_client, tracked_request):
     assert span.operation == "Elasticsearch/Ping"
 
 
-def test_search(elasticsearch_client, tracked_request):
+def test_search_v7_api(elasticsearch_client, tracked_request):
     elasticsearch_client.search()
 
     assert len(tracked_request.complete_spans) == 1
@@ -176,6 +190,15 @@ def test_search(elasticsearch_client, tracked_request):
     assert span.operation == "Elasticsearch/Unknown/Search"
 
 
+@skip_if_elasticsearch_v7
+def test_search(elasticsearch_client, tracked_request):
+    elasticsearch_client.options().search()
+
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Elasticsearch/Unknown/Search"
+
+
 @skip_if_python_2  # Cannot unwrap decorators on Python 2
 def test_search_arg_named_index(elasticsearch_client, tracked_request):
     with pytest.raises(elasticsearch.exceptions.NotFoundError):
@@ -232,9 +255,10 @@ def test_search_kwarg_index_list(elasticsearch_client, tracked_request):
 
 def test_perform_request_missing_url(elasticsearch_client, tracked_request):
     # Check Transport instrumentation doesn't crash if url is missing.
-    # This raises a TypeError when calling the original method.
-    with pytest.raises(TypeError):
-        elasticsearch_client.transport.perform_request("GET", params={}, body=None)
+    # This raises a TypeError and AttributeError when calling the original
+    # method.
+    with pytest.raises((TypeError, AttributeError)):
+        elasticsearch_client.transport.perform_request("GET", body=None)
 
     assert len(tracked_request.complete_spans) == 1
     span = tracked_request.complete_spans[0]
@@ -242,12 +266,10 @@ def test_perform_request_missing_url(elasticsearch_client, tracked_request):
 
 
 def test_perform_request_bad_url(elasticsearch_client, tracked_request):
-    with pytest.raises(TypeError):
+    with pytest.raises((TypeError, AttributeError)):
         # Transport instrumentation doesn't crash if url has the wrong type.
         # This raises a TypeError when calling the original method.
-        elasticsearch_client.transport.perform_request(
-            "GET", None, params={}, body=None
-        )
+        elasticsearch_client.transport.perform_request("GET", None, body=None)
 
     assert len(tracked_request.complete_spans) == 1
     span = tracked_request.complete_spans[0]
@@ -256,9 +278,7 @@ def test_perform_request_bad_url(elasticsearch_client, tracked_request):
 
 def test_perform_request_unknown_url(elasticsearch_client, tracked_request):
     # Transport instrumentation doesn't crash if url is unknown.
-    elasticsearch_client.transport.perform_request(
-        "GET", "/_nodes", params={}, body=None
-    )
+    elasticsearch_client.transport.perform_request("GET", "/_nodes", body=None)
 
     assert len(tracked_request.complete_spans) == 1
     span = tracked_request.complete_spans[0]
@@ -267,9 +287,7 @@ def test_perform_request_unknown_url(elasticsearch_client, tracked_request):
 
 def test_perform_request_known_url(elasticsearch_client, tracked_request):
     # Transport instrumentation doesn't crash if url is unknown.
-    elasticsearch_client.transport.perform_request(
-        "GET", "/_count", params={}, body=None
-    )
+    elasticsearch_client.transport.perform_request("GET", "/_count", body=None)
 
     assert len(tracked_request.complete_spans) == 1
     span = tracked_request.complete_spans[0]