Repository: scoutapp/scout_apm_python (python)

Issue:
Instrument non-index Elasticsearch client methods
There are many methods on the `elasticsearch-py` client class that we aren't instrumenting, for example `ping()`, because they don't operate on an index.

We should capture all these calls too - I'm sure many applications have such calls that would be good to show on traces.

https://github.com/elastic/elasticsearch-py/blob/master/elasticsearch/client/__init__.py

Pull request description:
Add tracking for missing Elasticsearch client methods
Fixes #435. Copying what #436 did for PyMongo.

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a175271..5f34fd8 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -12,6 +12,8 @@
   ([PR #431](https://github.com/scoutapp/scout_apm_python/pull/431)).
 - Added tracking for missing PyMongo Collection methods
   ([PR #436](https://github.com/scoutapp/scout_apm_python/pull/436)).
+- Added tracking for missing Elasticsearch client methods
+  ([PR #438](https://github.com/scoutapp/scout_apm_python/pull/438)).
 
 ### Removed
 
diff --git a/src/scout_apm/instruments/elasticsearch.py b/src/scout_apm/instruments/elasticsearch.py
index a262627..307df35 100644
--- a/src/scout_apm/instruments/elasticsearch.py
+++ b/src/scout_apm/instruments/elasticsearch.py
@@ -2,6 +2,7 @@
 from __future__ import absolute_import, division, print_function, unicode_literals
 
 import logging
+from collections import namedtuple
 
 import wrapt
 
@@ -26,31 +27,46 @@ def ensure_installed():
         ensure_transport_instrumented()
 
 
+ClientMethod = namedtuple("ClientMethod", ["name", "takes_index_argument"])
+
 CLIENT_METHODS = [
-    "bulk",
-    "count",
-    "create",
-    "delete",
-    "delete_by_query",
-    "exists",
-    "exists_source",
-    "explain",
-    "field_caps",
-    "get",
-    "get_source",
-    "index",
-    "mget",
-    "msearch",
-    "msearch_template",
-    "mtermvectors",
-    "reindex",
-    "reindex_rethrottle",
-    "search",
-    "search_shards",
-    "search_template",
-    "termvectors",
-    "update",
-    "update_by_query",
+    ClientMethod("bulk", True),
+    ClientMethod("clear_scroll", 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_source", True),
+    ClientMethod("index", True),
+    ClientMethod("info", False),
+    ClientMethod("mget", True),
+    ClientMethod("msearch", True),
+    ClientMethod("msearch_template", True),
+    ClientMethod("mtermvectors", 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_context", False),
+    ClientMethod("scripts_painless_execute", False),
+    ClientMethod("scroll", False),
+    ClientMethod("search", True),
+    ClientMethod("search_shards", True),
+    ClientMethod("search_template", True),
+    ClientMethod("termvectors", True),
+    ClientMethod("update", True),
+    ClientMethod("update_by_query", True),
+    ClientMethod("update_by_query_rethrottle", False),
 ]
 
 
@@ -61,13 +77,14 @@ def ensure_client_instrumented():
     global have_patched_client
 
     if not have_patched_client:
-        for name in CLIENT_METHODS:
+        for name, takes_index_argument in CLIENT_METHODS:
             try:
-                setattr(
-                    Elasticsearch,
-                    name,
-                    wrap_client_method(getattr(Elasticsearch, name)),
-                )
+                method = getattr(Elasticsearch, name)
+                if takes_index_argument:
+                    wrapped = wrap_client_index_method(method)
+                else:
+                    wrapped = wrap_client_method(method)
+                setattr(Elasticsearch, name, wrapped)
             except Exception as exc:
                 logger.warning(
                     "Unable to instrument elasticsearch.Elasticsearch.%s: %r",
@@ -80,7 +97,7 @@ def ensure_client_instrumented():
 
 
 @wrapt.decorator
-def wrap_client_method(wrapped, instance, args, kwargs):
+def wrap_client_index_method(wrapped, instance, args, kwargs):
     def _get_index(index, *args, **kwargs):
         return index
 
@@ -105,6 +122,19 @@ def wrap_client_method(wrapped, instance, args, kwargs):
         tracked_request.stop_span()
 
 
+@wrapt.decorator
+def wrap_client_method(wrapped, instance, args, kwargs):
+    camel_name = "".join(c.title() for c in wrapped.__name__.split("_"))
+    operation = "Elasticsearch/{}".format(camel_name)
+    tracked_request = TrackedRequest.instance()
+    tracked_request.start_span(operation=operation, ignore_children=True)
+
+    try:
+        return wrapped(*args, **kwargs)
+    finally:
+        tracked_request.stop_span()
+
+
 have_patched_transport = False
 
 


Test patch:
diff --git a/tests/integration/instruments/test_elasticsearch.py b/tests/integration/instruments/test_elasticsearch.py
index 602f6cb..241a801 100644
--- a/tests/integration/instruments/test_elasticsearch.py
+++ b/tests/integration/instruments/test_elasticsearch.py
@@ -1,6 +1,7 @@
 # coding=utf-8
 from __future__ import absolute_import, division, print_function, unicode_literals
 
+import inspect
 import logging
 import os
 
@@ -8,8 +9,9 @@ import elasticsearch
 import elasticsearch.exceptions
 import pytest
 
-from scout_apm.instruments.elasticsearch import ensure_installed
+from scout_apm.instruments.elasticsearch import CLIENT_METHODS, ensure_installed
 from tests.compat import mock
+from tests.tools import delete_attributes, skip_if_python_2
 
 
 @pytest.fixture(scope="module")
@@ -21,6 +23,36 @@ def elasticsearch_client():
     yield elasticsearch.Elasticsearch(os.environ["ELASTICSEARCH_URL"])
 
 
+def test_all_client_attributes_accounted_for():
+    public_attributes = {
+        m for m in dir(elasticsearch.Elasticsearch) if not m.startswith("_")
+    }
+    deliberately_ignored_attributes = set()
+    wrapped_methods = {m.name for m in CLIENT_METHODS}
+    assert (
+        public_attributes - deliberately_ignored_attributes - wrapped_methods
+    ) == set()
+
+
+@pytest.mark.parametrize(["method_name"], [[m.name] for m in CLIENT_METHODS])
+def test_all_client_methods_exist(method_name):
+    assert hasattr(elasticsearch.Elasticsearch, method_name)
+
+
+@skip_if_python_2
+@pytest.mark.parametrize(
+    ["method_name", "takes_index_argument"],
+    [[m.name, m.takes_index_argument] for m in CLIENT_METHODS],
+)
+def test_all_client_methods_match_index_argument(method_name, takes_index_argument):
+    signature = inspect.signature(getattr(elasticsearch.Elasticsearch, method_name))
+
+    if takes_index_argument:
+        assert "index" in signature.parameters
+    else:
+        assert "index" not in signature.parameters
+
+
 def test_ensure_installed_twice(caplog):
     ensure_installed()
     ensure_installed()
@@ -59,10 +91,8 @@ def test_ensure_installed_fail_no_client_bulk(caplog):
     mock_not_patched = mock.patch(
         "scout_apm.instruments.elasticsearch.have_patched_client", new=False
     )
-    mock_client = mock.patch("scout_apm.instruments.elasticsearch.Elasticsearch")
-    with mock_not_patched, mock_client as mocked_client:
-        del mocked_client.bulk
-
+    mock_no_bulk = delete_attributes(elasticsearch.Elasticsearch, "bulk")
+    with mock_not_patched, mock_no_bulk:
         ensure_installed()
 
     assert len(caplog.record_tuples) == 2
@@ -103,6 +133,14 @@ def test_ensure_installed_fail_no_transport_perform_request(caplog):
     )
 
 
+def test_ping(elasticsearch_client, tracked_request):
+    elasticsearch_client.ping()
+
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Elasticsearch/Ping"
+
+
 def test_search(elasticsearch_client, tracked_request):
     elasticsearch_client.search()