adding response_hook to elastic instrumentation

Author: ItayGibel-helios

Diff for: CHANGELOG.md

@@ -9,6 +9,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - `opentelemetry-sdk-extension-aws` Release AWS Python SDK Extension as 1.0.0
   ([#667](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/667))
 
+### Added
+- `opentelemetry-instrumentation-elasticsearch` Added `response_hook` callback
+
 ### Changed
 - `opentelemetry-instrumentation-botocore` Unpatch botocore Endpoint.prepare_request on uninstrument
   ([#664](https://github.com/open-telemetry/opentelemetry-python-contrib/pull/664))

Diff for: instrumentation/opentelemetry-instrumentation-elasticsearch/src/opentelemetry/instrumentation/elasticsearch/__init__.py

@@ -97,17 +97,20 @@ def _instrument(self, **kwargs):
         """
         tracer_provider = kwargs.get("tracer_provider")
         tracer = get_tracer(__name__, __version__, tracer_provider)
+        response_hook = kwargs.get("response_hook", _default_response_hook)
         _wrap(
             elasticsearch,
             "Transport.perform_request",
-            _wrap_perform_request(tracer, self._span_name_prefix),
+            _wrap_perform_request(
+                tracer, self._span_name_prefix, response_hook
+            ),
         )
 
     def _uninstrument(self, **kwargs):
         unwrap(elasticsearch.Transport, "perform_request")
 
 
-def _wrap_perform_request(tracer, span_name_prefix):
+def _wrap_perform_request(tracer, span_name_prefix, response_hook):
     # pylint: disable=R0912
     def wrapper(wrapped, _, args, kwargs):
         method = url = None
@@ -143,13 +146,16 @@ def wrapper(wrapped, _, args, kwargs):
                     span.set_attribute(key, value)
 
             rv = wrapped(*args, **kwargs)
-            if isinstance(rv, dict) and span.is_recording():
-                for member in _ATTRIBUTES_FROM_RESULT:
-                    if member in rv:
-                        span.set_attribute(
-                            "elasticsearch.{0}".format(member),
-                            str(rv[member]),
-                        )
+            response_hook(span, rv)
             return rv
 
     return wrapper
+
+
+def _default_response_hook(span, response):
+    if isinstance(response, dict) and span.is_recording():
+        for member in _ATTRIBUTES_FROM_RESULT:
+            if member in response:
+                span.set_attribute(
+                    "elasticsearch.{0}".format(member), str(response[member]),
+                )

Diff for: instrumentation/opentelemetry-instrumentation-elasticsearch/tests/test_elasticsearch.py

@@ -11,7 +11,7 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import json
 import os
 import threading
 from ast import literal_eval
@@ -316,3 +316,55 @@ def test_dsl_index(self, request_mock):
                 "title": "About searching",
             },
         )
+
+    def test_response_hook(self, request_mock):
+        response_attribute_name = "db.query_result"
+
+        def response_hook(span, response):
+            if span and span.is_recording():
+                span.set_attribute(
+                    response_attribute_name, json.dumps(response)
+                )
+
+        ElasticsearchInstrumentor().uninstrument()
+        ElasticsearchInstrumentor().instrument(response_hook=response_hook)
+
+        response_payload = {
+            "took": 9,
+            "timed_out": False,
+            "_shards": {
+                "total": 1,
+                "successful": 1,
+                "skipped": 0,
+                "failed": 0,
+            },
+            "hits": {
+                "total": {"value": 1, "relation": "eq"},
+                "max_score": 0.18232156,
+                "hits": [
+                    {
+                        "_index": "test-index",
+                        "_type": "tweet",
+                        "_id": "1",
+                        "_score": 0.18232156,
+                        "_source": {"name": "tester"},
+                    }
+                ],
+            },
+        }
+
+        request_mock.return_value = (
+            1,
+            {},
+            json.dumps(response_payload),
+        )
+        es = Elasticsearch()
+        es.get(index="test-index", doc_type="tweet", id=1)
+
+        spans = self.get_finished_spans()
+
+        self.assertEqual(1, len(spans))
+        self.assertEqual(
+            json.dumps(response_payload),
+            spans[0].attributes[response_attribute_name],
+        )