Repository: scoutapp/scout_apm_python
Pull request #439 (2020-01-02 12:23:00)

Track CherryPy static files
When testing the integration added in #431, I found that the `favicon.ico` route added by `cherrypy.quickstart()` was showing up as "Unknown" controller in the Scout UI. This should fix that.


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index a175271..460bac0 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,13 +5,17 @@
 ### Added
 
 - Use Python 3.8's standard library `importlib.metadata.distributions()` to
-  determine package versions, rather than `pkg_resources`. The backport
-  `importlib_metadata` is required on older Python versions
+  determine package versions, rather than `pkg_resources`. The [backport
+  package
+  `importlib_metadata`](https://pypi.org/project/importlib-metadata/) is
+  required on older Python versions
   ([PR #430](https://github.com/scoutapp/scout_apm_python/pull/430)).
 - Add CherryPy integration
   ([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/cherrypy.py b/src/scout_apm/cherrypy.py
index 226b527..a09e6da 100644
--- a/src/scout_apm/cherrypy.py
+++ b/src/scout_apm/cherrypy.py
@@ -108,6 +108,18 @@ def get_operation_name(request):
     while hasattr(real_handler, "callable"):
         real_handler = real_handler.callable
 
+    # Looks like it's from HandlerTool
+    if getattr(real_handler, "__name__", "") == "handle_func":
+        try:
+            wrapped_tool = real_handler.__closure__[2].cell_contents.callable
+        except (AttributeError, IndexError):
+            pass
+        else:
+            try:
+                return "Controller/{}".format(wrapped_tool.__name__)
+            except AttributeError:
+                pass
+
     # Not a method? Not from an exposed view then
     if not hasattr(real_handler, "__self__"):
         return None
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()
 
diff --git a/tests/integration/test_cherrypy.py b/tests/integration/test_cherrypy.py
index cf9ea13..fd77664 100644
--- a/tests/integration/test_cherrypy.py
+++ b/tests/integration/test_cherrypy.py
@@ -47,6 +47,9 @@ def app_with_scout(scout_config=None):
             cherrypy.response.status = 503
             return "Something went wrong"
 
+        # Serve this source file statically
+        static_file = cherrypy.tools.staticfile.handler(__file__)
+
     app = cherrypy.Application(Views(), "/", config=None)
 
     # Setup according to https://docs.scoutapm.com/#cherrypy
@@ -188,6 +191,20 @@ def test_return_error(tracked_requests):
     )
 
 
+def test_static_file(tracked_requests):
+    with app_with_scout() as app:
+        # CherryPy serves its built-in favicon by default
+        response = TestApp(app).get("/static-file/")
+
+    assert response.status_int == 200
+    assert response.headers["content-type"] == "text/x-python;charset=utf-8"
+    assert len(tracked_requests) == 1
+    tracked_request = tracked_requests[0]
+    assert len(tracked_request.complete_spans) == 1
+    span = tracked_request.complete_spans[0]
+    assert span.operation == "Controller/staticfile"
+
+
 def test_not_found(tracked_requests):
     with app_with_scout() as app:
         response = TestApp(app).get("/not-found/", expect_errors=True)