if possible, get pymongo address from socket, otherwise from nodes (#797)

Author: beniwohli

elasticapm/instrumentation/packages/pymongo.py

@@ -72,7 +72,11 @@ class PyMongoInstrumentation(AbstractInstrumentedModule):
  def call(self, module, method, wrapped, instance, args, kwargs):
  cls_name, method_name = method.split(".", 1)
  signature = ".".join([instance.full_name, method_name])
- host, port = instance.database.client.address
+ nodes = instance.database.client.nodes
+ if nodes:
+ host, port = list(nodes)[0]
+ else:
+ host, port = None, None
  destination_info = {
  "address": host,
  "port": port,
@@ -97,7 +101,13 @@ class PyMongoBulkInstrumentation(AbstractInstrumentedModule):
  def call(self, module, method, wrapped, instance, args, kwargs):
  collection = instance._BulkOperationBuilder__bulk.collection
  signature = ".".join([collection.full_name, "bulk.execute"])
- with capture_span(signature, span_type="db", span_subtype="mongodb", span_action="query"):
+ with capture_span(
+ signature,
+ span_type="db",
+ span_subtype="mongodb",
+ span_action="query",
+ extra={"destination": {"service": {"name": "mongodb", "resource": "mongodb", "type": "db"}}},
+ ):
  return wrapped(*args, **kwargs)
 
 
@@ -109,5 +119,18 @@ class PyMongoCursorInstrumentation(AbstractInstrumentedModule):
  def call(self, module, method, wrapped, instance, args, kwargs):
  collection = instance.collection
  signature = ".".join([collection.full_name, "cursor.refresh"])
- with capture_span(signature, span_type="db", span_subtype="mongodb", span_action="query"):
- return wrapped(*args, **kwargs)
+ with capture_span(
+ signature,
+ span_type="db",
+ span_subtype="mongodb",
+ span_action="query",
+ extra={"destination": {"service": {"name": "mongodb", "resource": "mongodb", "type": "db"}}},
+ ) as span:
+ response = wrapped(*args, **kwargs)
+ if span.context:
+ host, port = instance.address
+ span.context["destination"]["address"] = host
+ span.context["destination"]["port"] = port
+ else:
+ pass
+ return response

tests/instrumentation/pymongo_tests.py

@@ -192,17 +192,22 @@ def test_collection_find(instrument, elasticapm_client, mongo_database):
  blogposts.append({"author": "Tom", "comments": i})
  mongo_database.blogposts.insert(blogposts)
  r = mongo_database.blogposts.insert(blogpost)
- elasticapm_client.events[TRANSACTION]
  elasticapm_client.begin_transaction("transaction.test")
  r = list(mongo_database.blogposts.find({"comments": {"$gt": 995}}))
 
  elasticapm_client.end_transaction("transaction.test")
  transactions = elasticapm_client.events[TRANSACTION]
- span = _get_pymongo_span(elasticapm_client.spans_for_transaction(transactions[0]))
+ spans = elasticapm_client.spans_for_transaction(transactions[0])
+ span = _get_pymongo_span(spans)
  assert span["type"] == "db"
  assert span["subtype"] == "mongodb"
  assert span["action"] == "query"
  assert span["name"] == "elasticapm_test.blogposts.cursor.refresh"
+ assert span["context"]["destination"] == {
+ "address": os.environ.get("MONGODB_HOST", "localhost"),
+ "port": int(os.environ.get("MONGODB_PORT", 27017)),
+ "service": {"name": "mongodb", "resource": "mongodb", "type": "db"},
+ }
 
 
 @pytest.mark.integrationtest
@@ -220,6 +225,11 @@ def test_collection_find_one(instrument, elasticapm_client, mongo_database):
  assert span["subtype"] == "mongodb"
  assert span["action"] == "query"
  assert span["name"] == "elasticapm_test.blogposts.find_one"
+ assert span["context"]["destination"] == {
+ "address": os.environ.get("MONGODB_HOST", "localhost"),
+ "port": int(os.environ.get("MONGODB_PORT", 27017)),
+ "service": {"name": "mongodb", "resource": "mongodb", "type": "db"},
+ }
 
 
 @pytest.mark.integrationtest